<commit_before>// Copyright Project Harbor 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 awsecr

import (
	"errors"
	"fmt"
	"regexp"

	"github.com/aws/aws-sdk-go/aws/awserr"
	awsecrapi "github.com/aws/aws-sdk-go/service/ecr"
	"github.com/goharbor/harbor/src/lib/log"
	adp "github.com/goharbor/harbor/src/pkg/reg/adapter"
	"github.com/goharbor/harbor/src/pkg/reg/adapter/native"
	"github.com/goharbor/harbor/src/pkg/reg/model"
)

const (
	regionPattern = "https://(?:api|\\d+\\.dkr)\\.ecr\\.([\\w\\-]+)\\.amazonaws\\.com"
)

var (
	regionRegexp = regexp.MustCompile(regionPattern)
)

func init() {
	if err := adp.RegisterFactory(model.RegistryTypeAwsEcr, new(factory)); err != nil {
		log.Errorf("failed to register factory for %s: %v", model.RegistryTypeAwsEcr, err)
		return
	}
	log.Infof("the factory for adapter %s registered", model.RegistryTypeAwsEcr)
}

func newAdapter(registry *model.Registry) (*adapter, error) {
	region, err := parseRegion(registry.URL)
	if err != nil {
		return nil, err
	}
	svc, err := getAwsSvc(
		region, registry.Credential.AccessKey, registry.Credential.AccessSecret, registry.Insecure, nil)
	if err != nil {
		return nil, err
	}
	authorizer := NewAuth(registry.Credential.AccessKey, svc)
	return &adapter{
		registry: registry,
		Adapter:  native.NewAdapterWithAuthorizer(registry, authorizer),
		cacheSvc: svc,
	}, nil
}

func parseRegion(url string) (string, error) {
	rs := regionRegexp.FindStringSubmatch(url)
	if rs == nil {
		return "", errors.New("bad aws url")
	}
	return rs[1], nil
}

type factory struct {
}

// Create ...
func (f *factory) Create(r *model.Registry) (adp.Adapter, error) {
	return newAdapter(r)
}

// AdapterPattern ...
func (f *factory) AdapterPattern() *model.AdapterPattern {
	return getAdapterInfo()
}

var (
	_ adp.Adapter          = (*adapter)(nil)
	_ adp.ArtifactRegistry = (*adapter)(nil)
)

type adapter struct {
	*native.Adapter
	registry *model.Registry
	cacheSvc *awsecrapi.ECR
}

func (*adapter) Info() (info *model.RegistryInfo, err error) {
	return &model.RegistryInfo{
		Type: model.RegistryTypeAwsEcr,
		SupportedResourceTypes: []string{
			model.ResourceTypeImage,
		},
		SupportedResourceFilters: []*model.FilterStyle{
			{
				Type:  model.FilterTypeName,
				Style: model.FilterStyleTypeText,
			},
			{
				Type:  model.FilterTypeTag,
				Style: model.FilterStyleTypeText,
			},
		},
		SupportedTriggers: []string{
			model.TriggerTypeManual,
			model.TriggerTypeScheduled,
		},
	}, nil
}

func getAdapterInfo() *model.AdapterPattern {
	var endpoints []*model.Endpoint
	// https://docs.aws.amazon.com/general/latest/gr/rande.html#regional-endpoints
	for _, e := range []string{
		"us-east-2",
		"us-east-1",
		"us-west-1",
		"us-west-2",
		"af-south-1",
		"ap-east-1",
		"ap-south-1",
		"ap-northeast-3",
		"ap-northeast-2",
		"ap-southeast-1",
		"ap-southeast-2",
		"ap-northeast-1",
		"ca-central-1",
		"cn-north-1",
		"cn-northwest-1",
		"eu-central-1",
		"eu-west-1",
		"eu-west-2",
		"eu-south-1",
		"eu-west-3",
		"eu-north-1",
		"me-south-1",
		"sa-east-1",
	} {
		endpoints = append(endpoints, &model.Endpoint{
			Key:   e,
			Value: fmt.Sprintf("https://%s.amazonaws.com", e),
		})
	}
	info := &model.AdapterPattern{
		EndpointPattern: &model.EndpointPattern{
			EndpointType: model.EndpointPatternTypeList,
			Endpoints:    endpoints,
		},
	}
	return info
}

// HealthCheck checks health status of a registry
func (a *adapter) HealthCheck() (string, error) {
	if err := a.Ping(); err != nil {
		log.Errorf("failed to ping registry %s: %v", a.registry.URL, err)
		return model.Unhealthy, nil
	}
	return model.Healthy, nil
}

// PrepareForPush nothing need to do.
func (a *adapter) PrepareForPush(resources []*model.Resource) error {
	for _, resource := range resources {
		if resource == nil {
			return errors.New("the resource cannot be nil")
		}
		if resource.Metadata == nil {
			return errors.New("the metadata of resource cannot be nil")
		}
		if resource.Metadata.Repository == nil {
			return errors.New("the namespace of resource cannot be nil")
		}
		if len(resource.Metadata.Repository.Name) == 0 {
			return errors.New("the name of the namespace cannot be nil")
		}

		exist, err := a.checkRepository(resource.Metadata.Repository.Name)
		if err != nil {
			return err
		}
		if exist {
			log.Infof("Namespace %s already exist in AWS ECR, skip it.", resource.Metadata.Repository.Name)
			continue
		}
		err = a.createRepository(resource.Metadata.Repository.Name)
		if err != nil {
			return err
		}
	}
	return nil
}

func (a *adapter) checkRepository(repository string) (exists bool, err error) {
	out, err := a.cacheSvc.DescribeRepositories(&awsecrapi.DescribeRepositoriesInput{
		RepositoryNames: []*string{&repository},
	})
	if err != nil {
		if e, ok := err.(awserr.Error); ok {
			if e.Code() == awsecrapi.ErrCodeRepositoryNotFoundException {
				return false, nil
			}
		}
		return false, err
	}
	return len(out.Repositories) > 0, nil
}

func (a *adapter) createRepository(repository string) error {
	_, err := a.cacheSvc.CreateRepository(&awsecrapi.CreateRepositoryInput{
		RepositoryName: &repository,
	})
	if err != nil {
		if e, ok := err.(awserr.Error); ok {
			if e.Code() == awsecrapi.ErrCodeRepositoryAlreadyExistsException {
				return nil
			}
		}
		return err
	}
	return nil
}

// DeleteManifest ...
func (a *adapter) DeleteManifest(repository, reference string) error {
	_, err := a.cacheSvc.BatchDeleteImage(&awsecrapi.BatchDeleteImageInput{
		RepositoryName: &repository,
		ImageIds:       []*awsecrapi.ImageIdentifier{{ImageTag: &reference}},
	})
	return err
}
<commit_msg>fix aws adapter url<commit_after>// Copyright Project Harbor 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 awsecr

import (
	"errors"
	"fmt"
	"regexp"

	"github.com/aws/aws-sdk-go/aws/awserr"
	awsecrapi "github.com/aws/aws-sdk-go/service/ecr"
	"github.com/goharbor/harbor/src/lib/log"
	adp "github.com/goharbor/harbor/src/pkg/reg/adapter"
	"github.com/goharbor/harbor/src/pkg/reg/adapter/native"
	"github.com/goharbor/harbor/src/pkg/reg/model"
)

const (
	regionPattern = "https://(?:api|\\d+\\.dkr)\\.ecr\\.([\\w\\-]+)\\.amazonaws\\.com"
)

var (
	regionRegexp = regexp.MustCompile(regionPattern)
)

func init() {
	if err := adp.RegisterFactory(model.RegistryTypeAwsEcr, new(factory)); err != nil {
		log.Errorf("failed to register factory for %s: %v", model.RegistryTypeAwsEcr, err)
		return
	}
	log.Infof("the factory for adapter %s registered", model.RegistryTypeAwsEcr)
}

func newAdapter(registry *model.Registry) (*adapter, error) {
	region, err := parseRegion(registry.URL)
	if err != nil {
		return nil, err
	}
	svc, err := getAwsSvc(
		region, registry.Credential.AccessKey, registry.Credential.AccessSecret, registry.Insecure, nil)
	if err != nil {
		return nil, err
	}
	authorizer := NewAuth(registry.Credential.AccessKey, svc)
	return &adapter{
		registry: registry,
		Adapter:  native.NewAdapterWithAuthorizer(registry, authorizer),
		cacheSvc: svc,
	}, nil
}

func parseRegion(url string) (string, error) {
	rs := regionRegexp.FindStringSubmatch(url)
	if rs == nil {
		return "", errors.New("bad aws url")
	}
	return rs[1], nil
}

type factory struct {
}

// Create ...
func (f *factory) Create(r *model.Registry) (adp.Adapter, error) {
	return newAdapter(r)
}

// AdapterPattern ...
func (f *factory) AdapterPattern() *model.AdapterPattern {
	return getAdapterInfo()
}

var (
	_ adp.Adapter          = (*adapter)(nil)
	_ adp.ArtifactRegistry = (*adapter)(nil)
)

type adapter struct {
	*native.Adapter
	registry *model.Registry
	cacheSvc *awsecrapi.ECR
}

func (*adapter) Info() (info *model.RegistryInfo, err error) {
	return &model.RegistryInfo{
		Type: model.RegistryTypeAwsEcr,
		SupportedResourceTypes: []string{
			model.ResourceTypeImage,
		},
		SupportedResourceFilters: []*model.FilterStyle{
			{
				Type:  model.FilterTypeName,
				Style: model.FilterStyleTypeText,
			},
			{
				Type:  model.FilterTypeTag,
				Style: model.FilterStyleTypeText,
			},
		},
		SupportedTriggers: []string{
			model.TriggerTypeManual,
			model.TriggerTypeScheduled,
		},
	}, nil
}

func getAdapterInfo() *model.AdapterPattern {
	var endpoints []*model.Endpoint
	// https://docs.aws.amazon.com/general/latest/gr/rande.html#regional-endpoints
	for _, e := range []string{
		"us-east-2",
		"us-east-1",
		"us-west-1",
		"us-west-2",
		"af-south-1",
		"ap-east-1",
		"ap-south-1",
		"ap-northeast-3",
		"ap-northeast-2",
		"ap-southeast-1",
		"ap-southeast-2",
		"ap-northeast-1",
		"ca-central-1",
		"cn-north-1",
		"cn-northwest-1",
		"eu-central-1",
		"eu-west-1",
		"eu-west-2",
		"eu-south-1",
		"eu-west-3",
		"eu-north-1",
		"me-south-1",
		"sa-east-1",
	} {
		endpoints = append(endpoints, &model.Endpoint{
			Key:   e,
			Value: fmt.Sprintf("https://api.ecr.%s.amazonaws.com", e),
		})
	}
	info := &model.AdapterPattern{
		EndpointPattern: &model.EndpointPattern{
			EndpointType: model.EndpointPatternTypeList,
			Endpoints:    endpoints,
		},
	}
	return info
}

// HealthCheck checks health status of a registry
func (a *adapter) HealthCheck() (string, error) {
	if err := a.Ping(); err != nil {
		log.Errorf("failed to ping registry %s: %v", a.registry.URL, err)
		return model.Unhealthy, nil
	}
	return model.Healthy, nil
}

// PrepareForPush nothing need to do.
func (a *adapter) PrepareForPush(resources []*model.Resource) error {
	for _, resource := range resources {
		if resource == nil {
			return errors.New("the resource cannot be nil")
		}
		if resource.Metadata == nil {
			return errors.New("the metadata of resource cannot be nil")
		}
		if resource.Metadata.Repository == nil {
			return errors.New("the namespace of resource cannot be nil")
		}
		if len(resource.Metadata.Repository.Name) == 0 {
			return errors.New("the name of the namespace cannot be nil")
		}

		exist, err := a.checkRepository(resource.Metadata.Repository.Name)
		if err != nil {
			return err
		}
		if exist {
			log.Infof("Namespace %s already exist in AWS ECR, skip it.", resource.Metadata.Repository.Name)
			continue
		}
		err = a.createRepository(resource.Metadata.Repository.Name)
		if err != nil {
			return err
		}
	}
	return nil
}

func (a *adapter) checkRepository(repository string) (exists bool, err error) {
	out, err := a.cacheSvc.DescribeRepositories(&awsecrapi.DescribeRepositoriesInput{
		RepositoryNames: []*string{&repository},
	})
	if err != nil {
		if e, ok := err.(awserr.Error); ok {
			if e.Code() == awsecrapi.ErrCodeRepositoryNotFoundException {
				return false, nil
			}
		}
		return false, err
	}
	return len(out.Repositories) > 0, nil
}

func (a *adapter) createRepository(repository string) error {
	_, err := a.cacheSvc.CreateRepository(&awsecrapi.CreateRepositoryInput{
		RepositoryName: &repository,
	})
	if err != nil {
		if e, ok := err.(awserr.Error); ok {
			if e.Code() == awsecrapi.ErrCodeRepositoryAlreadyExistsException {
				return nil
			}
		}
		return err
	}
	return nil
}

// DeleteManifest ...
func (a *adapter) DeleteManifest(repository, reference string) error {
	_, err := a.cacheSvc.BatchDeleteImage(&awsecrapi.BatchDeleteImageInput{
		RepositoryName: &repository,
		ImageIds:       []*awsecrapi.ImageIdentifier{{ImageTag: &reference}},
	})
	return err
}
<|endoftext|>


<commit_before>package cgo

import (
	"github.com/stretchr/testify/assert"
	"go/ast"
	"testing"
)

func TestSingleImport(t *testing.T) {
	subject := Imports("C")
	assert.NotNil(t, subject)
	genDecl := subject
	assert.Equal(t, 1, len(genDecl.Specs))
	importSpec := genDecl.Specs[0].(*ast.ImportSpec)
	assert.Equal(t, "\"C\"", importSpec.Path.Value)
}

func TestMultipleImports(t *testing.T) {
	subject := Imports("foo", "bar", "baz")
	assert.NotNil(t, subject)
	genDecl := subject
	assert.Equal(t, 3, len(genDecl.Specs))
}
<commit_msg>clean import test<commit_after>package cgo

import (
	"github.com/stretchr/testify/assert"
	"go/ast"
	"testing"
)

func TestSingleImport(t *testing.T) {
	subject := Imports("C")
	assert.NotNil(t, subject)
	assert.Equal(t, 1, len(subject.Specs))
	importSpec := subject.Specs[0].(*ast.ImportSpec)
	assert.Equal(t, "\"C\"", importSpec.Path.Value)
}

func TestMultipleImports(t *testing.T) {
	subject := Imports("foo", "bar", "baz")
	assert.NotNil(t, subject)
	assert.Equal(t, 3, len(subject.Specs))
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	stdlog "log"
	"os"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/metrics/prometheus"
	_ "github.com/lib/pq"
	"github.com/piotrkowalczuk/charon"
	"github.com/piotrkowalczuk/mnemosyne"
	"github.com/piotrkowalczuk/sklog"
	stdprometheus "github.com/prometheus/client_golang/prometheus"
	"google.golang.org/grpc"
)

const (
	loggerAdapterStdOut = "stdout"
	loggerFormatJSON    = "json"
	loggerFormatHumane  = "humane"
	loggerFormatLogFmt  = "logfmt"
)

func initLogger(adapter, format string, level int, context ...interface{}) log.Logger {
	var l log.Logger

	if adapter != loggerAdapterStdOut {
		stdlog.Fatal("service: unsupported logger adapter")
	}

	switch format {
	case loggerFormatHumane:
		l = sklog.NewHumaneLogger(os.Stdout, sklog.DefaultHTTPFormatter)
	case loggerFormatJSON:
		l = log.NewJSONLogger(os.Stdout)
	case loggerFormatLogFmt:
		l = log.NewLogfmtLogger(os.Stdout)
	default:
		stdlog.Fatal("charond: unsupported logger format")
	}

	l = log.NewContext(l).With(context...)

	sklog.Info(l, "logger has been initialized successfully", "adapter", adapter, "format", format, "level", level)

	return l
}

func initPostgres(connectionString string, logger log.Logger) *sql.DB {
	postgres, err := sql.Open("postgres", connectionString)
	if err != nil {
		sklog.Fatal(logger, err)
	}

	err = setupDatabase(postgres)
	if err != nil {
		sklog.Fatal(logger, err)
	}
	sklog.Info(logger, "postgres connection has been established", "address", connectionString)

	return postgres
}

func initMnemosyne(address string, logger log.Logger) (*grpc.ClientConn, mnemosyne.Mnemosyne) {
	conn, err := grpc.Dial(address, grpc.WithInsecure())
	if err != nil {
		sklog.Fatal(logger, err, "address", address)
	}

	sklog.Info(logger, "rpc connection to mnemosyne has been established", "address", address)

	return conn, mnemosyne.New(conn, mnemosyne.MnemosyneOpts{})
}

func initPasswordHasher(cost int, logger log.Logger) charon.PasswordHasher {
	bh, err := charon.NewBCryptPasswordHasher(cost)
	if err != nil {
		sklog.Fatal(logger, err)
	}

	return bh
}

func initPermissionRegistry(r PermissionRepository, permissions charon.Permissions, logger log.Logger) (pr PermissionRegistry) {
	pr = newPermissionRegistry(r)
	created, untouched, removed, err := pr.Register(permissions)
	if err != nil {
		sklog.Fatal(logger, err)
	}

	sklog.Info(logger, "charon permissions has been registered", "created", created, "untouched", untouched, "removed", removed)

	return
}

const (
	monitoringEnginePrometheus = "prometheus"
)

func initMonitoring(fn func() (monitoring, error), logger log.Logger) monitoring {
	m, err := fn()
	if err != nil {
		sklog.Fatal(logger, err)
	}

	return m
}

func initPrometheus(namespace, subsystem string, constLabels stdprometheus.Labels) func() (monitoring, error) {
	return func() (monitoring, error) {
		rpcRequests := prometheus.NewCounter(
			stdprometheus.CounterOpts{
				Namespace:   namespace,
				Subsystem:   subsystem,
				Name:        "rpc_requests_total",
				Help:        "Total number of RPC requests made.",
				ConstLabels: constLabels,
			},
			monitoringRPCLabels,
		)
		rpcErrors := prometheus.NewCounter(
			stdprometheus.CounterOpts{
				Namespace:   namespace,
				Subsystem:   subsystem,
				Name:        "rpc_errors_total",
				Help:        "Total number of errors that happen during RPC calles.",
				ConstLabels: constLabels,
			},
			monitoringRPCLabels,
		)

		postgresQueries := prometheus.NewCounter(
			stdprometheus.CounterOpts{
				Namespace:   namespace,
				Subsystem:   subsystem,
				Name:        "postgres_queries_total",
				Help:        "Total number of SQL queries made.",
				ConstLabels: constLabels,
			},
			monitoringPostgresLabels,
		)
		postgresErrors := prometheus.NewCounter(
			stdprometheus.CounterOpts{
				Namespace:   namespace,
				Subsystem:   subsystem,
				Name:        "postgres_errors_total",
				Help:        "Total number of errors that happen during SQL queries.",
				ConstLabels: constLabels,
			},
			monitoringPostgresLabels,
		)

		return monitoring{
			enabled: true,
			rpc: monitoringRPC{
				enabled:  true,
				requests: rpcRequests,
				errors:   rpcErrors,
			},
			postgres: monitoringPostgres{
				enabled: true,
				queries: postgresQueries,
				errors:  postgresErrors,
			},
		}, nil
	}
}
<commit_msg>custom context error func for sklog<commit_after>package main

import (
	"database/sql"
	stdlog "log"
	"os"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/metrics/prometheus"
	_ "github.com/lib/pq"
	"github.com/piotrkowalczuk/charon"
	"github.com/piotrkowalczuk/mnemosyne"
	"github.com/piotrkowalczuk/sklog"
	"github.com/piotrkowalczuk/sklog/ctxgrpc"
	stdprometheus "github.com/prometheus/client_golang/prometheus"
	"google.golang.org/grpc"
)

const (
	loggerAdapterStdOut = "stdout"
	loggerFormatJSON    = "json"
	loggerFormatHumane  = "humane"
	loggerFormatLogFmt  = "logfmt"
)

func initLogger(adapter, format string, level int, context ...interface{}) log.Logger {
	sklog.SetContextErrorFunc(ctxgrpc.NewContextErrorGeneric)
	var l log.Logger

	if adapter != loggerAdapterStdOut {
		stdlog.Fatal("service: unsupported logger adapter")
	}

	switch format {
	case loggerFormatHumane:
		l = sklog.NewHumaneLogger(os.Stdout, sklog.DefaultHTTPFormatter)
	case loggerFormatJSON:
		l = log.NewJSONLogger(os.Stdout)
	case loggerFormatLogFmt:
		l = log.NewLogfmtLogger(os.Stdout)
	default:
		stdlog.Fatal("charond: unsupported logger format")
	}

	l = log.NewContext(l).With(context...)

	sklog.Info(l, "logger has been initialized successfully", "adapter", adapter, "format", format, "level", level)

	return l
}

func initPostgres(connectionString string, logger log.Logger) *sql.DB {
	postgres, err := sql.Open("postgres", connectionString)
	if err != nil {
		sklog.Fatal(logger, err)
	}

	err = setupDatabase(postgres)
	if err != nil {
		sklog.Fatal(logger, err)
	}
	sklog.Info(logger, "postgres connection has been established", "address", connectionString)

	return postgres
}

func initMnemosyne(address string, logger log.Logger) (*grpc.ClientConn, mnemosyne.Mnemosyne) {
	conn, err := grpc.Dial(address, grpc.WithInsecure())
	if err != nil {
		sklog.Fatal(logger, err, "address", address)
	}

	sklog.Info(logger, "rpc connection to mnemosyne has been established", "address", address)

	return conn, mnemosyne.New(conn, mnemosyne.MnemosyneOpts{})
}

func initPasswordHasher(cost int, logger log.Logger) charon.PasswordHasher {
	bh, err := charon.NewBCryptPasswordHasher(cost)
	if err != nil {
		sklog.Fatal(logger, err)
	}

	return bh
}

func initPermissionRegistry(r PermissionRepository, permissions charon.Permissions, logger log.Logger) (pr PermissionRegistry) {
	pr = newPermissionRegistry(r)
	created, untouched, removed, err := pr.Register(permissions)
	if err != nil {
		sklog.Fatal(logger, err)
	}

	sklog.Info(logger, "charon permissions has been registered", "created", created, "untouched", untouched, "removed", removed)

	return
}

const (
	monitoringEnginePrometheus = "prometheus"
)

func initMonitoring(fn func() (monitoring, error), logger log.Logger) monitoring {
	m, err := fn()
	if err != nil {
		sklog.Fatal(logger, err)
	}

	return m
}

func initPrometheus(namespace, subsystem string, constLabels stdprometheus.Labels) func() (monitoring, error) {
	return func() (monitoring, error) {
		rpcRequests := prometheus.NewCounter(
			stdprometheus.CounterOpts{
				Namespace:   namespace,
				Subsystem:   subsystem,
				Name:        "rpc_requests_total",
				Help:        "Total number of RPC requests made.",
				ConstLabels: constLabels,
			},
			monitoringRPCLabels,
		)
		rpcErrors := prometheus.NewCounter(
			stdprometheus.CounterOpts{
				Namespace:   namespace,
				Subsystem:   subsystem,
				Name:        "rpc_errors_total",
				Help:        "Total number of errors that happen during RPC calles.",
				ConstLabels: constLabels,
			},
			monitoringRPCLabels,
		)

		postgresQueries := prometheus.NewCounter(
			stdprometheus.CounterOpts{
				Namespace:   namespace,
				Subsystem:   subsystem,
				Name:        "postgres_queries_total",
				Help:        "Total number of SQL queries made.",
				ConstLabels: constLabels,
			},
			monitoringPostgresLabels,
		)
		postgresErrors := prometheus.NewCounter(
			stdprometheus.CounterOpts{
				Namespace:   namespace,
				Subsystem:   subsystem,
				Name:        "postgres_errors_total",
				Help:        "Total number of errors that happen during SQL queries.",
				ConstLabels: constLabels,
			},
			monitoringPostgresLabels,
		)

		return monitoring{
			enabled: true,
			rpc: monitoringRPC{
				enabled:  true,
				requests: rpcRequests,
				errors:   rpcErrors,
			},
			postgres: monitoringPostgres{
				enabled: true,
				queries: postgresQueries,
				errors:  postgresErrors,
			},
		}, nil
	}
}
<|endoftext|>


<commit_before>package checker

import (
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/asciimoo/privacyscore/result"
	"github.com/asciimoo/privacyscore/utils"
)

const RESOURCE_LIMIT = 64
const TIMEOUT = 5
const USER_AGENT = "Mozilla/5.0 (compatible) PrivacyScore Checker v0.1.0"
const maxResponseBodySize = 1024 * 1024 * 5

var mutex = &sync.Mutex{}

type Checker interface {
	Check(*CheckJob, *PageInfo)
}

type CheckJob struct {
	sync.RWMutex
	Result    *result.Result
	Resources map[string]*PageInfo
	Chan      chan bool
}

type PageInfo struct {
	ResponseBody   []byte
	ContentType    string
	StatusCode     int
	URL            *url.URL
	OriginalURL    *url.URL
	Cookies        []*http.Cookie
	Domain         string
	ResponseHeader *http.Header
}

var checkers []Checker = []Checker{
	&CookieChecker{},
	&HTMLChecker{},
	&HTTPSChecker{},
	&SecureHeaderChecker{},
	&CSSChecker{},
}

func Run(URL string) (*CheckJob, error) {
	if !strings.HasPrefix(URL, "http://") && !strings.HasPrefix(URL, "https://") {
		URL = "http://" + URL
	}
	c := newCheckJob(URL)
	finishedResources := 0
	errorCount := 0
	c.CheckURL(URL)
	for finishedResources != len(c.Resources) && finishedResources < RESOURCE_LIMIT {
		select {
		case ret := <-c.Chan:
			if ret == false {
				errorCount += 1
			}
			finishedResources += 1
		}
	}
	if finishedResources == 0 || (errorCount > 0 && errorCount == finishedResources) {
		return c, errors.New("Could not download host")
	}
	if r, found := c.Resources[URL]; found && r != nil {
		c.Result.BaseURL = c.Resources[URL].URL.String()
	}
	return c, nil
}

func newCheckJob(URL string) *CheckJob {
	return &CheckJob{
		Result:    result.New(URL),
		Resources: make(map[string]*PageInfo),
		Chan:      make(chan bool, RESOURCE_LIMIT),
	}
}

func (c *CheckJob) CheckURL(URL string) {
	if URL == "" {
		return
	}
	// URL already added
	c.RLock()
	r, found := c.Resources[URL]
	c.RUnlock()
	if found && r != nil {
		return
	}
	// resource limit exceeded
	if len(c.Resources) >= RESOURCE_LIMIT {
		return
	}
	var p *PageInfo
	c.Lock()
	c.Resources[URL] = p
	c.Unlock()
	go func() {
		r, err := fetchURL(URL)
		if err != nil {
			c.Result.AddError(err)
			c.Chan <- false
			return
		}
		var body []byte
		contentType := r.Header.Get("Content-Type")
		if r.StatusCode == 200 && (strings.Contains(contentType, "text/") || strings.Contains(contentType, "html")) {
			body, err = ioutil.ReadAll(io.LimitReader(r.Body, maxResponseBodySize))
			if err != nil {
				c.Result.AddError(err)
			}
		} else {
			body = []byte{}
		}
		r.Body.Close()

		u, _ := url.Parse(URL)
		p = &PageInfo{
			body,
			r.Header.Get("Content-Type"),
			r.StatusCode,
			r.Request.URL,
			u,
			r.Cookies(),
			utils.CropSubdomains(r.Request.URL.Host),
			&r.Header,
		}
		c.Lock()
		c.Resources[URL] = p
		c.Unlock()
		for _, ch := range checkers {
			ch.Check(c, p)
		}
		c.Chan <- true
	}()
	return
}

func fetchURL(URL string) (*http.Response, error) {
	client := http.Client{
		Timeout: time.Duration(TIMEOUT * time.Second),
		Transport: &http.Transport{
			DisableKeepAlives:     true,
			ResponseHeaderTimeout: TIMEOUT * time.Second,
		},
	}
	req, err := http.NewRequest("GET", URL, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", USER_AGENT)
	response, err := client.Do(req)
	return response, err
}
<commit_msg>[fix] locking<commit_after>package checker

import (
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/asciimoo/privacyscore/result"
	"github.com/asciimoo/privacyscore/utils"
)

const RESOURCE_LIMIT = 64
const TIMEOUT = 5
const USER_AGENT = "Mozilla/5.0 (compatible) PrivacyScore Checker v0.1.0"
const maxResponseBodySize = 1024 * 1024 * 5

var mutex = &sync.Mutex{}

type Checker interface {
	Check(*CheckJob, *PageInfo)
}

type CheckJob struct {
	sync.RWMutex
	Result    *result.Result
	Resources map[string]*PageInfo
	Chan      chan bool
}

type PageInfo struct {
	ResponseBody   []byte
	ContentType    string
	StatusCode     int
	URL            *url.URL
	OriginalURL    *url.URL
	Cookies        []*http.Cookie
	Domain         string
	ResponseHeader *http.Header
}

var checkers []Checker = []Checker{
	&CookieChecker{},
	&HTMLChecker{},
	&HTTPSChecker{},
	&SecureHeaderChecker{},
	&CSSChecker{},
}

func Run(URL string) (*CheckJob, error) {
	if !strings.HasPrefix(URL, "http://") && !strings.HasPrefix(URL, "https://") {
		URL = "http://" + URL
	}
	c := newCheckJob(URL)
	finishedResources := 0
	errorCount := 0
	c.CheckURL(URL)
	for finishedResources != len(c.Resources) && finishedResources < RESOURCE_LIMIT {
		select {
		case ret := <-c.Chan:
			if ret == false {
				errorCount += 1
			}
			finishedResources += 1
		}
	}
	if finishedResources == 0 || (errorCount > 0 && errorCount == finishedResources) {
		return c, errors.New("Could not download host")
	}
	c.RLock()
	r, found := c.Resources[URL]
	c.RUnlock()
	if found && r != nil {
		c.Result.BaseURL = r.URL.String()
	}
	return c, nil
}

func newCheckJob(URL string) *CheckJob {
	return &CheckJob{
		Result:    result.New(URL),
		Resources: make(map[string]*PageInfo),
		Chan:      make(chan bool, RESOURCE_LIMIT),
	}
}

func (c *CheckJob) CheckURL(URL string) {
	if URL == "" {
		return
	}
	// URL already added
	c.RLock()
	r, found := c.Resources[URL]
	c.RUnlock()
	if found && r != nil {
		return
	}
	// resource limit exceeded
	if len(c.Resources) >= RESOURCE_LIMIT {
		return
	}
	var p *PageInfo
	c.Lock()
	c.Resources[URL] = p
	c.Unlock()
	go func() {
		r, err := fetchURL(URL)
		if err != nil {
			c.Result.AddError(err)
			c.Chan <- false
			return
		}
		var body []byte
		contentType := r.Header.Get("Content-Type")
		if r.StatusCode == 200 && (strings.Contains(contentType, "text/") || strings.Contains(contentType, "html")) {
			body, err = ioutil.ReadAll(io.LimitReader(r.Body, maxResponseBodySize))
			if err != nil {
				c.Result.AddError(err)
			}
		} else {
			body = []byte{}
		}
		r.Body.Close()

		u, _ := url.Parse(URL)
		p = &PageInfo{
			body,
			r.Header.Get("Content-Type"),
			r.StatusCode,
			r.Request.URL,
			u,
			r.Cookies(),
			utils.CropSubdomains(r.Request.URL.Host),
			&r.Header,
		}
		c.Lock()
		c.Resources[URL] = p
		c.Unlock()
		for _, ch := range checkers {
			ch.Check(c, p)
		}
		c.Chan <- true
	}()
	return
}

func fetchURL(URL string) (*http.Response, error) {
	client := http.Client{
		Timeout: time.Duration(TIMEOUT * time.Second),
		Transport: &http.Transport{
			DisableKeepAlives:     true,
			ResponseHeaderTimeout: TIMEOUT * time.Second,
		},
	}
	req, err := http.NewRequest("GET", URL, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", USER_AGENT)
	response, err := client.Do(req)
	return response, err
}
<|endoftext|>


<commit_before>package server

import (
	"bufio"
	"bytes"
	"fmt"
	"hash/fnv"
	"io"
	"math/rand"
	"os"
	"path"
	"reflect"
	"strings"
	"sync"
	"testing"

	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/require"
	"github.com/pachyderm/pachyderm/src/client/pps"
	versionlib "github.com/pachyderm/pachyderm/src/client/version"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
	tu "github.com/pachyderm/pachyderm/src/server/pkg/testutil"
	"github.com/pachyderm/pachyderm/src/server/pkg/workload"

	"github.com/golang/snappy"
)

const (
	KB = 1024
	MB = 1024 * KB
)

var pachClient *client.APIClient
var getPachClientOnce sync.Once

func getPachClient(t testing.TB) *client.APIClient {
	getPachClientOnce.Do(func() {
		var err error
		if addr := os.Getenv("PACHD_PORT_650_TCP_ADDR"); addr != "" {
			pachClient, err = client.NewInCluster()
		} else {
			pachClient, err = client.NewOnUserMachine(false, false, "user")
		}
		require.NoError(t, err)
	})
	return pachClient
}

func collectCommitInfos(t testing.TB, commitInfoIter client.CommitInfoIterator) []*pfs.CommitInfo {
	var commitInfos []*pfs.CommitInfo
	for {
		commitInfo, err := commitInfoIter.Next()
		if err == io.EOF {
			return commitInfos
		}
		require.NoError(t, err)
		commitInfos = append(commitInfos, commitInfo)
	}
}

// TODO(msteffen) equivalent to funciton in src/server/auth/server/admin_test.go.
// These should be unified.
func RepoInfoToName(repoInfo interface{}) interface{} {
	return repoInfo.(*pfs.RepoInfo).Repo.Name
}

// testExtractRestored effectively implements both TestExtractRestoreObjects
// TestExtractRestoreNoObjects, as their logic is mostly the same
func testExtractRestore(t *testing.T, testObjects bool) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestoreObjects-in-")
	require.NoError(t, c.CreateRepo(dataRepo))

	// Create input data
	nCommits := 2
	r := rand.New(rand.NewSource(45))
	fileHashes := make([]string, 0, nCommits)
	for i := 0; i < nCommits; i++ {
		hash := fnv.New64a()
		fileContent := workload.RandString(r, 40*MB)
		_, err := c.PutFile(dataRepo, "master", fmt.Sprintf("file-%d", i),
			io.TeeReader(strings.NewReader(fileContent), hash))
		require.NoError(t, err)
		fileHashes = append(fileHashes, string(hash.Sum(nil)))
	}

	// Create test pipelines
	numPipelines := 3
	var input, pipeline string
	input = dataRepo
	for i := 0; i < numPipelines; i++ {
		pipeline = tu.UniqueString(fmt.Sprintf("TestExtractRestoreObjects-P%d-", i))
		require.NoError(t, c.CreatePipeline(
			pipeline,
			"",
			[]string{"bash"},
			[]string{
				fmt.Sprintf("cp /pfs/%s/* /pfs/out/", input),
			},
			&pps.ParallelismSpec{
				Constant: 1,
			},
			client.NewPFSInput(input, "/*"),
			"",
			false,
		))
		input = pipeline
	}

	// Wait for pipelines to process input data
	commitIter, err := c.FlushCommit([]*pfs.Commit{client.NewCommit(dataRepo, "master")}, nil)
	require.NoError(t, err)
	commitInfos := collectCommitInfos(t, commitIter)
	require.Equal(t, numPipelines, len(commitInfos))

	// Extract existing cluster state
	ops, err := c.ExtractAll(testObjects)
	require.NoError(t, err)

	// Delete existing metadata
	require.NoError(t, c.DeleteAll())

	if testObjects {
		// Delete existing objects
		require.NoError(t, c.GarbageCollect(10000))
	}

	// Restore metadata and possibly objects
	require.NoError(t, c.Restore(ops))

	// Wait for re-created pipelines to process recreated input data
	commitIter, err = c.FlushCommit([]*pfs.Commit{client.NewCommit(dataRepo, "master")}, nil)
	require.NoError(t, err)
	commitInfos = collectCommitInfos(t, commitIter)
	require.Equal(t, numPipelines, len(commitInfos))

	// Make sure recreated jobs all succeeded
	jis, err := c.ListJob("", nil, nil) // make sure jobs all succeeded
	require.NoError(t, err)
	for _, ji := range jis {
		require.Equal(t, pps.JobState_JOB_SUCCESS, ji.State)
	}

	// Make sure all branches were recreated
	bis, err := c.ListBranch(dataRepo)
	require.NoError(t, err)
	require.Equal(t, 1, len(bis))

	// Check input data
	// This check uses a backoff because sometimes GetFile causes pachd to OOM
	var restoredFileHashes []string
	require.NoError(t, backoff.Retry(func() error {
		restoredFileHashes = make([]string, 0, nCommits)
		for i := 0; i < nCommits; i++ {
			hash := fnv.New64a()
			err := c.GetFile(dataRepo, "master", fmt.Sprintf("file-%d", i), 0, 0, hash)
			if err != nil {
				return err
			}
			restoredFileHashes = append(restoredFileHashes, string(hash.Sum(nil)))
		}
		return nil
	}, backoff.NewTestingBackOff()))
	require.ElementsEqual(t, fileHashes, restoredFileHashes)

	// Check output data
	// This check uses a backoff because sometimes GetFile causes pachd to OOM
	require.NoError(t, backoff.Retry(func() error {
		restoredFileHashes = make([]string, 0, nCommits)
		for i := 0; i < nCommits; i++ {
			hash := fnv.New64a()
			err := c.GetFile(pipeline, "master", fmt.Sprintf("file-%d", i), 0, 0, hash)
			if err != nil {
				return err
			}
			restoredFileHashes = append(restoredFileHashes, string(hash.Sum(nil)))
		}
		return nil
	}, backoff.NewTestingBackOff()))
	require.ElementsEqual(t, fileHashes, restoredFileHashes)
}

// TestExtractRestoreNoObjects tests extraction and restoration in the case
// where existing objects are re-used (common for cloud deployments, as objects
// are stored outside of kubernetes, in object store)
func TestExtractRestoreNoObjects(t *testing.T) {
	testExtractRestore(t, false)
}

// TestExtractRestoreObjects tests extraction and restoration of objects. Note
// that since 1.8, only data in input repos is referenced by objects, so this
// tests extracting/restoring an input repo.
func TestExtractRestoreObjects(t *testing.T) {
	testExtractRestore(t, true)
}

func TestExtractRestoreHeadlessBranches(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestore_data")
	require.NoError(t, c.CreateRepo(dataRepo))

	// create a headless branch
	require.NoError(t, c.CreateBranch(dataRepo, "headless", "", nil))

	ops, err := c.ExtractAll(false)
	require.NoError(t, err)
	require.NoError(t, c.DeleteAll())
	require.NoError(t, c.Restore(ops))

	bis, err := c.ListBranch(dataRepo)
	require.NoError(t, err)
	require.Equal(t, 1, len(bis))
	require.Equal(t, "headless", bis[0].Branch.Name)
}

func TestExtractVersion(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestore_data")
	require.NoError(t, c.CreateRepo(dataRepo))

	r := rand.New(rand.NewSource(45))
	_, err := c.PutFile(dataRepo, "master", "file", strings.NewReader(workload.RandString(r, 40*MB)))
	require.NoError(t, err)

	pipeline := tu.UniqueString("TestExtractRestore")
	require.NoError(t, c.CreatePipeline(
		pipeline,
		"",
		[]string{"bash"},
		[]string{
			fmt.Sprintf("cp /pfs/%s/* /pfs/out/", dataRepo),
		},
		&pps.ParallelismSpec{
			Constant: 1,
		},
		client.NewPFSInput(dataRepo, "/*"),
		"",
		false,
	))

	ops, err := c.ExtractAll(false)
	require.NoError(t, err)
	require.True(t, len(ops) > 0)

	// Check that every Op looks right; the version set matches pachd's version
	for _, op := range ops {
		opV := reflect.ValueOf(op).Elem()
		var versions, nonemptyVersions int
		for i := 0; i < opV.NumField(); i++ {
			fDesc := opV.Type().Field(i)
			if !strings.HasPrefix(fDesc.Name, "Op") {
				continue
			}
			versions++
			if strings.HasSuffix(fDesc.Name,
				fmt.Sprintf("%d_%d", versionlib.MajorVersion, versionlib.MinorVersion)) {
				require.False(t, opV.Field(i).IsNil())
				nonemptyVersions++
			} else {
				require.True(t, opV.Field(i).IsNil())
			}
		}
		require.Equal(t, 1, nonemptyVersions)
		require.True(t, versions > 1)
	}
}

func TestMigrateFrom1_7(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	// Clear pachyderm cluster (so that next cluster starts up in a clean environment)
	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	// Restore dumped metadata (now that objects are present)
	md, err := os.Open(path.Join(os.Getenv("GOPATH"),
		"src/github.com/pachyderm/pachyderm/etc/testing/migration/1_7/sort.metadata"))
	require.NoError(t, err)
	require.NoError(t, c.RestoreReader(snappy.NewReader(md)))
	require.NoError(t, md.Close())

	// Wait for final imported commit to be processed
	commitIter, err := c.FlushCommit([]*pfs.Commit{client.NewCommit("left", "master")}, nil)
	require.NoError(t, err)
	commitInfos := collectCommitInfos(t, commitIter)
	// filter-left and filter-right both compute a join of left and
	// right--depending on when the final commit to 'left' was added, it may have
	// been processed multiple times (should be n * 3, as there are 3 pipelines)
	require.True(t, len(commitInfos) >= 3)

	// Inspect input
	commits, err := c.ListCommit("left", "master", "", 0)
	require.NoError(t, err)
	require.Equal(t, 3, len(commits))
	commits, err = c.ListCommit("right", "master", "", 0)
	require.NoError(t, err)
	require.Equal(t, 3, len(commits))

	// Inspect output
	repos, err := c.ListRepo()
	require.NoError(t, err)
	require.ElementsEqualUnderFn(t,
		[]string{"left", "right", "copy", "sort"},
		repos, RepoInfoToName)

	// make sure all numbers 0-99 are in /nums
	var buf bytes.Buffer
	require.NoError(t, c.GetFile("sort", "master", "/nums", 0, 0, &buf))
	s := bufio.NewScanner(&buf)
	numbers := make(map[string]struct{})
	for s.Scan() {
		numbers[s.Text()] = struct{}{}
	}
	require.Equal(t, 100, len(numbers)) // job processed all inputs

	// Confirm stats commits are present
	commits, err = c.ListCommit("sort", "stats", "", 0)
	require.NoError(t, err)
	require.Equal(t, 6, len(commits))
}

func int64p(i int64) *int64 {
	return &i
}
<commit_msg>Prevent race in test<commit_after>package server

import (
	"bufio"
	"bytes"
	"fmt"
	"hash/fnv"
	"io"
	"math/rand"
	"os"
	"path"
	"reflect"
	"strings"
	"sync"
	"testing"

	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/require"
	"github.com/pachyderm/pachyderm/src/client/pps"
	versionlib "github.com/pachyderm/pachyderm/src/client/version"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
	tu "github.com/pachyderm/pachyderm/src/server/pkg/testutil"
	"github.com/pachyderm/pachyderm/src/server/pkg/workload"

	"github.com/golang/snappy"
)

const (
	KB = 1024
	MB = 1024 * KB
)

var pachClient *client.APIClient
var getPachClientOnce sync.Once

func getPachClient(t testing.TB) *client.APIClient {
	getPachClientOnce.Do(func() {
		var err error
		if addr := os.Getenv("PACHD_PORT_650_TCP_ADDR"); addr != "" {
			pachClient, err = client.NewInCluster()
		} else {
			pachClient, err = client.NewOnUserMachine(false, false, "user")
		}
		require.NoError(t, err)
	})
	return pachClient
}

func collectCommitInfos(t testing.TB, commitInfoIter client.CommitInfoIterator) []*pfs.CommitInfo {
	var commitInfos []*pfs.CommitInfo
	for {
		commitInfo, err := commitInfoIter.Next()
		if err == io.EOF {
			return commitInfos
		}
		require.NoError(t, err)
		commitInfos = append(commitInfos, commitInfo)
	}
}

// TODO(msteffen) equivalent to funciton in src/server/auth/server/admin_test.go.
// These should be unified.
func RepoInfoToName(repoInfo interface{}) interface{} {
	return repoInfo.(*pfs.RepoInfo).Repo.Name
}

// testExtractRestored effectively implements both TestExtractRestoreObjects
// TestExtractRestoreNoObjects, as their logic is mostly the same
func testExtractRestore(t *testing.T, testObjects bool) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestoreObjects-in-")
	require.NoError(t, c.CreateRepo(dataRepo))

	// Create input data
	nCommits := 2
	r := rand.New(rand.NewSource(45))
	fileHashes := make([]string, 0, nCommits)
	for i := 0; i < nCommits; i++ {
		hash := fnv.New64a()
		fileContent := workload.RandString(r, 40*MB)
		_, err := c.PutFile(dataRepo, "master", fmt.Sprintf("file-%d", i),
			io.TeeReader(strings.NewReader(fileContent), hash))
		require.NoError(t, err)
		fileHashes = append(fileHashes, string(hash.Sum(nil)))
	}

	// Create test pipelines
	numPipelines := 3
	var input, pipeline string
	input = dataRepo
	for i := 0; i < numPipelines; i++ {
		pipeline = tu.UniqueString(fmt.Sprintf("TestExtractRestoreObjects-P%d-", i))
		require.NoError(t, c.CreatePipeline(
			pipeline,
			"",
			[]string{"bash"},
			[]string{
				fmt.Sprintf("cp /pfs/%s/* /pfs/out/", input),
			},
			&pps.ParallelismSpec{
				Constant: 1,
			},
			client.NewPFSInput(input, "/*"),
			"",
			false,
		))
		input = pipeline
	}

	// Wait for pipelines to process input data
	commitIter, err := c.FlushCommit([]*pfs.Commit{client.NewCommit(dataRepo, "master")}, nil)
	require.NoError(t, err)
	commitInfos := collectCommitInfos(t, commitIter)
	require.Equal(t, numPipelines, len(commitInfos))

	// Extract existing cluster state
	ops, err := c.ExtractAll(testObjects)
	require.NoError(t, err)

	// Delete existing metadata
	require.NoError(t, c.DeleteAll())

	if testObjects {
		// Delete existing objects
		require.NoError(t, c.GarbageCollect(10000))
	}

	// Restore metadata and possibly objects
	require.NoError(t, c.Restore(ops))

	// Wait for re-created pipelines to process recreated input data
	commitIter, err = c.FlushCommit([]*pfs.Commit{client.NewCommit(dataRepo, "master")}, nil)
	require.NoError(t, err)
	commitInfos = collectCommitInfos(t, commitIter)
	require.Equal(t, numPipelines, len(commitInfos))

	// Make sure recreated jobs all succeeded
	jis, err := c.ListJob("", nil, nil) // make sure jobs all succeeded
	require.NoError(t, err)
	for _, ji := range jis {
		// prevent race--we may call listJob between when a job's output commit is
		// closed and when its state is updated
		backoff.Retry(func() error {
			if ji.State.String() == "JOB_MERGING" {
				return fmt.Errorf("output commit is closed but job state hasn't been updated")
			}
			return nil
		}, backoff.NewTestingBackOff())
		// Job must ultimately succeed
		require.Equal(t, "JOB_SUCCESS", ji.State.String())
	}

	// Make sure all branches were recreated
	bis, err := c.ListBranch(dataRepo)
	require.NoError(t, err)
	require.Equal(t, 1, len(bis))

	// Check input data
	// This check uses a backoff because sometimes GetFile causes pachd to OOM
	var restoredFileHashes []string
	require.NoError(t, backoff.Retry(func() error {
		restoredFileHashes = make([]string, 0, nCommits)
		for i := 0; i < nCommits; i++ {
			hash := fnv.New64a()
			err := c.GetFile(dataRepo, "master", fmt.Sprintf("file-%d", i), 0, 0, hash)
			if err != nil {
				return err
			}
			restoredFileHashes = append(restoredFileHashes, string(hash.Sum(nil)))
		}
		return nil
	}, backoff.NewTestingBackOff()))
	require.ElementsEqual(t, fileHashes, restoredFileHashes)

	// Check output data
	// This check uses a backoff because sometimes GetFile causes pachd to OOM
	require.NoError(t, backoff.Retry(func() error {
		restoredFileHashes = make([]string, 0, nCommits)
		for i := 0; i < nCommits; i++ {
			hash := fnv.New64a()
			err := c.GetFile(pipeline, "master", fmt.Sprintf("file-%d", i), 0, 0, hash)
			if err != nil {
				return err
			}
			restoredFileHashes = append(restoredFileHashes, string(hash.Sum(nil)))
		}
		return nil
	}, backoff.NewTestingBackOff()))
	require.ElementsEqual(t, fileHashes, restoredFileHashes)
}

// TestExtractRestoreNoObjects tests extraction and restoration in the case
// where existing objects are re-used (common for cloud deployments, as objects
// are stored outside of kubernetes, in object store)
func TestExtractRestoreNoObjects(t *testing.T) {
	testExtractRestore(t, false)
}

// TestExtractRestoreObjects tests extraction and restoration of objects. Note
// that since 1.8, only data in input repos is referenced by objects, so this
// tests extracting/restoring an input repo.
func TestExtractRestoreObjects(t *testing.T) {
	testExtractRestore(t, true)
}

func TestExtractRestoreHeadlessBranches(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestore_data")
	require.NoError(t, c.CreateRepo(dataRepo))

	// create a headless branch
	require.NoError(t, c.CreateBranch(dataRepo, "headless", "", nil))

	ops, err := c.ExtractAll(false)
	require.NoError(t, err)
	require.NoError(t, c.DeleteAll())
	require.NoError(t, c.Restore(ops))

	bis, err := c.ListBranch(dataRepo)
	require.NoError(t, err)
	require.Equal(t, 1, len(bis))
	require.Equal(t, "headless", bis[0].Branch.Name)
}

func TestExtractVersion(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	dataRepo := tu.UniqueString("TestExtractRestore_data")
	require.NoError(t, c.CreateRepo(dataRepo))

	r := rand.New(rand.NewSource(45))
	_, err := c.PutFile(dataRepo, "master", "file", strings.NewReader(workload.RandString(r, 40*MB)))
	require.NoError(t, err)

	pipeline := tu.UniqueString("TestExtractRestore")
	require.NoError(t, c.CreatePipeline(
		pipeline,
		"",
		[]string{"bash"},
		[]string{
			fmt.Sprintf("cp /pfs/%s/* /pfs/out/", dataRepo),
		},
		&pps.ParallelismSpec{
			Constant: 1,
		},
		client.NewPFSInput(dataRepo, "/*"),
		"",
		false,
	))

	ops, err := c.ExtractAll(false)
	require.NoError(t, err)
	require.True(t, len(ops) > 0)

	// Check that every Op looks right; the version set matches pachd's version
	for _, op := range ops {
		opV := reflect.ValueOf(op).Elem()
		var versions, nonemptyVersions int
		for i := 0; i < opV.NumField(); i++ {
			fDesc := opV.Type().Field(i)
			if !strings.HasPrefix(fDesc.Name, "Op") {
				continue
			}
			versions++
			if strings.HasSuffix(fDesc.Name,
				fmt.Sprintf("%d_%d", versionlib.MajorVersion, versionlib.MinorVersion)) {
				require.False(t, opV.Field(i).IsNil())
				nonemptyVersions++
			} else {
				require.True(t, opV.Field(i).IsNil())
			}
		}
		require.Equal(t, 1, nonemptyVersions)
		require.True(t, versions > 1)
	}
}

func TestMigrateFrom1_7(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	// Clear pachyderm cluster (so that next cluster starts up in a clean environment)
	c := getPachClient(t)
	require.NoError(t, c.DeleteAll())

	// Restore dumped metadata (now that objects are present)
	md, err := os.Open(path.Join(os.Getenv("GOPATH"),
		"src/github.com/pachyderm/pachyderm/etc/testing/migration/1_7/sort.metadata"))
	require.NoError(t, err)
	require.NoError(t, c.RestoreReader(snappy.NewReader(md)))
	require.NoError(t, md.Close())

	// Wait for final imported commit to be processed
	commitIter, err := c.FlushCommit([]*pfs.Commit{client.NewCommit("left", "master")}, nil)
	require.NoError(t, err)
	commitInfos := collectCommitInfos(t, commitIter)
	// filter-left and filter-right both compute a join of left and
	// right--depending on when the final commit to 'left' was added, it may have
	// been processed multiple times (should be n * 3, as there are 3 pipelines)
	require.True(t, len(commitInfos) >= 3)

	// Inspect input
	commits, err := c.ListCommit("left", "master", "", 0)
	require.NoError(t, err)
	require.Equal(t, 3, len(commits))
	commits, err = c.ListCommit("right", "master", "", 0)
	require.NoError(t, err)
	require.Equal(t, 3, len(commits))

	// Inspect output
	repos, err := c.ListRepo()
	require.NoError(t, err)
	require.ElementsEqualUnderFn(t,
		[]string{"left", "right", "copy", "sort"},
		repos, RepoInfoToName)

	// make sure all numbers 0-99 are in /nums
	var buf bytes.Buffer
	require.NoError(t, c.GetFile("sort", "master", "/nums", 0, 0, &buf))
	s := bufio.NewScanner(&buf)
	numbers := make(map[string]struct{})
	for s.Scan() {
		numbers[s.Text()] = struct{}{}
	}
	require.Equal(t, 100, len(numbers)) // job processed all inputs

	// Confirm stats commits are present
	commits, err = c.ListCommit("sort", "stats", "", 0)
	require.NoError(t, err)
	require.Equal(t, 6, len(commits))
}

func int64p(i int64) *int64 {
	return &i
}
<|endoftext|>


<commit_before>package jsonp

import (
	"net/http"
	"testing"
)

type TestAAAResponseWriter struct {
	hdr    http.Header
	buf    []byte
	status int
}

func (ho TestAAAResponseWriter) Header() http.Header {
	return http.Header(ho.hdr)
}

func (ho TestAAAResponseWriter) Write(b []byte) (int, error) {
	return 0, nil
}

func (ho TestAAAResponseWriter) WriteHeader(code int) {
	ho.status = code
}

func Test301Redirect(t *testing.T) {

}

func Test_Exists(t *testing.T) {
	/*
		if Exists("./no-such-file.do-not-create") {
			t.Errorf("Expected false, file should not exist\n")
		}
		if !Exists("./test_data") {
			t.Errorf("Expected true, check of directory\n")
		}
		if !Exists("./test_data/file-exists") {
			t.Errorf("Expected true, check of directory\n")
		}
	*/
	SetJsonPrefix(")]}',\n")
	if JSON_Prefix != ")]}',\n" {
		t.Errorf("Expected prefix to be set, not set correctly\n")
	}

	var bob = "{\"bob\":\"bob\"}\n"

	if PrependPrefix(bob) != ")]}',\n{\"bob\":\"bob\"}\n" {
		t.Errorf("Expected prefix to be used, not used correctly\n")
	}

	SetJsonPrefix("")

	res := TestAAAResponseWriter{
		hdr:    make(http.Header),
		buf:    make([]byte, 0, 200),
		status: 0,
	}

	s := JsonP_Param(bob, res, "callback000")
	// fmt.Printf("s >%s<\n", s)
	if s != "callback000({\"bob\":\"bob\"}\n);" {
		t.Errorf("Expected JSONp callback, got >%s< \n", s)
	}

	// TODO: check that headers are set corectly

	// req, _ := http.NewRequest("GET", "http://localhost:8204/api/status?id=xyzzy", nil)		// Test requries a server at this locaiton
	req, _ := http.NewRequest("GET", "http://google.com/", nil)
	_ = req

	// TODO: func JsonP(s string, res http.ResponseWriter, req *http.Request) string {

}
<commit_msg>Added missing tests with JsonP validate return value and headers<commit_after>package jsonp

import (
	"net/http"
	"testing"
)

type TestAAAResponseWriter struct {
	hdr    http.Header
	buf    []byte
	status int
}

func (ho TestAAAResponseWriter) Header() http.Header {
	return http.Header(ho.hdr)
}

func (ho TestAAAResponseWriter) Write(b []byte) (int, error) {
	return 0, nil
}

func (ho TestAAAResponseWriter) WriteHeader(code int) {
	ho.status = code
}

func Test301Redirect(t *testing.T) {

}

func Test_Exists(t *testing.T) {
	/*
		if Exists("./no-such-file.do-not-create") {
			t.Errorf("Expected false, file should not exist\n")
		}
		if !Exists("./test_data") {
			t.Errorf("Expected true, check of directory\n")
		}
		if !Exists("./test_data/file-exists") {
			t.Errorf("Expected true, check of directory\n")
		}
	*/
	SetJsonPrefix(")]}',\n")
	if JSON_Prefix != ")]}',\n" {
		t.Errorf("Expected prefix to be set, not set correctly\n")
	}

	var bob = "{\"bob\":\"bob\"}\n"

	if PrependPrefix(bob) != ")]}',\n{\"bob\":\"bob\"}\n" {
		t.Errorf("Expected prefix to be used, not used correctly\n")
	}

	SetJsonPrefix("")

	res := TestAAAResponseWriter{
		hdr:    make(http.Header),
		buf:    make([]byte, 0, 200),
		status: 0,
	}

	s := JsonP_Param(bob, res, "callback000")
	// fmt.Printf("s >%s<\n", s)
	if s != "callback000({\"bob\":\"bob\"}\n);" {
		t.Errorf("Expected JSONp callback, got >%s< \n", s)
	}

	res1 := TestAAAResponseWriter{
		hdr:    make(http.Header),
		buf:    make([]byte, 0, 200),
		status: 0,
	}

	req := http.Request{
		Method:     "GET",
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		RemoteAddr: "[::1]:64770",
		RequestURI: "/api/status?callback=callback12323123123",
		Host:       "localhost:8204",
		Header:     make(http.Header),
	}

	u := JsonP(`{"status":"success"}`, res1, &req)
	// Check return value
	// fmt.Printf("s= ->%s<-\n", u)
	if u != `callback12323123123({"status":"success"});` {
		t.Errorf("Expected JSONp callback, got >%s< \n", u)
	}
	// check that headers are set corectly
	// fmt.Printf("hdr= ->%s<-\n", res1.hdr)
	if _, ok := res1.hdr["Content-Type"]; !ok {
		t.Errorf("Expected to have a Content-Type header set , but did not find it.")
	} else {
		if v := res1.hdr["Content-Type"]; len(v) > 0 && v[0] != "application/javascript" {
			t.Errorf("Expected to have a Content-Type header set to appliation/javascript but got %s instead\n", v)
		}
	}

}
<|endoftext|>


<commit_before>package lifted

// Quantifier removal

func (d *Domain) ExpandQuants(objs []TypedName) {
	f := newExpandFrame(objs)
	acts := make([]Action, 0, len(d.Actions))
	for i, _ := range d.Actions {
		a := &d.Actions[i]
		a.Precondition = a.Precondition.ExpandQuants(f)
		a.Effect = a.Effect.ExpandQuants(f)
		acts = append(acts, expandParams(f, a, a.Parameters)...)
	}
	d.Actions = acts
}

func expandParams(f *expandFrame, a *Action, ps []TypedName) (acts []Action) {
	if len(ps) == 0 {
		act := Action{
			Name: a.Name,
			Parameters: make([]TypedName, len(a.Parameters)),
			Precondition: a.Precondition.ExpandQuants(f),
			Effect: a.Effect.ExpandQuants(f),
		}
		copy(act.Parameters, a.Parameters)
		return []Action{ act }
	}

	pnum := len(a.Parameters) - len(ps)
	saved := a.Parameters[pnum]

	objs := objsOfType(f, saved.Type)
	for _, obj := range objs {
		a.Parameters[pnum].Name = obj
		g := f.push(saved.Name, obj)
		acts = append(acts, expandParams(g, a, ps[1:])...)
	}
 	a.Parameters[pnum] = saved

	return
}

func (l *Literal) ExpandQuants(f *expandFrame) *Literal {
	res := &Literal{
		Positive:   l.Positive,
		Name:       l.Name,
		Parameters: make([]Term, len(l.Parameters)),
	}
	for i, t := range l.Parameters {
		if t.Kind == TermConstant {
			res.Parameters[i] = t
			continue
		}
		o, ok := f.lookup(t.Name)
		if !ok {
			res.Parameters[i] = t
			continue
		}
		res.Parameters[i] = Term{
			Kind: TermConstant,
			Name: o,
			Loc: t.Loc,
		}
	}
	return res
}

func (e ExprTrue) ExpandQuants(*expandFrame) Expr {
	return e
}

func (e ExprFalse) ExpandQuants(*expandFrame) Expr {
	return e
}

func (e *ExprAnd) ExpandQuants(f *expandFrame) Expr {
	return ExprConj(e.Left.ExpandQuants(f), e.Right.ExpandQuants(f))
}

func (e *ExprOr) ExpandQuants(f *expandFrame) Expr {
	return ExprDisj(e.Left.ExpandQuants(f), e.Right.ExpandQuants(f))
}

func (e *ExprNot) ExpandQuants(f *expandFrame) Expr {
	return ExprNeg(e.Expr.ExpandQuants(f))
}

func (e *ExprQuant) ExpandQuants(f *expandFrame, base Expr,
comb func(Expr, Expr) Expr) Expr {
	objs := objsOfType(f, e.Variable.Type)
	seq := base
	for _, obj := range objs {
		g := f.push(e.Variable.Name, obj)
		seq = comb(e.Expr.ExpandQuants(g), seq)
	}
	return seq
}

func (e *ExprForall) ExpandQuants(f *expandFrame) Expr {
	return (*ExprQuant)(e).ExpandQuants(f, ExprTrue(1), ExprConj)
}

func (e *ExprExists) ExpandQuants(f *expandFrame) Expr {
	return (*ExprQuant)(e).ExpandQuants(f, ExprFalse(0), ExprDisj)
}

func (e *ExprLiteral) ExpandQuants(f *expandFrame) Expr {
	return (*ExprLiteral)((*Literal)(e).ExpandQuants(f))
}

func (EffectNone) ExpandQuants(*expandFrame) Effect {
	return EffectNone(0)
}

func (e *EffectAnd) ExpandQuants(f *expandFrame) Effect {
	return EffectConj(e.Left.ExpandQuants(f), e.Right.ExpandQuants(f))
}

func (e *EffectForall) ExpandQuants(f *expandFrame) Effect {
	objs := objsOfType(f, e.Variable.Type)
	seq := Effect(EffectNone(0))
	for _, obj := range objs {
		g := f.push(e.Variable.Name, obj)
		seq = EffectConj(e.Effect.ExpandQuants(g), seq)
	}
	return seq
}

func (e *EffectWhen) ExpandQuants(f *expandFrame) Effect {
	return &EffectWhen {
		Condition: e.Condition.ExpandQuants(f),
		EffectUnary: EffectUnary{ e.Effect.ExpandQuants(f) },
	}
}

func (e *EffectLiteral) ExpandQuants(f *expandFrame) Effect {
	return (*EffectLiteral)((*Literal)(e).ExpandQuants(f))
}

func (e *EffectAssign) ExpandQuants(f *expandFrame) Effect {
	return &EffectAssign {
		Op: e.Op,
		Lval: e.Lval,
		Rval: e.Rval,
	}
}

type expandFrame struct {
	objsByType map[string][]string
	variable   string
	object     string
	up         *expandFrame
}

func newExpandFrame(objs []TypedName) *expandFrame {
	objsByType := make(map[string][]string)

	for _, obj := range objs {
		for _, t := range obj.Type {
			lst, _ := objsByType[t]
			objsByType[t] = append(lst, obj.Name)
		}
	}

	return &expandFrame{ objsByType: objsByType }
}

func (f *expandFrame) push(vr string, obj string) *expandFrame {
	return &expandFrame{
		objsByType: f.objsByType,
		variable:   vr,
		object:     obj,
		up:         f,
	}
}

func (f *expandFrame) lookup(vr string) (string, bool) {
	if f == nil {
		return "", false
	}
	if f.variable == vr {
		return f.object, true
	}
	return f.up.lookup(vr)
}

func objsOfType(f *expandFrame, typ []string) (objs []string) {
	seen := make(map[string]bool)
	for _, t := range typ {
		for _, o := range f.objsByType[t] {
			if _, ok := seen[o]; ok {
				continue
			}
			seen[o] = true
			objs = append(objs, o)
		}
	}
	return
}
<commit_msg>gofmt.<commit_after>package lifted

// Quantifier removal

func (d *Domain) ExpandQuants(objs []TypedName) {
	f := newExpandFrame(objs)
	acts := make([]Action, 0, len(d.Actions))
	for i, _ := range d.Actions {
		a := &d.Actions[i]
		a.Precondition = a.Precondition.ExpandQuants(f)
		a.Effect = a.Effect.ExpandQuants(f)
		acts = append(acts, expandParams(f, a, a.Parameters)...)
	}
	d.Actions = acts
}

func expandParams(f *expandFrame, a *Action, ps []TypedName) (acts []Action) {
	if len(ps) == 0 {
		act := Action{
			Name:         a.Name,
			Parameters:   make([]TypedName, len(a.Parameters)),
			Precondition: a.Precondition.ExpandQuants(f),
			Effect:       a.Effect.ExpandQuants(f),
		}
		copy(act.Parameters, a.Parameters)
		return []Action{act}
	}

	pnum := len(a.Parameters) - len(ps)
	saved := a.Parameters[pnum]

	objs := objsOfType(f, saved.Type)
	for _, obj := range objs {
		a.Parameters[pnum].Name = obj
		g := f.push(saved.Name, obj)
		acts = append(acts, expandParams(g, a, ps[1:])...)
	}
	a.Parameters[pnum] = saved

	return
}

func (l *Literal) ExpandQuants(f *expandFrame) *Literal {
	res := &Literal{
		Positive:   l.Positive,
		Name:       l.Name,
		Parameters: make([]Term, len(l.Parameters)),
	}
	for i, t := range l.Parameters {
		if t.Kind == TermConstant {
			res.Parameters[i] = t
			continue
		}
		o, ok := f.lookup(t.Name)
		if !ok {
			res.Parameters[i] = t
			continue
		}
		res.Parameters[i] = Term{
			Kind: TermConstant,
			Name: o,
			Loc:  t.Loc,
		}
	}
	return res
}

func (e ExprTrue) ExpandQuants(*expandFrame) Expr {
	return e
}

func (e ExprFalse) ExpandQuants(*expandFrame) Expr {
	return e
}

func (e *ExprAnd) ExpandQuants(f *expandFrame) Expr {
	return ExprConj(e.Left.ExpandQuants(f), e.Right.ExpandQuants(f))
}

func (e *ExprOr) ExpandQuants(f *expandFrame) Expr {
	return ExprDisj(e.Left.ExpandQuants(f), e.Right.ExpandQuants(f))
}

func (e *ExprNot) ExpandQuants(f *expandFrame) Expr {
	return ExprNeg(e.Expr.ExpandQuants(f))
}

func (e *ExprQuant) ExpandQuants(f *expandFrame, base Expr,
comb func(Expr, Expr) Expr) Expr {
	objs := objsOfType(f, e.Variable.Type)
	seq := base
	for _, obj := range objs {
		g := f.push(e.Variable.Name, obj)
		seq = comb(e.Expr.ExpandQuants(g), seq)
	}
	return seq
}

func (e *ExprForall) ExpandQuants(f *expandFrame) Expr {
	return (*ExprQuant)(e).ExpandQuants(f, ExprTrue(1), ExprConj)
}

func (e *ExprExists) ExpandQuants(f *expandFrame) Expr {
	return (*ExprQuant)(e).ExpandQuants(f, ExprFalse(0), ExprDisj)
}

func (e *ExprLiteral) ExpandQuants(f *expandFrame) Expr {
	return (*ExprLiteral)((*Literal)(e).ExpandQuants(f))
}

func (EffectNone) ExpandQuants(*expandFrame) Effect {
	return EffectNone(0)
}

func (e *EffectAnd) ExpandQuants(f *expandFrame) Effect {
	return EffectConj(e.Left.ExpandQuants(f), e.Right.ExpandQuants(f))
}

func (e *EffectForall) ExpandQuants(f *expandFrame) Effect {
	objs := objsOfType(f, e.Variable.Type)
	seq := Effect(EffectNone(0))
	for _, obj := range objs {
		g := f.push(e.Variable.Name, obj)
		seq = EffectConj(e.Effect.ExpandQuants(g), seq)
	}
	return seq
}

func (e *EffectWhen) ExpandQuants(f *expandFrame) Effect {
	return &EffectWhen{
		Condition:   e.Condition.ExpandQuants(f),
		EffectUnary: EffectUnary{e.Effect.ExpandQuants(f)},
	}
}

func (e *EffectLiteral) ExpandQuants(f *expandFrame) Effect {
	return (*EffectLiteral)((*Literal)(e).ExpandQuants(f))
}

func (e *EffectAssign) ExpandQuants(f *expandFrame) Effect {
	return &EffectAssign{
		Op:   e.Op,
		Lval: e.Lval,
		Rval: e.Rval,
	}
}

type expandFrame struct {
	objsByType map[string][]string
	variable   string
	object     string
	up         *expandFrame
}

func newExpandFrame(objs []TypedName) *expandFrame {
	objsByType := make(map[string][]string)

	for _, obj := range objs {
		for _, t := range obj.Type {
			lst, _ := objsByType[t]
			objsByType[t] = append(lst, obj.Name)
		}
	}

	return &expandFrame{objsByType: objsByType}
}

func (f *expandFrame) push(vr string, obj string) *expandFrame {
	return &expandFrame{
		objsByType: f.objsByType,
		variable:   vr,
		object:     obj,
		up:         f,
	}
}

func (f *expandFrame) lookup(vr string) (string, bool) {
	if f == nil {
		return "", false
	}
	if f.variable == vr {
		return f.object, true
	}
	return f.up.lookup(vr)
}

func objsOfType(f *expandFrame, typ []string) (objs []string) {
	seen := make(map[string]bool)
	for _, t := range typ {
		for _, o := range f.objsByType[t] {
			if _, ok := seen[o]; ok {
				continue
			}
			seen[o] = true
			objs = append(objs, o)
		}
	}
	return
}
<|endoftext|>


<commit_before>package spartan

import (
	"errors"

	"github.com/nsf/termbox-go"
	"github.com/sparkymat/spartan/direction"
	"github.com/sparkymat/spartan/gravity"
	"github.com/sparkymat/spartan/size"
)

type LinearLayout struct {
	parent     ViewGroup
	views      []View
	x          uint32
	y          uint32
	width      size.Size
	height     size.Size
	direction  direction.Type
	isBordered bool
	title      string
}

func (layout *LinearLayout) SetTitle(title string) {
	layout.title = title
}

func (layout LinearLayout) GetTitle() string {
	return layout.title
}

func (layout *LinearLayout) EnableBorder() {
	layout.isBordered = true
}

func (layout *LinearLayout) DisableBorder() {
	layout.isBordered = false
}

func (layout *LinearLayout) SetWidth(width size.Size) {
	layout.width = width
}

func (layout *LinearLayout) SetHeight(height size.Size) {
	layout.height = height
}

func (layout *LinearLayout) SetLeftMargin(x uint32) {
	layout.x = x
}

func (layout *LinearLayout) SetTopMargin(y uint32) {
	layout.y = y
}

func (layout *LinearLayout) AddView(view View) {
	layout.views = append(layout.views, view)
}

func (layout LinearLayout) GetChildCount() uint32 {
	return uint32(len(layout.views))
}

func (layout LinearLayout) GetChildAt(index uint32) (View, error) {
	if index >= layout.GetChildCount() {
		return nil, errors.New("index out of bounds")
	}

	return layout.views[index], nil
}

func (layout LinearLayout) GetAbsoluteX() uint32 {
	if layout.GetParent() == nil {
		return layout.x
	} else {
		return layout.GetParent().GetAbsoluteX() + layout.x
	}
}

func (layout LinearLayout) GetAbsoluteY() uint32 {
	if layout.GetParent() == nil {
		return layout.y
	} else {
		return layout.GetParent().GetAbsoluteY() + layout.y
	}
}

func (layout LinearLayout) GetAbsoluteWidth() uint32 {
	if layout.width == size.MatchParent {
		if layout.parent == nil {
			width, _ := termbox.Size()
			return uint32(width)
		} else {
			return layout.parent.GetAbsoluteWidth()
		}
	} else {
		return uint32(layout.width)
	}
}

func (layout LinearLayout) GetAbsoluteHeight() uint32 {
	if layout.height == size.MatchParent {
		if layout.parent == nil {
			_, height := termbox.Size()
			return uint32(height)
		} else {
			return layout.parent.GetAbsoluteHeight()
		}
	} else {
		return uint32(layout.height)
	}
}

func (layout LinearLayout) draw() {
	containerLeft := layout.GetAbsoluteX()
	containerTop := layout.GetAbsoluteWidth()
	containerWidth := layout.GetAbsoluteWidth()
	containerHeight := layout.GetAbsoluteHeight()
	containerRight := containerLeft + containerWidth
	containerBottom := containerTop + containerHeight

	if layout.isBordered {
		containerLeft += 1
		containerRight -= 1
		containerWidth -= 2

		containerTop += 1
		containerBottom -= 1
		containerHeight -= 2
	}

	if layout.direction == direction.Vertical {
		heights := make([]uint32, len(layout.views), len(layout.views))

		stretchiesCount := uint32(0)
		totalFixedHeight := uint32(0)

		for i, view := range layout.views {
			if view.GetHeight() == size.MatchParent {
				stretchiesCount += 1
			} else {
				heights[i] = uint32(view.GetHeight())
				totalFixedHeight += uint32(view.GetHeight())
			}
		}

		if stretchiesCount > 0 {

			for i, view := range layout.views {
				if view.GetHeight() == size.MatchParent {
					heights[i] = (containerHeight - totalFixedHeight) / stretchiesCount
				}
			}
		}

		currentTop := containerTop

		for i, view := range layout.views {
			height := heights[i]
			width := view.GetAbsoluteWidth()

			currentLeft := uint32(0)

			if view.GetLayoutGravity() == gravity.Left {
				currentLeft = containerLeft + view.GetLeftMargin()
			} else if view.GetLayoutGravity() == gravity.Right {
				currentLeft = containerLeft + (containerWidth - width - view.GetRightMargin())
			} else if view.GetLayoutGravity() == gravity.Middle {
				currentLeft = containerLeft + (containerWidth-width)/2
			}

			currentRight := currentLeft + width - 1
			currentBottom := currentTop + height - 1

			// Clip to container dimensions
			if currentLeft < containerLeft {
				currentLeft = containerLeft
			}

			if currentRight > containerRight {
				currentRight = containerRight
			}

			if currentTop < containerTop {
				currentTop = containerTop
			}

			if currentBottom > containerBottom {
				currentBottom = containerBottom
			}

			view.draw(currentLeft, currentTop, currentRight, currentBottom)

			currentTop += height
		}
	} else if layout.direction == direction.Horizontal {
		widths := make([]uint32, len(layout.views), len(layout.views))

		stretchiesCount := uint32(0)
		totalFixedWidth := uint32(0)

		for i, view := range layout.views {
			if view.GetWidth() == size.MatchParent {
				stretchiesCount += 1
			} else {
				widths[i] = uint32(view.GetWidth())
				totalFixedWidth += uint32(view.GetWidth())
			}
		}

		if stretchiesCount > 0 {

			for i, view := range layout.views {
				if view.GetWidth() == size.MatchParent {
					widths[i] = (containerWidth - totalFixedWidth) / stretchiesCount
				}
			}
		}

		currentLeft := containerLeft

		for i, view := range layout.views {
			width := widths[i]
			height := view.GetAbsoluteHeight()

			currentTop := uint32(0)

			if view.GetLayoutGravity() == gravity.Top {
				currentTop = containerTop + view.GetTopMargin()
			} else if view.GetLayoutGravity() == gravity.Bottom {
				currentTop = containerTop + (containerHeight - height - view.GetBottomMargin())
			} else if view.GetLayoutGravity() == gravity.Center {
				currentTop = containerTop + (containerHeight-height)/2
			}

			currentRight := currentLeft + width - 1
			currentBottom := currentTop + height - 1

			// Clip to container dimensions
			if currentLeft < containerLeft {
				currentLeft = containerLeft
			}

			if currentRight > containerRight {
				currentRight = containerRight
			}

			if currentTop < containerTop {
				currentTop = containerTop
			}

			if currentBottom > containerBottom {
				currentBottom = containerBottom
			}

			view.draw(currentLeft, currentTop, currentLeft+width, currentTop+height)
			currentLeft += width
		}
	}
}

func (layout LinearLayout) GetParent() ViewGroup {
	return layout.parent
}

func (layout *LinearLayout) SetParent(parent ViewGroup) {
	layout.parent = parent
}

func (layout *LinearLayout) SetDirection(direction direction.Type) {
	layout.direction = direction
}

func (layout LinearLayout) GetDirection() direction.Type {
	return layout.direction
}
<commit_msg>That was really stupid :/<commit_after>package spartan

import (
	"errors"

	"github.com/nsf/termbox-go"
	"github.com/sparkymat/spartan/direction"
	"github.com/sparkymat/spartan/gravity"
	"github.com/sparkymat/spartan/size"
)

type LinearLayout struct {
	parent     ViewGroup
	views      []View
	x          uint32
	y          uint32
	width      size.Size
	height     size.Size
	direction  direction.Type
	isBordered bool
	title      string
}

func (layout *LinearLayout) SetTitle(title string) {
	layout.title = title
}

func (layout LinearLayout) GetTitle() string {
	return layout.title
}

func (layout *LinearLayout) EnableBorder() {
	layout.isBordered = true
}

func (layout *LinearLayout) DisableBorder() {
	layout.isBordered = false
}

func (layout *LinearLayout) SetWidth(width size.Size) {
	layout.width = width
}

func (layout *LinearLayout) SetHeight(height size.Size) {
	layout.height = height
}

func (layout *LinearLayout) SetLeftMargin(x uint32) {
	layout.x = x
}

func (layout *LinearLayout) SetTopMargin(y uint32) {
	layout.y = y
}

func (layout *LinearLayout) AddView(view View) {
	layout.views = append(layout.views, view)
}

func (layout LinearLayout) GetChildCount() uint32 {
	return uint32(len(layout.views))
}

func (layout LinearLayout) GetChildAt(index uint32) (View, error) {
	if index >= layout.GetChildCount() {
		return nil, errors.New("index out of bounds")
	}

	return layout.views[index], nil
}

func (layout LinearLayout) GetAbsoluteX() uint32 {
	if layout.GetParent() == nil {
		return layout.x
	} else {
		return layout.GetParent().GetAbsoluteX() + layout.x
	}
}

func (layout LinearLayout) GetAbsoluteY() uint32 {
	if layout.GetParent() == nil {
		return layout.y
	} else {
		return layout.GetParent().GetAbsoluteY() + layout.y
	}
}

func (layout LinearLayout) GetAbsoluteWidth() uint32 {
	if layout.width == size.MatchParent {
		if layout.parent == nil {
			width, _ := termbox.Size()
			return uint32(width)
		} else {
			return layout.parent.GetAbsoluteWidth()
		}
	} else {
		return uint32(layout.width)
	}
}

func (layout LinearLayout) GetAbsoluteHeight() uint32 {
	if layout.height == size.MatchParent {
		if layout.parent == nil {
			_, height := termbox.Size()
			return uint32(height)
		} else {
			return layout.parent.GetAbsoluteHeight()
		}
	} else {
		return uint32(layout.height)
	}
}

func (layout LinearLayout) draw() {
	containerLeft := layout.GetAbsoluteX()
	containerTop := layout.GetAbsoluteY()
	containerWidth := layout.GetAbsoluteWidth()
	containerHeight := layout.GetAbsoluteHeight()
	containerRight := containerLeft + containerWidth
	containerBottom := containerTop + containerHeight

	if layout.isBordered {
		containerLeft += 1
		containerRight -= 1
		containerWidth -= 2

		containerTop += 1
		containerBottom -= 1
		containerHeight -= 2
	}

	if layout.direction == direction.Vertical {
		heights := make([]uint32, len(layout.views), len(layout.views))

		stretchiesCount := uint32(0)
		totalFixedHeight := uint32(0)

		for i, view := range layout.views {
			if view.GetHeight() == size.MatchParent {
				stretchiesCount += 1
			} else {
				heights[i] = uint32(view.GetHeight())
				totalFixedHeight += uint32(view.GetHeight())
			}
		}

		if stretchiesCount > 0 {

			for i, view := range layout.views {
				if view.GetHeight() == size.MatchParent {
					heights[i] = (containerHeight - totalFixedHeight) / stretchiesCount
				}
			}
		}

		currentTop := containerTop

		for i, view := range layout.views {
			height := heights[i]
			width := view.GetAbsoluteWidth()

			currentLeft := uint32(0)

			if view.GetLayoutGravity() == gravity.Left {
				currentLeft = containerLeft + view.GetLeftMargin()
			} else if view.GetLayoutGravity() == gravity.Right {
				currentLeft = containerLeft + (containerWidth - width - view.GetRightMargin())
			} else if view.GetLayoutGravity() == gravity.Middle {
				currentLeft = containerLeft + (containerWidth-width)/2
			}

			currentRight := currentLeft + width - 1
			currentBottom := currentTop + height - 1

			// Clip to container dimensions
			if currentLeft < containerLeft {
				currentLeft = containerLeft
			}

			if currentRight > containerRight {
				currentRight = containerRight
			}

			if currentTop < containerTop {
				currentTop = containerTop
			}

			if currentBottom > containerBottom {
				currentBottom = containerBottom
			}

			view.draw(currentLeft, currentTop, currentRight, currentBottom)

			currentTop += height
		}
	} else if layout.direction == direction.Horizontal {
		widths := make([]uint32, len(layout.views), len(layout.views))

		stretchiesCount := uint32(0)
		totalFixedWidth := uint32(0)

		for i, view := range layout.views {
			if view.GetWidth() == size.MatchParent {
				stretchiesCount += 1
			} else {
				widths[i] = uint32(view.GetWidth())
				totalFixedWidth += uint32(view.GetWidth())
			}
		}

		if stretchiesCount > 0 {

			for i, view := range layout.views {
				if view.GetWidth() == size.MatchParent {
					widths[i] = (containerWidth - totalFixedWidth) / stretchiesCount
				}
			}
		}

		currentLeft := containerLeft

		for i, view := range layout.views {
			width := widths[i]
			height := view.GetAbsoluteHeight()

			currentTop := uint32(0)

			if view.GetLayoutGravity() == gravity.Top {
				currentTop = containerTop + view.GetTopMargin()
			} else if view.GetLayoutGravity() == gravity.Bottom {
				currentTop = containerTop + (containerHeight - height - view.GetBottomMargin())
			} else if view.GetLayoutGravity() == gravity.Center {
				currentTop = containerTop + (containerHeight-height)/2
			}

			currentRight := currentLeft + width - 1
			currentBottom := currentTop + height - 1

			// Clip to container dimensions
			if currentLeft < containerLeft {
				currentLeft = containerLeft
			}

			if currentRight > containerRight {
				currentRight = containerRight
			}

			if currentTop < containerTop {
				currentTop = containerTop
			}

			if currentBottom > containerBottom {
				currentBottom = containerBottom
			}

			view.draw(currentLeft, currentTop, currentLeft+width, currentTop+height)
			currentLeft += width
		}
	}
}

func (layout LinearLayout) GetParent() ViewGroup {
	return layout.parent
}

func (layout *LinearLayout) SetParent(parent ViewGroup) {
	layout.parent = parent
}

func (layout *LinearLayout) SetDirection(direction direction.Type) {
	layout.direction = direction
}

func (layout LinearLayout) GetDirection() direction.Type {
	return layout.direction
}
<|endoftext|>


<commit_before>package devastator

import (
	"crypto/tls"
	"crypto/x509"
	"strconv"
	"testing"
)

func TestLen(t *testing.T) {
	a, _ := strconv.Atoi("12344324")
	t.Log(a)
}

// todo: if we are going to expose raw Listener, this should be in integration tests, otherwise Listener should be private
func TestListener(t *testing.T) {
	msg1 := "Lorem ipsum dolor sit amet, consectetur adipiscing elit."
	msg2 := "In sit amet lectus felis, at pellentesque turpis."
	msg3 := "Nunc urna enim, cursus varius aliquet ac, imperdiet eget tellus."
	msg4 := randString(45000)
	msg5 := randString(500000)

	host := "localhost:" + Conf.App.Port
	cert, privKey, _ := genCert("localhost", 0, nil, nil, 512, "localhost", "devastator")
	l, err := Listen(cert, privKey, host, false)
	if err != nil {
		t.Fatal(err)
	}
	defer l.Close()

	go l.Accept(func(conn *Conn, session *Session, msg []byte) {
		certs := conn.ConnectionState().PeerCertificates
		if len(certs) > 0 {
			t.Logf("Client connected with client certificate subject: %v\n", certs[0].Subject)
		}

		m := string(msg)
		if m != msg1 && m != msg2 && m != msg3 && m != msg4 && m != msg5 {
			t.Fatal("Sent and incoming message did not match for message:", m)
		}
	}, func(conn *Conn, session *Session) {
	})

	roots := x509.NewCertPool()
	ok := roots.AppendCertsFromPEM(cert)
	if !ok {
		panic("failed to parse root certificate")
	}

	tlsConf := &tls.Config{RootCAs: roots}
	conn, err := tls.Dial("tcp", host, tlsConf)
	if err != nil {
		t.Fatal(err)
	}
	defer l.connWG.Wait()
	defer conn.Close()

	newconn := NewConn(conn, 0, 0, 0, false)

	send(t, newconn, "ping")
	send(t, newconn, msg1)
	send(t, newconn, msg1)
	send(t, newconn, msg2)
	send(t, newconn, msg3)
	send(t, newconn, msg4)
	send(t, newconn, msg1)
	send(t, newconn, msg5)
	send(t, newconn, msg1)
	send(t, newconn, "close")

	// t.Logf("\nconn:\n%+v\n\n", conn)
	// t.Logf("\nconn.ConnectionState():\n%+v\n\n", conn.ConnectionState())
	// t.Logf("\ntls.Config:\n%+v\n\n", tlsConf)
}

// func TestClientDisconnect(t *testing.T) {
// 	// todo: we need to verify that events occur in the order that we want them (either via event hooks or log analysis)
// 	// this seems like a listener test than a integration test from a client perspective
// 	s := getServer(t)
// 	c := getClientConnWithClientCert(t)
// 	if err := c.Close(); err != nil {
// 		t.Fatal("Failed to close the client connection:", err)
// 	}
// 	if err := s.Stop(); err != nil {
// 		t.Fatal("Failed to stop the server:", err)
// 	}
// 	wg.Wait()
// }

func send(t *testing.T, conn *Conn, msg string) {
	n, err := conn.Write([]byte(msg))
	if err != nil {
		t.Fatal(err)
	}
	if n < 100 {
		t.Logf("Sent message to listener from client: %v (%v bytes)", msg, n)
	} else {
		t.Logf("Sent message to listener from client: ... (%v bytes)", n)
	}
}
<commit_msg>tidy up connect testing<commit_after>package devastator

import (
	"crypto/tls"
	"crypto/x509"
	"strconv"
	"testing"
)

func TestLen(t *testing.T) {
	a, _ := strconv.Atoi("12344324")
	t.Log(a)
}

// todo: if we are going to expose raw Listener, this should be in integration tests, otherwise Listener should be private
func TestListener(t *testing.T) {
	msg1 := "Lorem ipsum dolor sit amet, consectetur adipiscing elit."
	msg2 := "In sit amet lectus felis, at pellentesque turpis."
	msg3 := "Nunc urna enim, cursus varius aliquet ac, imperdiet eget tellus."
	msg4 := randString(45000)
	msg5 := randString(500000)

	host := "localhost:" + Conf.App.Port
	cert, privKey, _ := genCert("localhost", 0, nil, nil, 512, "localhost", "devastator")
	l, err := Listen(cert, privKey, host, false)
	if err != nil {
		t.Fatal(err)
	}
	defer l.Close()

	go l.Accept(func(conn *Conn, session *Session, msg []byte) {
		certs := conn.ConnectionState().PeerCertificates
		if len(certs) > 0 {
			t.Logf("Client connected with client certificate subject: %v\n", certs[0].Subject)
		}

		m := string(msg)
		if m != msg1 && m != msg2 && m != msg3 && m != msg4 && m != msg5 {
			t.Fatal("Sent and incoming message did not match for message:", m)
		}
	}, func(conn *Conn, session *Session) {
	})

	roots := x509.NewCertPool()
	ok := roots.AppendCertsFromPEM(cert)
	if !ok {
		panic("failed to parse root certificate")
	}

	tlsConf := &tls.Config{RootCAs: roots}
	conn, err := tls.Dial("tcp", host, tlsConf)
	if err != nil {
		t.Fatal(err)
	}

	newconn := NewConn(conn, 0, 0, 0, false)
	// defer newconn.Close()

	send(t, newconn, "ping")
	send(t, newconn, msg1)
	send(t, newconn, msg1)
	send(t, newconn, msg2)
	send(t, newconn, msg3)
	send(t, newconn, msg4)
	send(t, newconn, msg1)
	send(t, newconn, msg5)
	send(t, newconn, msg1)
	send(t, newconn, "close")

	// t.Logf("\nconn:\n%+v\n\n", conn)
	// t.Logf("\nconn.ConnectionState():\n%+v\n\n", conn.ConnectionState())
	// t.Logf("\ntls.Config:\n%+v\n\n", tlsConf)
}

// func TestClientDisconnect(t *testing.T) {
// 	// todo: we need to verify that events occur in the order that we want them (either via event hooks or log analysis)
// 	// this seems like a listener test than a integration test from a client perspective
// 	s := getServer(t)
// 	c := getClientConnWithClientCert(t)
// 	if err := c.Close(); err != nil {
// 		t.Fatal("Failed to close the client connection:", err)
// 	}
// 	if err := s.Stop(); err != nil {
// 		t.Fatal("Failed to stop the server:", err)
// 	}
// 	wg.Wait()
// }

func send(t *testing.T, conn *Conn, msg string) {
	n, err := conn.Write([]byte(msg))
	if err != nil {
		t.Fatal(err)
	}
	if n < 100 {
		t.Logf("Sent message to listener from client: %v (%v bytes)", msg, n)
	} else {
		t.Logf("Sent message to listener from client: ... (%v bytes)", n)
	}
}
<|endoftext|>


<commit_before>package lock

import (
	"container/list"
	"fmt"
	"hash/crc32"
	"path"
	"sort"
	"strings"
	"sync"
	"time"
)

const defaultPathSlotSize = 4096

type pendingNode struct {
	path string
	ch   chan struct{}
}

func newPendingNode(path string) *pendingNode {
	n := new(pendingNode)
	n.path = path
	n.ch = make(chan struct{}, 1)
	return n
}

type pathLockerSet struct {
	m sync.Mutex

	set *refValueSet

	pendingLock sync.Mutex
	pendingPath *list.List
}

func newPathLockerSet() *pathLockerSet {
	s := new(pathLockerSet)
	s.pendingPath = list.New()

	f := func(v *refValue) {
		v.v = false
	}

	s.set = newRefValueSet(f)
	return s
}

// a/b/c/ return ["a/", "a/b/", "a/b/c/"]
func (s *pathLockerSet) makeAncestorPaths(path string) []string {
	items := make([]string, 0, 4)

	pos := 0
	for {
		index := strings.IndexByte(path[pos:], '/')
		if index == -1 {
			break
		}

		item := path[0 : pos+index+1]
		items = append(items, item)

		pos += index + 1
		if pos >= len(path) {
			break
		}
	}

	return items
}

func (s *pathLockerSet) tryLock(path string) *pendingNode {
	items := s.makeAncestorPaths(path)

	vs := make([]*refValue, 0, len(items))

	s.m.Lock()
	defer s.m.Unlock()

	var v *refValue
	locked := true

	for _, item := range items {
		v = s.set.Get(item)
		vs = append(vs, v)

		if v.v == true {
			// other lock the ancestor path
			locked = false
			break
		}
	}

	if v.ref != 1 {
		// other lock the children path
		locked = false
	} else {
		v.v = true
		locked = true
	}

	if !locked {
		for i, v := range vs {
			s.set.Put(items[i], v)
		}
		return newPendingNode(path)
	} else {
		return nil
	}
}

func (s *pathLockerSet) LockTimeout(path string, t *time.Timer) bool {
	for {
		n := s.tryLock(path)
		if n == nil {
			return true
		}
		s.addPendingNode(n)

		select {
		case <-n.ch:
		case <-t.C:
			return false
		}
	}
}

func (s *pathLockerSet) addPendingNode(n *pendingNode) {
	s.pendingLock.Lock()
	s.pendingPath.PushBack(n)
	s.pendingLock.Unlock()
}

func (s *pathLockerSet) noticePendingNode(path string) {
	s.pendingLock.Lock()

	var next *list.Element
	for e := s.pendingPath.Front(); e != nil; e = next {
		m := e.Value.(*pendingNode)
		next = e.Next()

		if strings.Contains(path, m.path) || strings.Contains(m.path, path) {
			s.pendingPath.Remove(e)
			m.ch <- struct{}{}
		}
	}

	s.pendingLock.Unlock()
}

func (s *pathLockerSet) Unlock(path string) {
	s.unlock(path)

	s.noticePendingNode(path)
}

func (s *pathLockerSet) unlock(path string) {
	s.m.Lock()
	defer s.m.Unlock()

	items := s.makeAncestorPaths(path)

	for i := len(items) - 1; i >= 0; i-- {
		key := items[i]

		v := s.set.RawGet(key)
		if v == nil {
			panic(fmt.Sprintf("%s is not locked, panic", path))
		}

		v.v = false
		s.set.Put(key, v)
	}
}

type PathLockerGroup struct {
	set []*pathLockerSet
}

func (g *PathLockerGroup) Lock(paths ...string) {
	// use a very long timeout
	b := g.LockTimeout(InfiniteTimeout, paths...)
	if !b {
		panic("Wait lock too long, panic")
	}
}

func (g *PathLockerGroup) LockTimeout(timeout time.Duration, paths ...string) bool {
	if len(paths) == 0 {
		panic("empty paths, panic")
	}

	paths = g.canoicalizePaths(paths...)

	timer := time.NewTimer(timeout)
	defer timer.Stop()

	grapNum := 0
	for _, path := range paths {
		b := g.getSet(path).LockTimeout(path, timer)
		if !b {
			g.Unlock(paths[0:grapNum]...)
			return false
		} else {
			grapNum++
		}
	}

	return true
}

func (g *PathLockerGroup) Unlock(paths ...string) {
	if len(paths) == 0 {
		return
	}

	paths = g.canoicalizePaths(paths...)
	sort.Sort(sort.Reverse(sort.StringSlice(paths)))

	for _, path := range paths {
		g.getSet(path).Unlock(path)
	}
}

func (g *PathLockerGroup) getSet(path string) *pathLockerSet {
	base := strings.SplitN(path, "/", 2)
	index := crc32.ChecksumIEEE([]byte(base[0])) % uint32(defaultPathSlotSize)
	return g.set[index]
}

func (g *PathLockerGroup) canonicalizePath(p string) string {
	p = path.Clean(p)

	// remove first, so /a/b/c will be a/b/c
	p = strings.TrimPrefix(p, "/")

	// add / suffix, path Clean will remove the / suffix
	p = p + "/"

	return p
}

func (g *PathLockerGroup) canoicalizePaths(paths ...string) []string {
	p := make([]string, 0, len(paths))

	for i, path := range paths {
		paths[i] = g.canonicalizePath(path)
		if paths[i] == "/" {
			panic("invalid path, can not empty")
		}
	}

	sort.Strings(paths)

	p = append(p, paths[0])

	for i := 1; i < len(paths); i++ {
		for j := 0; j < len(p); j++ {
			if strings.Contains(paths[i], p[j]) {
				// if we want to lock a/b and a/b/c at same time, we only
				// need to lock the parent path a/b
				break
			}
		}

		p = append(p, paths[i])
	}

	return p
}

func NewPathLockerGroup() *PathLockerGroup {
	g := new(PathLockerGroup)
	g.set = make([]*pathLockerSet, defaultPathSlotSize)
	for i := 0; i < defaultPathSlotSize; i++ {
		g.set[i] = newPathLockerSet()
	}

	return g
}
<commit_msg>check unlock<commit_after>package lock

import (
	"container/list"
	"fmt"
	"hash/crc32"
	"path"
	"sort"
	"strings"
	"sync"
	"time"
)

const defaultPathSlotSize = 4096

type pendingNode struct {
	path string
	ch   chan struct{}
}

func newPendingNode(path string) *pendingNode {
	n := new(pendingNode)
	n.path = path
	n.ch = make(chan struct{}, 1)
	return n
}

type pathLockerSet struct {
	m sync.Mutex

	set *refValueSet

	pendingLock sync.Mutex
	pendingPath *list.List
}

func newPathLockerSet() *pathLockerSet {
	s := new(pathLockerSet)
	s.pendingPath = list.New()

	f := func(v *refValue) {
		v.v = false
	}

	s.set = newRefValueSet(f)
	return s
}

// a/b/c/ return ["a/", "a/b/", "a/b/c/"]
func (s *pathLockerSet) makeAncestorPaths(path string) []string {
	items := make([]string, 0, 4)

	pos := 0
	for {
		index := strings.IndexByte(path[pos:], '/')
		if index == -1 {
			break
		}

		item := path[0 : pos+index+1]
		items = append(items, item)

		pos += index + 1
		if pos >= len(path) {
			break
		}
	}

	return items
}

func (s *pathLockerSet) tryLock(path string) *pendingNode {
	items := s.makeAncestorPaths(path)

	vs := make([]*refValue, 0, len(items))

	s.m.Lock()
	defer s.m.Unlock()

	var v *refValue
	locked := true

	for _, item := range items {
		v = s.set.Get(item)
		vs = append(vs, v)

		if v.v == true {
			// other lock the ancestor path
			locked = false
			break
		}
	}

	if v.ref != 1 {
		// other lock the children path
		locked = false
	} else {
		v.v = true
		locked = true
	}

	if !locked {
		for i, v := range vs {
			s.set.Put(items[i], v)
		}
		return newPendingNode(path)
	} else {
		return nil
	}
}

func (s *pathLockerSet) LockTimeout(path string, t *time.Timer) bool {
	for {
		n := s.tryLock(path)
		if n == nil {
			return true
		}
		s.addPendingNode(n)

		select {
		case <-n.ch:
		case <-t.C:
			return false
		}
	}
}

func (s *pathLockerSet) addPendingNode(n *pendingNode) {
	s.pendingLock.Lock()
	s.pendingPath.PushBack(n)
	s.pendingLock.Unlock()
}

func (s *pathLockerSet) noticePendingNode(path string) {
	s.pendingLock.Lock()

	var next *list.Element
	for e := s.pendingPath.Front(); e != nil; e = next {
		m := e.Value.(*pendingNode)
		next = e.Next()

		if strings.Contains(path, m.path) || strings.Contains(m.path, path) {
			s.pendingPath.Remove(e)
			m.ch <- struct{}{}
		}
	}

	s.pendingLock.Unlock()
}

func (s *pathLockerSet) Unlock(path string) {
	s.unlock(path)

	s.noticePendingNode(path)
}

func (s *pathLockerSet) unlock(path string) {
	items := s.makeAncestorPaths(path)
	vs := make([]*refValue, 0, len(items))

	s.m.Lock()
	defer s.m.Unlock()

	for i := 0; i < len(items); i++ {
		key := items[i]

		v := s.set.RawGet(key)
		if v == nil {
			panic(fmt.Sprintf("%s is not locked, panic", path))
		} else if i == len(items)-1 && v.v != true {
			panic(fmt.Sprintf("%s is not locked, panic", path))
		} else if i != len(items)-1 && v.v == true {
			panic(fmt.Sprintf("%s ancestor %s is locked, panic", path, items[i]))
		}

		vs = append(vs, v)
	}

	for i, v := range vs {
		v.v = false
		s.set.Put(items[i], v)
	}
}

type PathLockerGroup struct {
	set []*pathLockerSet
}

func (g *PathLockerGroup) Lock(paths ...string) {
	// use a very long timeout
	b := g.LockTimeout(InfiniteTimeout, paths...)
	if !b {
		panic("Wait lock too long, panic")
	}
}

func (g *PathLockerGroup) LockTimeout(timeout time.Duration, paths ...string) bool {
	if len(paths) == 0 {
		panic("empty paths, panic")
	}

	paths = g.canoicalizePaths(paths...)

	timer := time.NewTimer(timeout)
	defer timer.Stop()

	grapNum := 0
	for _, path := range paths {
		b := g.getSet(path).LockTimeout(path, timer)
		if !b {
			g.Unlock(paths[0:grapNum]...)
			return false
		} else {
			grapNum++
		}
	}

	return true
}

func (g *PathLockerGroup) Unlock(paths ...string) {
	if len(paths) == 0 {
		return
	}

	paths = g.canoicalizePaths(paths...)
	sort.Sort(sort.Reverse(sort.StringSlice(paths)))

	for _, path := range paths {
		g.getSet(path).Unlock(path)
	}
}

func (g *PathLockerGroup) getSet(path string) *pathLockerSet {
	base := strings.SplitN(path, "/", 2)
	index := crc32.ChecksumIEEE([]byte(base[0])) % uint32(defaultPathSlotSize)
	return g.set[index]
}

func (g *PathLockerGroup) canonicalizePath(p string) string {
	p = path.Clean(p)

	// remove first, so /a/b/c will be a/b/c
	p = strings.TrimPrefix(p, "/")

	// add / suffix, path Clean will remove the / suffix
	p = p + "/"

	return p
}

func (g *PathLockerGroup) canoicalizePaths(paths ...string) []string {
	p := make([]string, 0, len(paths))

	for i, path := range paths {
		paths[i] = g.canonicalizePath(path)
		if paths[i] == "/" {
			panic("invalid path, can not empty")
		}
	}

	sort.Strings(paths)

	p = append(p, paths[0])

	for i := 1; i < len(paths); i++ {
		for j := 0; j < len(p); j++ {
			if strings.Contains(paths[i], p[j]) {
				// if we want to lock a/b and a/b/c at same time, we only
				// need to lock the parent path a/b
				break
			}
		}

		p = append(p, paths[i])
	}

	return p
}

func NewPathLockerGroup() *PathLockerGroup {
	g := new(PathLockerGroup)
	g.set = make([]*pathLockerSet, defaultPathSlotSize)
	for i := 0; i < defaultPathSlotSize; i++ {
		g.set[i] = newPathLockerSet()
	}

	return g
}
<|endoftext|>