<commit_before>package models

import (
	"context"
	"crypto/rand"
	"time"

	"github.com/gofrs/uuid"
	"github.com/pkg/errors"
	"golang.org/x/crypto/nacl/secretbox"

	"github.com/dollarshaveclub/furan/pkg/generated/furanrpc"
)

//go:generate stringer -type=BuildStatus

type BuildStatus int

const (
	// Invalid or unknown status
	BuildStatusUnknown BuildStatus = iota
	// Build has been requested but not started yet
	BuildStatusNotStarted
	// Build was requested but determined to be unnecessary
	BuildStatusSkipped
	// Build is currently running in a k8s job
	BuildStatusRunning
	// Build failed or internal error
	BuildStatusFailure
	// Build successfully completed & pushed
	BuildStatusSuccess
	// Build cancellation was requested but build has not yet aborted
	BuildStatusCancelRequested
	// Build was aborted due to cancellation request
	BuildStatusCancelled
)

func (bs BuildStatus) State() furanrpc.BuildState {
	switch bs {
	case BuildStatusUnknown:
		return furanrpc.BuildState_UNKNOWN
	case BuildStatusNotStarted:
		return furanrpc.BuildState_NOTSTARTED
	case BuildStatusSkipped:
		return furanrpc.BuildState_SKIPPED
	case BuildStatusRunning:
		return furanrpc.BuildState_RUNNING
	case BuildStatusFailure:
		return furanrpc.BuildState_FAILURE
	case BuildStatusSuccess:
		return furanrpc.BuildState_SUCCESS
	case BuildStatusCancelRequested:
		return furanrpc.BuildState_CANCEL_REQUESTED
	case BuildStatusCancelled:
		return furanrpc.BuildState_CANCELLED
	default:
		return furanrpc.BuildState_UNKNOWN
	}
}

func BuildStatusFromState(s furanrpc.BuildState) BuildStatus {
	switch s {
	case furanrpc.BuildState_UNKNOWN:
		return BuildStatusUnknown
	case furanrpc.BuildState_NOTSTARTED:
		return BuildStatusNotStarted
	case furanrpc.BuildState_SKIPPED:
		return BuildStatusSkipped
	case furanrpc.BuildState_RUNNING:
		return BuildStatusRunning
	case furanrpc.BuildState_FAILURE:
		return BuildStatusFailure
	case furanrpc.BuildState_SUCCESS:
		return BuildStatusSuccess
	case furanrpc.BuildState_CANCEL_REQUESTED:
		return BuildStatusCancelRequested
	case furanrpc.BuildState_CANCELLED:
		return BuildStatusCancelled
	default:
		return BuildStatusUnknown
	}
}

// TerminalState returns whether the status is in a final (terminal) state that will not change
func (bs BuildStatus) TerminalState() bool {
	switch bs {
	case BuildStatusSuccess:
		fallthrough
	case BuildStatusFailure:
		fallthrough
	case BuildStatusCancelled:
		return true
	default:
		return false
	}
}

type Build struct {
	ID                          uuid.UUID
	Created, Updated, Completed time.Time
	GitHubRepo, GitHubRef       string
	EncryptedGitHubCredential   []byte
	ImageRepos                  []string
	Tags                        []string
	CommitSHATag                bool
	BuildOptions                BuildOpts
	Request                     furanrpc.BuildRequest
	Status                      BuildStatus
	Events                      []string
}

// CanAddEvent indicates whether b is in a state where events can be added
func (b Build) CanAddEvent() bool {
	return b.EventListenable()
}

// EventListenable indicates where b is in a state where events can be listened for
func (b Build) EventListenable() bool {
	return !b.Status.TerminalState()
}

func (b Build) Running() bool {
	return b.Status == BuildStatusRunning
}

// TimeFromRPCTimestamp returns a UTC time.Time for an RPC timestamp
func TimeFromRPCTimestamp(ts furanrpc.Timestamp) time.Time {
	return time.Unix(ts.Seconds, int64(ts.Nanos)).UTC()
}

// RPCTimestampFromTime takes a time.Time and returns an RPC timestamp
func RPCTimestampFromTime(t time.Time) furanrpc.Timestamp {
	return furanrpc.Timestamp{
		Seconds: t.Unix(),
		Nanos:   int32(t.Nanosecond()),
	}
}

// EncryptAndSetGitHubCredential takes a GitHub credential, encrypts it and sets EncryptedGitHubCredential accordingly
func (b *Build) EncryptAndSetGitHubCredential(cred []byte, key [32]byte) error {
	var nonce [24]byte
	if n, err := rand.Read(nonce[:]); err != nil || n != len(nonce) {
		return errors.Wrapf(err, "error reading random bytes for nonce (read: %v)", n)
	}
	b.EncryptedGitHubCredential = secretbox.Seal(nonce[:], cred, &nonce, &key)
	return nil
}

// GetGitHubCredential returns the decrypted user token using key or error
func (b Build) GetGitHubCredential(key [32]byte) (string, error) {
	var nonce [24]byte
	copy(nonce[:], b.EncryptedGitHubCredential[:24])
	tkn, ok := secretbox.Open(nil, b.EncryptedGitHubCredential[24:], &nonce, &key)
	if !ok {
		return "", errors.New("decryption error (incorrect key?)")
	}
	if tkn == nil {
		return "", errors.New("decrypted token was nil")
	}
	return string(tkn), nil
}

// BuildOpts models all options required to perform a build
type BuildOpts struct {
	BuildID                uuid.UUID               `json:"-"`
	ContextPath, CommitSHA string                  `json:"-"` // set by Builder
	RelativeDockerfilePath string                  `json:"relative_dockerfile_path"`
	BuildArgs              map[string]string       `json:"build_args"`
	Cache                  furanrpc.BuildCacheOpts `json:"cache_opts"`
}

// Job describes methods on a single abstract build job
type Job interface {
	// Error returns a channel that will contain any errors associated with this Job
	Error() chan error
	// Running returns a channel that signals that the build the Job is executing has been updated to status Running
	// This indicates that the Furan sidecar has started and is executing successfully and will take responsibility for
	// tracking the build status from this point forward
	Running() chan struct{}
	// Logs returns all pod logs associated with the Job
	Logs() (map[string]map[string][]byte, error)
}

type JobRunner interface {
	Run(build Build) (Job, error)
}

// CacheFetcher describes an object that fetches and saves build cache
type CacheFetcher interface {
	// Fetch fetches the build cache for a build and returns a local filesystem
	// path where it was written. Caller is responsible for cleaning up the path when finished.
	Fetch(ctx context.Context, b Build) (string, error)
	// Save persists the build cache for a build located at path.
	// Caller is responsible for cleaning up the path afterward.
	Save(ctx context.Context, b Build, path string) error
}

// CodeFetcher represents an object capable of fetching code
type CodeFetcher interface {
	GetCommitSHA(ctx context.Context, repo, ref string) (string, error)
	Fetch(ctx context.Context, repo, ref, destinationPath string) error
}

// Builder describes an image build backend
type Builder interface {
	Build(ctx context.Context, opts BuildOpts) error
}

// BuilderManager describes an object that manages builds
type BuildManager interface {
	Start(ctx context.Context, opts BuildOpts) error
	Run(ctx context.Context, id uuid.UUID) error
}

type TagChecker interface {
	AllTagsExist(tags []string, repo string) (bool, []string, error)
}

// APIKey models a user-created API key
type APIKey struct {
	ID                            uuid.UUID
	Created                       time.Time
	GitHubUser, Name, Description string
	ReadOnly                      bool
}
<commit_msg>fix for skipped builds<commit_after>package models

import (
	"context"
	"crypto/rand"
	"time"

	"github.com/gofrs/uuid"
	"github.com/pkg/errors"
	"golang.org/x/crypto/nacl/secretbox"

	"github.com/dollarshaveclub/furan/pkg/generated/furanrpc"
)

//go:generate stringer -type=BuildStatus

type BuildStatus int

const (
	// Invalid or unknown status
	BuildStatusUnknown BuildStatus = iota
	// Build has been requested but not started yet
	BuildStatusNotStarted
	// Build was requested but determined to be unnecessary
	BuildStatusSkipped
	// Build is currently running in a k8s job
	BuildStatusRunning
	// Build failed or internal error
	BuildStatusFailure
	// Build successfully completed & pushed
	BuildStatusSuccess
	// Build cancellation was requested but build has not yet aborted
	BuildStatusCancelRequested
	// Build was aborted due to cancellation request
	BuildStatusCancelled
)

func (bs BuildStatus) State() furanrpc.BuildState {
	switch bs {
	case BuildStatusUnknown:
		return furanrpc.BuildState_UNKNOWN
	case BuildStatusNotStarted:
		return furanrpc.BuildState_NOTSTARTED
	case BuildStatusSkipped:
		return furanrpc.BuildState_SKIPPED
	case BuildStatusRunning:
		return furanrpc.BuildState_RUNNING
	case BuildStatusFailure:
		return furanrpc.BuildState_FAILURE
	case BuildStatusSuccess:
		return furanrpc.BuildState_SUCCESS
	case BuildStatusCancelRequested:
		return furanrpc.BuildState_CANCEL_REQUESTED
	case BuildStatusCancelled:
		return furanrpc.BuildState_CANCELLED
	default:
		return furanrpc.BuildState_UNKNOWN
	}
}

func BuildStatusFromState(s furanrpc.BuildState) BuildStatus {
	switch s {
	case furanrpc.BuildState_UNKNOWN:
		return BuildStatusUnknown
	case furanrpc.BuildState_NOTSTARTED:
		return BuildStatusNotStarted
	case furanrpc.BuildState_SKIPPED:
		return BuildStatusSkipped
	case furanrpc.BuildState_RUNNING:
		return BuildStatusRunning
	case furanrpc.BuildState_FAILURE:
		return BuildStatusFailure
	case furanrpc.BuildState_SUCCESS:
		return BuildStatusSuccess
	case furanrpc.BuildState_CANCEL_REQUESTED:
		return BuildStatusCancelRequested
	case furanrpc.BuildState_CANCELLED:
		return BuildStatusCancelled
	default:
		return BuildStatusUnknown
	}
}

// TerminalState returns whether the status is in a final (terminal) state that will not change
func (bs BuildStatus) TerminalState() bool {
	switch bs {
	case BuildStatusSuccess:
		fallthrough
	case BuildStatusFailure:
		fallthrough
	case BuildStatusSkipped:
		fallthrough
	case BuildStatusCancelled:
		return true
	default:
		return false
	}
}

type Build struct {
	ID                          uuid.UUID
	Created, Updated, Completed time.Time
	GitHubRepo, GitHubRef       string
	EncryptedGitHubCredential   []byte
	ImageRepos                  []string
	Tags                        []string
	CommitSHATag                bool
	BuildOptions                BuildOpts
	Request                     furanrpc.BuildRequest
	Status                      BuildStatus
	Events                      []string
}

// CanAddEvent indicates whether b is in a state where events can be added
func (b Build) CanAddEvent() bool {
	return b.EventListenable()
}

// EventListenable indicates where b is in a state where events can be listened for
func (b Build) EventListenable() bool {
	return !b.Status.TerminalState()
}

func (b Build) Running() bool {
	return b.Status == BuildStatusRunning
}

// TimeFromRPCTimestamp returns a UTC time.Time for an RPC timestamp
func TimeFromRPCTimestamp(ts furanrpc.Timestamp) time.Time {
	return time.Unix(ts.Seconds, int64(ts.Nanos)).UTC()
}

// RPCTimestampFromTime takes a time.Time and returns an RPC timestamp
func RPCTimestampFromTime(t time.Time) furanrpc.Timestamp {
	return furanrpc.Timestamp{
		Seconds: t.Unix(),
		Nanos:   int32(t.Nanosecond()),
	}
}

// EncryptAndSetGitHubCredential takes a GitHub credential, encrypts it and sets EncryptedGitHubCredential accordingly
func (b *Build) EncryptAndSetGitHubCredential(cred []byte, key [32]byte) error {
	var nonce [24]byte
	if n, err := rand.Read(nonce[:]); err != nil || n != len(nonce) {
		return errors.Wrapf(err, "error reading random bytes for nonce (read: %v)", n)
	}
	b.EncryptedGitHubCredential = secretbox.Seal(nonce[:], cred, &nonce, &key)
	return nil
}

// GetGitHubCredential returns the decrypted user token using key or error
func (b Build) GetGitHubCredential(key [32]byte) (string, error) {
	var nonce [24]byte
	copy(nonce[:], b.EncryptedGitHubCredential[:24])
	tkn, ok := secretbox.Open(nil, b.EncryptedGitHubCredential[24:], &nonce, &key)
	if !ok {
		return "", errors.New("decryption error (incorrect key?)")
	}
	if tkn == nil {
		return "", errors.New("decrypted token was nil")
	}
	return string(tkn), nil
}

// BuildOpts models all options required to perform a build
type BuildOpts struct {
	BuildID                uuid.UUID               `json:"-"`
	ContextPath, CommitSHA string                  `json:"-"` // set by Builder
	RelativeDockerfilePath string                  `json:"relative_dockerfile_path"`
	BuildArgs              map[string]string       `json:"build_args"`
	Cache                  furanrpc.BuildCacheOpts `json:"cache_opts"`
}

// Job describes methods on a single abstract build job
type Job interface {
	// Error returns a channel that will contain any errors associated with this Job
	Error() chan error
	// Running returns a channel that signals that the build the Job is executing has been updated to status Running
	// This indicates that the Furan sidecar has started and is executing successfully and will take responsibility for
	// tracking the build status from this point forward
	Running() chan struct{}
	// Logs returns all pod logs associated with the Job
	Logs() (map[string]map[string][]byte, error)
}

type JobRunner interface {
	Run(build Build) (Job, error)
}

// CacheFetcher describes an object that fetches and saves build cache
type CacheFetcher interface {
	// Fetch fetches the build cache for a build and returns a local filesystem
	// path where it was written. Caller is responsible for cleaning up the path when finished.
	Fetch(ctx context.Context, b Build) (string, error)
	// Save persists the build cache for a build located at path.
	// Caller is responsible for cleaning up the path afterward.
	Save(ctx context.Context, b Build, path string) error
}

// CodeFetcher represents an object capable of fetching code
type CodeFetcher interface {
	GetCommitSHA(ctx context.Context, repo, ref string) (string, error)
	Fetch(ctx context.Context, repo, ref, destinationPath string) error
}

// Builder describes an image build backend
type Builder interface {
	Build(ctx context.Context, opts BuildOpts) error
}

// BuilderManager describes an object that manages builds
type BuildManager interface {
	Start(ctx context.Context, opts BuildOpts) error
	Run(ctx context.Context, id uuid.UUID) error
}

type TagChecker interface {
	AllTagsExist(tags []string, repo string) (bool, []string, error)
}

// APIKey models a user-created API key
type APIKey struct {
	ID                            uuid.UUID
	Created                       time.Time
	GitHubUser, Name, Description string
	ReadOnly                      bool
}
<|endoftext|>


<commit_before>package store

import (
	"crypto/x509"
	"errors"
	"fmt"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/adamdecaf/cert-manage/pkg/certutil"
	"github.com/adamdecaf/cert-manage/pkg/file"
	"github.com/adamdecaf/cert-manage/pkg/whitelist"
)

var (
	openSSLCertPaths = []string{
		"/usr/local/etc/openssl/certs", // Darwin/OSX
	}

	openSSLRehashPaths = []string{
		"/usr/local/opt/openssl/bin/c_rehash", // Darwin/OSX
	}
)

type opensslStore struct{}

func OpenSSLStore() Store {
	return opensslStore{}
}

func (s opensslStore) Add(certs []*x509.Certificate) error {
	dir, err := s.findCertPath()
	if err != nil {
		return err
	}

	for i := range certs {
		fp := certutil.GetHexSHA256Fingerprint(*certs[i])
		path := filepath.Join(dir, fmt.Sprintf("%s.crt", fp))
		err = certutil.ToFile(path, certs[i:i+1])
		if err != nil {
			return err
		}
	}
	if len(certs) > 0 {
		return s.rehash()
	}
	return nil
}

func (s opensslStore) Backup() error {
	return nil
}

func (s opensslStore) GetInfo() *Info {
	out, err := exec.Command("openssl", "version").CombinedOutput()
	if err != nil {
		return &Info{ // just return something non-nil
			Name: "OpenSSL",
		}
	}

	// 'LibreSSL 2.2.7' or 'OpenSSL 1.0.2g  1 Mar 2016'
	parts := strings.Split(string(out), " ")

	return &Info{
		Name:    strings.TrimSpace(parts[0]),
		Version: strings.TrimSpace(parts[1]),
	}
}

func (s opensslStore) List() ([]*x509.Certificate, error) {
	return nil, nil
}

func (s opensslStore) Remove(wh whitelist.Whitelist) error {
	return nil
}

func (s opensslStore) Restore(where string) error {
	return nil
}

// TODO(adam): What do we do if multiple exist
func (s opensslStore) findCertPath() (string, error) {
	for i := range openSSLCertPaths {
		if file.Exists(openSSLCertPaths[i]) {
			return openSSLCertPaths[i], nil
		}
	}
	return "", errors.New("unable to find openssl cert directory")
}

func (s opensslStore) rehash() error {
	var bin string
	for i := range openSSLRehashPaths {
		if file.Exists(openSSLRehashPaths[i]) {
			bin = openSSLRehashPaths[i]
			break
		}
	}
	if bin == "" {
		return errors.New("unable to find openssl c_rehash binary")
	}

	// run c_rehash
	cmd := exec.Command(bin)
	out, err := cmd.CombinedOutput()
	if err != nil {
		if debug {
			fmt.Printf("store/openssl: Command was: %s\n", strings.Join(cmd.Args, " "))
			fmt.Printf("store/openssl: Output: %q\n", string(out))
		}
		return err
	}
	return nil
}
<commit_msg>store/openssl: add ubuntu openssl paths<commit_after>package store

import (
	"crypto/x509"
	"errors"
	"fmt"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/adamdecaf/cert-manage/pkg/certutil"
	"github.com/adamdecaf/cert-manage/pkg/file"
	"github.com/adamdecaf/cert-manage/pkg/whitelist"
)

var (
	openSSLCertPaths = []string{
		"/etc/ssl/certs/ca-certificates.crt", // Ubuntu
		"/usr/local/etc/openssl/certs", // Darwin/OSX
	}

	openSSLRehashPaths = []string{
		"/usr/bin/c_rehash", // Ubuntu
		"/usr/local/opt/openssl/bin/c_rehash", // Darwin/OSX
	}
)

type opensslStore struct{}

func OpenSSLStore() Store {
	return opensslStore{}
}

func (s opensslStore) Add(certs []*x509.Certificate) error {
	dir, err := s.findCertPath()
	if err != nil {
		return err
	}

	for i := range certs {
		fp := certutil.GetHexSHA256Fingerprint(*certs[i])
		path := filepath.Join(dir, fmt.Sprintf("%s.crt", fp))
		err = certutil.ToFile(path, certs[i:i+1])
		if err != nil {
			return err
		}
	}
	if len(certs) > 0 {
		return s.rehash()
	}
	return nil
}

func (s opensslStore) Backup() error {
	return nil
}

func (s opensslStore) GetInfo() *Info {
	out, err := exec.Command("openssl", "version").CombinedOutput()
	if err != nil {
		return &Info{ // just return something non-nil
			Name: "OpenSSL",
		}
	}

	// 'LibreSSL 2.2.7' or 'OpenSSL 1.0.2g  1 Mar 2016'
	parts := strings.Split(string(out), " ")

	return &Info{
		Name:    strings.TrimSpace(parts[0]),
		Version: strings.TrimSpace(parts[1]),
	}
}

func (s opensslStore) List() ([]*x509.Certificate, error) {
	return nil, nil
}

func (s opensslStore) Remove(wh whitelist.Whitelist) error {
	return nil
}

func (s opensslStore) Restore(where string) error {
	return nil
}

// TODO(adam): What do we do if multiple exist
func (s opensslStore) findCertPath() (string, error) {
	for i := range openSSLCertPaths {
		if file.Exists(openSSLCertPaths[i]) {
			return openSSLCertPaths[i], nil
		}
	}
	return "", errors.New("unable to find openssl cert directory")
}

func (s opensslStore) rehash() error {
	var bin string
	for i := range openSSLRehashPaths {
		if file.Exists(openSSLRehashPaths[i]) {
			bin = openSSLRehashPaths[i]
			break
		}
	}
	if bin == "" {
		return errors.New("unable to find openssl c_rehash binary")
	}

	// run c_rehash
	cmd := exec.Command(bin)
	out, err := cmd.CombinedOutput()
	if err != nil {
		if debug {
			fmt.Printf("store/openssl: Command was: %s\n", strings.Join(cmd.Args, " "))
			fmt.Printf("store/openssl: Output: %q\n", string(out))
		}
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 CodisLabs. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package rpc

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

	"github.com/CodisLabs/codis/pkg/utils/errors"
	"github.com/CodisLabs/codis/pkg/utils/log"
	"github.com/CodisLabs/codis/pkg/utils/sync2/atomic2"
	"github.com/CodisLabs/codis/pkg/utils/trace"
)

const (
	MethodGet = "GET"
	MethodPut = "PUT"
)

var client *http.Client

func init() {
	var dials atomic2.Int64
	tr := &http.Transport{}
	tr.Dial = func(network, addr string) (net.Conn, error) {
		c, err := net.DialTimeout(network, addr, time.Second*10)
		if err == nil {
			log.Debugf("rpc: dial new connection to [%d] %s - %s",
				dials.Incr()-1, network, addr)
		}
		return c, err
	}
	client = &http.Client{
		Transport: tr,
		Timeout:   time.Second * 30,
	}
	go func() {
		for {
			time.Sleep(time.Minute)
			tr.CloseIdleConnections()
		}
	}()
}

type RemoteError struct {
	Cause string
	Stack trace.Stack
}

func (e *RemoteError) Error() string {
	return e.Cause
}

func (e *RemoteError) TracedError() error {
	return &errors.TracedError{
		Cause: errors.New("[Remote Error] " + e.Cause),
		Stack: e.Stack,
	}
}

func NewRemoteError(err error) *RemoteError {
	if err == nil {
		return nil
	}
	if v, ok := err.(*RemoteError); ok {
		return v
	}
	return &RemoteError{
		Cause: err.Error(),
		Stack: errors.Stack(err),
	}
}

func responseBodyAsBytes(rsp *http.Response) ([]byte, error) {
	b, err := ioutil.ReadAll(rsp.Body)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return b, nil
}

func responseBodyAsError(rsp *http.Response) (error, error) {
	b, err := responseBodyAsBytes(rsp)
	if err != nil {
		return nil, err
	}
	if len(b) == 0 {
		return nil, errors.Errorf("remote error is empty")
	}
	e := &RemoteError{}
	if err := json.Unmarshal(b, e); err != nil {
		return nil, errors.Trace(err)
	}
	return e.TracedError(), nil
}

func apiMarshalJson(v interface{}) ([]byte, error) {
	return json.Marshal(v)
}

func apiRequestJson(method string, url string, args, reply interface{}) error {
	var body []byte
	if args != nil {
		b, err := apiMarshalJson(args)
		if err != nil {
			return errors.Trace(err)
		}
		body = b
	}

	req, err := http.NewRequest(method, url, bytes.NewReader(body))
	if err != nil {
		return errors.Trace(err)
	}
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Content-Length", strconv.Itoa(len(body)))
	}

	var start = time.Now()

	rsp, err := client.Do(req)
	if err != nil {
		return errors.Trace(err)
	}
	defer func() {
		io.Copy(ioutil.Discard, rsp.Body)
		rsp.Body.Close()
		log.Debugf("call rpc [%s] %s in %v", method, url, time.Since(start))
	}()

	switch rsp.StatusCode {
	case 200:
		b, err := responseBodyAsBytes(rsp)
		if err != nil {
			return err
		}
		if reply == nil {
			return nil
		}
		if err := json.Unmarshal(b, reply); err != nil {
			return errors.Trace(err)
		} else {
			return nil
		}
	case 1500:
		e, err := responseBodyAsError(rsp)
		if err != nil {
			return err
		} else {
			return e
		}
	default:
		return errors.Errorf("[%d] %s - %s", rsp.StatusCode, http.StatusText(rsp.StatusCode), url)
	}
}

func ApiGetJson(url string, reply interface{}) error {
	return apiRequestJson(MethodGet, url, nil, reply)
}

func ApiPutJson(url string, args, reply interface{}) error {
	return apiRequestJson(MethodPut, url, args, reply)
}

func ApiResponseError(err error) (int, string) {
	if err == nil {
		return 1500, ""
	}
	b, err := apiMarshalJson(NewRemoteError(err))
	if err != nil {
		return 1500, ""
	} else {
		return 1500, string(b)
	}
}

func ApiResponseJson(v interface{}) (int, string) {
	b, err := apiMarshalJson(v)
	if err != nil {
		return ApiResponseError(errors.Trace(err))
	} else {
		return 200, string(b)
	}
}

func EncodeURL(host string, format string, args ...interface{}) string {
	var u url.URL
	u.Scheme = "http"
	u.Host = host
	u.Path = fmt.Sprintf(format, args...)
	return u.String()
}
<commit_msg>Fix, change status code for rpc failure<commit_after>// Copyright 2016 CodisLabs. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package rpc

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

	"github.com/CodisLabs/codis/pkg/utils/errors"
	"github.com/CodisLabs/codis/pkg/utils/log"
	"github.com/CodisLabs/codis/pkg/utils/sync2/atomic2"
	"github.com/CodisLabs/codis/pkg/utils/trace"
)

const (
	MethodGet = "GET"
	MethodPut = "PUT"
)

var client *http.Client

func init() {
	var dials atomic2.Int64
	tr := &http.Transport{}
	tr.Dial = func(network, addr string) (net.Conn, error) {
		c, err := net.DialTimeout(network, addr, time.Second*10)
		if err == nil {
			log.Debugf("rpc: dial new connection to [%d] %s - %s",
				dials.Incr()-1, network, addr)
		}
		return c, err
	}
	client = &http.Client{
		Transport: tr,
		Timeout:   time.Second * 30,
	}
	go func() {
		for {
			time.Sleep(time.Minute)
			tr.CloseIdleConnections()
		}
	}()
}

type RemoteError struct {
	Cause string
	Stack trace.Stack
}

func (e *RemoteError) Error() string {
	return e.Cause
}

func (e *RemoteError) TracedError() error {
	return &errors.TracedError{
		Cause: errors.New("[Remote Error] " + e.Cause),
		Stack: e.Stack,
	}
}

func NewRemoteError(err error) *RemoteError {
	if err == nil {
		return nil
	}
	if v, ok := err.(*RemoteError); ok {
		return v
	}
	return &RemoteError{
		Cause: err.Error(),
		Stack: errors.Stack(err),
	}
}

func responseBodyAsBytes(rsp *http.Response) ([]byte, error) {
	b, err := ioutil.ReadAll(rsp.Body)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return b, nil
}

func responseBodyAsError(rsp *http.Response) (error, error) {
	b, err := responseBodyAsBytes(rsp)
	if err != nil {
		return nil, err
	}
	if len(b) == 0 {
		return nil, errors.Errorf("remote error is empty")
	}
	e := &RemoteError{}
	if err := json.Unmarshal(b, e); err != nil {
		return nil, errors.Trace(err)
	}
	return e.TracedError(), nil
}

func apiMarshalJson(v interface{}) ([]byte, error) {
	return json.Marshal(v)
}

func apiRequestJson(method string, url string, args, reply interface{}) error {
	var body []byte
	if args != nil {
		b, err := apiMarshalJson(args)
		if err != nil {
			return errors.Trace(err)
		}
		body = b
	}

	req, err := http.NewRequest(method, url, bytes.NewReader(body))
	if err != nil {
		return errors.Trace(err)
	}
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Content-Length", strconv.Itoa(len(body)))
	}

	var start = time.Now()

	rsp, err := client.Do(req)
	if err != nil {
		return errors.Trace(err)
	}
	defer func() {
		io.Copy(ioutil.Discard, rsp.Body)
		rsp.Body.Close()
		log.Debugf("call rpc [%s] %s in %v", method, url, time.Since(start))
	}()

	switch rsp.StatusCode {
	case 200:
		b, err := responseBodyAsBytes(rsp)
		if err != nil {
			return err
		}
		if reply == nil {
			return nil
		}
		if err := json.Unmarshal(b, reply); err != nil {
			return errors.Trace(err)
		} else {
			return nil
		}
	case 800, 1500:
		e, err := responseBodyAsError(rsp)
		if err != nil {
			return err
		} else {
			return e
		}
	default:
		return errors.Errorf("[%d] %s - %s", rsp.StatusCode, http.StatusText(rsp.StatusCode), url)
	}
}

func ApiGetJson(url string, reply interface{}) error {
	return apiRequestJson(MethodGet, url, nil, reply)
}

func ApiPutJson(url string, args, reply interface{}) error {
	return apiRequestJson(MethodPut, url, args, reply)
}

func ApiResponseError(err error) (int, string) {
	if err == nil {
		return 800, ""
	}
	b, err := apiMarshalJson(NewRemoteError(err))
	if err != nil {
		return 800, ""
	} else {
		return 800, string(b)
	}
}

func ApiResponseJson(v interface{}) (int, string) {
	b, err := apiMarshalJson(v)
	if err != nil {
		return ApiResponseError(errors.Trace(err))
	} else {
		return 200, string(b)
	}
}

func EncodeURL(host string, format string, args ...interface{}) string {
	var u url.URL
	u.Scheme = "http"
	u.Host = host
	u.Path = fmt.Sprintf(format, args...)
	return u.String()
}
<|endoftext|>


<commit_before>// Copyright Microsoft Corp.
// All rights reserved.

package ipam

import (
	"fmt"
	"net"
	"strings"
)

const (
	// Default address space IDs.
	localDefaultAddressSpaceId  = "LocalDefaultAddressSpace"
	globalDefaultAddressSpaceId = "GlobalDefaultAddressSpace"

	// Address space scopes.
	localScope  = "local"
	globalScope = "global"
)

// Represents the key to an address pool.
type addressPoolId struct {
	AsId        string
	Subnet      string
	ChildSubnet string
}

// Represents a set of non-overlapping address pools.
type addressSpace struct {
	Id    string
	Scope string
	Pools map[string]*addressPool
	epoch int
}

// Represents a subnet and the set of addresses in it.
type addressPool struct {
	as        *addressSpace
	Id        string
	IfName    string
	Subnet    net.IPNet
	Addresses map[string]*addressRecord
	IsIPv6    bool
	Priority  int
	InUse     bool
	epoch     int
}

// Represents an IP address in a pool.
type addressRecord struct {
	Addr  net.IP
	InUse bool
	epoch int
}

//
// AddressPoolId
//

// Creates a new address pool ID object.
func NewAddressPoolId(asId string, subnet string, childSubnet string) *addressPoolId {
	return &addressPoolId{
		AsId:        asId,
		Subnet:      subnet,
		ChildSubnet: childSubnet,
	}
}

// Creates a new pool ID from a string representation.
func NewAddressPoolIdFromString(s string) (*addressPoolId, error) {
	var pid addressPoolId

	p := strings.Split(s, "|")
	if len(p) > 3 {
		return nil, errInvalidPoolId
	}

	pid.AsId = p[0]
	if len(p) >= 2 {
		pid.Subnet = p[1]
	}
	if len(p) == 3 {
		pid.ChildSubnet = p[2]
	}

	return &pid, nil
}

// Returns the string representation of a pool ID.
func (pid *addressPoolId) String() string {
	s := fmt.Sprintf("%s|%s", pid.AsId, pid.Subnet)
	if pid.ChildSubnet != "" {
		s = fmt.Sprintf("%s|%s", s, pid.ChildSubnet)
	}
	return s
}

//
// AddressSpace
//

// Creates a new addressSpace object.
func (am *addressManager) newAddressSpace(id string, scope string) (*addressSpace, error) {
	if scope != localScope && scope != globalScope {
		return nil, errInvalidScope
	}

	return &addressSpace{
		Id:    id,
		Scope: scope,
		Pools: make(map[string]*addressPool),
	}, nil
}

// Returns the address space with the given ID.
func (am *addressManager) getAddressSpace(id string) (*addressSpace, error) {
	as := am.AddrSpaces[id]
	if as == nil {
		return nil, errInvalidAddressSpace
	}

	return as, nil
}

// Sets a new or updates an existing address space.
func (am *addressManager) setAddressSpace(as *addressSpace) error {
	as1, ok := am.AddrSpaces[as.Id]
	if !ok {
		am.AddrSpaces[as.Id] = as
	} else {
		as1.merge(as)
	}

	// Notify NetPlugin of external interfaces.
	if am.netApi != nil {
		for _, ap := range as.Pools {
			am.netApi.AddExternalInterface(ap.IfName, ap.Subnet.String())
		}
	}

	am.save()

	return nil
}

// Merges a new address space to an existing one.
func (as *addressSpace) merge(newas *addressSpace) {
	// The new epoch after the merge.
	as.epoch++

	// Add new pools and addresses.
	for pk, pv := range newas.Pools {
		ap := as.Pools[pk]

		if ap == nil {
			// This is a new address pool.
			// Merge it to the existing address space.
			as.Pools[pk] = pv
			pv.as = as
			pv.epoch = as.epoch
		} else {
			// This pool already exists.
			// Compare address records one by one.
			for ak, av := range pv.Addresses {
				ar := ap.Addresses[ak]

				if ar == nil {
					// This is a new address record.
					// Merge it to the existing address pool.
					ap.Addresses[ak] = av
					av.epoch = as.epoch
				} else {
					// This address record already exists.
					ar.epoch = as.epoch
				}

				delete(pv.Addresses, ak)
			}

			pv.as = nil
		}

		delete(newas.Pools, pk)
	}

	// Cleanup stale pools and addresses from the old epoch.
	// Those currently in use will be deleted after they are released.
	for pk, pv := range as.Pools {
		if pv.epoch < as.epoch {
			// This pool may have stale addresses.
			for ak, av := range pv.Addresses {
				if av.epoch == as.epoch || av.InUse {
					// Pool has at least one valid or in-use address.
					pv.epoch = as.epoch
				} else {
					// This address is no longer available.
					delete(pv.Addresses, ak)
				}
			}

			// Delete the pool if it has no addresses left.
			if pv.epoch < as.epoch && !pv.InUse {
				pv.as = nil
				delete(as.Pools, pk)
			}
		}
	}

	return
}

// Creates a new addressPool object.
func (as *addressSpace) newAddressPool(ifName string, priority int, subnet *net.IPNet) (*addressPool, error) {
	id := subnet.String()

	pool, ok := as.Pools[id]
	if ok {
		return pool, errAddressPoolExists
	}

	v6 := (subnet.IP.To4() == nil)

	pool = &addressPool{
		as:        as,
		Id:        id,
		IfName:    ifName,
		Subnet:    *subnet,
		Addresses: make(map[string]*addressRecord),
		IsIPv6:    v6,
		Priority:  priority,
		epoch:     as.epoch,
	}

	as.Pools[id] = pool

	return pool, nil
}

// Returns the address pool with the given pool ID.
func (as *addressSpace) getAddressPool(poolId string) (*addressPool, error) {
	ap := as.Pools[poolId]
	if ap == nil {
		return nil, errInvalidPoolId
	}

	return ap, nil
}

// Requests a new address pool from the address space.
func (as *addressSpace) requestPool(poolId string, subPoolId string, options map[string]string, v6 bool) (*addressPool, error) {
	var ap *addressPool

	if poolId != "" {
		// Return the specific address pool requested.
		ap = as.Pools[poolId]
		if ap == nil {
			return nil, errAddressPoolNotFound
		}

		// Fail if requested pool is already in use.
		if ap.InUse {
			return nil, errAddressPoolInUse
		}
	} else {
		// Return any available address pool.
		highestPriority := -1

		for _, pool := range as.Pools {
			// Skip if pool is already in use.
			if pool.InUse {
				continue
			}

			// Pick a pool from the same address family.
			if pool.IsIPv6 != v6 {
				continue
			}

			// Pick the pool with the highest priority.
			if pool.Priority > highestPriority {
				highestPriority = pool.Priority
				ap = pool
			}
		}

		if ap == nil {
			return nil, errNoAvailableAddressPools
		}
	}

	ap.InUse = true

	return ap, nil
}

// Releases a previously requested address pool back to its address space.
func (as *addressSpace) releasePool(poolId string) error {
	ap, ok := as.Pools[poolId]
	if !ok {
		return errAddressPoolNotFound
	}

	if !ap.InUse {
		return errAddressPoolNotInUse
	}

	ap.InUse = false

	// Delete address pool if it is no longer available.
	if ap.epoch < as.epoch {
		delete(as.Pools, poolId)
	}

	return nil
}

//
// AddressPool
//

// Creates a new addressRecord object.
func (ap *addressPool) newAddressRecord(addr *net.IP) (*addressRecord, error) {
	id := addr.String()

	if !ap.Subnet.Contains(*addr) {
		return nil, errInvalidAddress
	}

	ar, ok := ap.Addresses[id]
	if ok {
		return ar, errAddressExists
	}

	ar = &addressRecord{
		Addr:  *addr,
		epoch: ap.epoch,
	}

	ap.Addresses[id] = ar

	return ar, nil
}

// Requests a new address from the address pool.
func (ap *addressPool) requestAddress(address string, options map[string]string) (string, error) {
	var ar *addressRecord

	if address != "" {
		// Return the specific address requested.
		ar = ap.Addresses[address]
		if ar == nil {
			return "", errAddressNotFound
		}
		if ar.InUse {
			return "", errAddressInUse
		}
	} else {
		// Return any available address.
		for _, ar = range ap.Addresses {
			if !ar.InUse {
				break
			}
		}

		if ar == nil {
			return "", errNoAvailableAddresses
		}
	}

	ar.InUse = true

	// Return address in CIDR notation.
	addr := net.IPNet{
		IP:   ar.Addr,
		Mask: ap.Subnet.Mask,
	}

	return addr.String(), nil
}

// Releases a previously requested address back to its address pool.
func (ap *addressPool) releaseAddress(address string) error {
	ar := ap.Addresses[address]
	if ar == nil {
		return errAddressNotFound
	}
	if !ar.InUse {
		return errAddressNotInUse
	}

	ar.InUse = false

	// Delete address record if it is no longer available.
	if ar.epoch < ap.as.epoch {
		delete(ap.Addresses, address)
	}

	return nil
}
<commit_msg>Added reference counting for shared pools<commit_after>// Copyright Microsoft Corp.
// All rights reserved.

package ipam

import (
	"fmt"
	"net"
	"strings"
)

const (
	// Default address space IDs.
	localDefaultAddressSpaceId  = "LocalDefaultAddressSpace"
	globalDefaultAddressSpaceId = "GlobalDefaultAddressSpace"

	// Address space scopes.
	localScope  = "local"
	globalScope = "global"
)

// Represents the key to an address pool.
type addressPoolId struct {
	AsId        string
	Subnet      string
	ChildSubnet string
}

// Represents a set of non-overlapping address pools.
type addressSpace struct {
	Id    string
	Scope string
	Pools map[string]*addressPool
	epoch int
}

// Represents a subnet and the set of addresses in it.
type addressPool struct {
	as        *addressSpace
	Id        string
	IfName    string
	Subnet    net.IPNet
	Addresses map[string]*addressRecord
	IsIPv6    bool
	Priority  int
	RefCount  int
	epoch     int
}

// Represents an IP address in a pool.
type addressRecord struct {
	Addr  net.IP
	InUse bool
	epoch int
}

//
// AddressPoolId
//

// Creates a new address pool ID object.
func NewAddressPoolId(asId string, subnet string, childSubnet string) *addressPoolId {
	return &addressPoolId{
		AsId:        asId,
		Subnet:      subnet,
		ChildSubnet: childSubnet,
	}
}

// Creates a new pool ID from a string representation.
func NewAddressPoolIdFromString(s string) (*addressPoolId, error) {
	var pid addressPoolId

	p := strings.Split(s, "|")
	if len(p) > 3 {
		return nil, errInvalidPoolId
	}

	pid.AsId = p[0]
	if len(p) >= 2 {
		pid.Subnet = p[1]
	}
	if len(p) == 3 {
		pid.ChildSubnet = p[2]
	}

	return &pid, nil
}

// Returns the string representation of a pool ID.
func (pid *addressPoolId) String() string {
	s := fmt.Sprintf("%s|%s", pid.AsId, pid.Subnet)
	if pid.ChildSubnet != "" {
		s = fmt.Sprintf("%s|%s", s, pid.ChildSubnet)
	}
	return s
}

//
// AddressSpace
//

// Creates a new addressSpace object.
func (am *addressManager) newAddressSpace(id string, scope string) (*addressSpace, error) {
	if scope != localScope && scope != globalScope {
		return nil, errInvalidScope
	}

	return &addressSpace{
		Id:    id,
		Scope: scope,
		Pools: make(map[string]*addressPool),
	}, nil
}

// Returns the address space with the given ID.
func (am *addressManager) getAddressSpace(id string) (*addressSpace, error) {
	as := am.AddrSpaces[id]
	if as == nil {
		return nil, errInvalidAddressSpace
	}

	return as, nil
}

// Sets a new or updates an existing address space.
func (am *addressManager) setAddressSpace(as *addressSpace) error {
	as1, ok := am.AddrSpaces[as.Id]
	if !ok {
		am.AddrSpaces[as.Id] = as
	} else {
		as1.merge(as)
	}

	// Notify NetPlugin of external interfaces.
	if am.netApi != nil {
		for _, ap := range as.Pools {
			am.netApi.AddExternalInterface(ap.IfName, ap.Subnet.String())
		}
	}

	am.save()

	return nil
}

// Merges a new address space to an existing one.
func (as *addressSpace) merge(newas *addressSpace) {
	// The new epoch after the merge.
	as.epoch++

	// Add new pools and addresses.
	for pk, pv := range newas.Pools {
		ap := as.Pools[pk]

		if ap == nil {
			// This is a new address pool.
			// Merge it to the existing address space.
			as.Pools[pk] = pv
			pv.as = as
			pv.epoch = as.epoch
		} else {
			// This pool already exists.
			// Compare address records one by one.
			for ak, av := range pv.Addresses {
				ar := ap.Addresses[ak]

				if ar == nil {
					// This is a new address record.
					// Merge it to the existing address pool.
					ap.Addresses[ak] = av
					av.epoch = as.epoch
				} else {
					// This address record already exists.
					ar.epoch = as.epoch
				}

				delete(pv.Addresses, ak)
			}

			pv.as = nil
		}

		delete(newas.Pools, pk)
	}

	// Cleanup stale pools and addresses from the old epoch.
	// Those currently in use will be deleted after they are released.
	for pk, pv := range as.Pools {
		if pv.epoch < as.epoch {
			// This pool may have stale addresses.
			for ak, av := range pv.Addresses {
				if av.epoch == as.epoch || av.InUse {
					// Pool has at least one valid or in-use address.
					pv.epoch = as.epoch
				} else {
					// This address is no longer available.
					delete(pv.Addresses, ak)
				}
			}

			// Delete the pool if it has no addresses left.
			if pv.epoch < as.epoch && !pv.isInUse() {
				pv.as = nil
				delete(as.Pools, pk)
			}
		}
	}

	return
}

// Creates a new addressPool object.
func (as *addressSpace) newAddressPool(ifName string, priority int, subnet *net.IPNet) (*addressPool, error) {
	id := subnet.String()

	pool, ok := as.Pools[id]
	if ok {
		return pool, errAddressPoolExists
	}

	v6 := (subnet.IP.To4() == nil)

	pool = &addressPool{
		as:        as,
		Id:        id,
		IfName:    ifName,
		Subnet:    *subnet,
		Addresses: make(map[string]*addressRecord),
		IsIPv6:    v6,
		Priority:  priority,
		epoch:     as.epoch,
	}

	as.Pools[id] = pool

	return pool, nil
}

// Returns the address pool with the given pool ID.
func (as *addressSpace) getAddressPool(poolId string) (*addressPool, error) {
	ap := as.Pools[poolId]
	if ap == nil {
		return nil, errInvalidPoolId
	}

	return ap, nil
}

// Requests a new address pool from the address space.
func (as *addressSpace) requestPool(poolId string, subPoolId string, options map[string]string, v6 bool) (*addressPool, error) {
	var ap *addressPool

	if poolId != "" {
		// Return the specific address pool requested.
		ap = as.Pools[poolId]
		if ap == nil {
			return nil, errAddressPoolNotFound
		}
	} else {
		// Return any available address pool.
		highestPriority := -1

		for _, pool := range as.Pools {
			// Skip if pool is already in use.
			if pool.isInUse() {
				continue
			}

			// Pick a pool from the same address family.
			if pool.IsIPv6 != v6 {
				continue
			}

			// Pick the pool with the highest priority.
			if pool.Priority > highestPriority {
				highestPriority = pool.Priority
				ap = pool
			}
		}

		if ap == nil {
			return nil, errNoAvailableAddressPools
		}
	}

	ap.RefCount++

	return ap, nil
}

// Releases a previously requested address pool back to its address space.
func (as *addressSpace) releasePool(poolId string) error {
	ap, ok := as.Pools[poolId]
	if !ok {
		return errAddressPoolNotFound
	}

	if !ap.isInUse() {
		return errAddressPoolNotInUse
	}

	ap.RefCount--

	// Delete address pool if it is no longer available.
	if ap.epoch < as.epoch && !ap.isInUse() {
		delete(as.Pools, poolId)
	}

	return nil
}

//
// AddressPool
//

// Returns if an addess pool is currently in use.
func (ap *addressPool) isInUse() bool {
	return ap.RefCount > 0
}

// Creates a new addressRecord object.
func (ap *addressPool) newAddressRecord(addr *net.IP) (*addressRecord, error) {
	id := addr.String()

	if !ap.Subnet.Contains(*addr) {
		return nil, errInvalidAddress
	}

	ar, ok := ap.Addresses[id]
	if ok {
		return ar, errAddressExists
	}

	ar = &addressRecord{
		Addr:  *addr,
		epoch: ap.epoch,
	}

	ap.Addresses[id] = ar

	return ar, nil
}

// Requests a new address from the address pool.
func (ap *addressPool) requestAddress(address string, options map[string]string) (string, error) {
	var ar *addressRecord

	if address != "" {
		// Return the specific address requested.
		ar = ap.Addresses[address]
		if ar == nil {
			return "", errAddressNotFound
		}
		if ar.InUse {
			return "", errAddressInUse
		}
	} else {
		// Return any available address.
		for _, ar = range ap.Addresses {
			if !ar.InUse {
				break
			}
		}

		if ar == nil {
			return "", errNoAvailableAddresses
		}
	}

	ar.InUse = true

	// Return address in CIDR notation.
	addr := net.IPNet{
		IP:   ar.Addr,
		Mask: ap.Subnet.Mask,
	}

	return addr.String(), nil
}

// Releases a previously requested address back to its address pool.
func (ap *addressPool) releaseAddress(address string) error {
	ar := ap.Addresses[address]
	if ar == nil {
		return errAddressNotFound
	}
	if !ar.InUse {
		return errAddressNotInUse
	}

	ar.InUse = false

	// Delete address record if it is no longer available.
	if ar.epoch < ap.as.epoch {
		delete(ap.Addresses, address)
	}

	return nil
}
<|endoftext|>


<commit_before>package consumergroup

import (
	"crypto/rand"
	"fmt"
	"io"
	"os"
	"sort"

	"github.com/wvanbergen/kazoo-go"
)

func retrievePartitionLeaders(partitions kazoo.PartitionList) (partitionLeaders, error) {

	pls := make(partitionLeaders, 0, len(partitions))
	for _, partition := range partitions {
		leader, err := partition.Leader()
		if err != nil {
			return nil, err
		}

		pl := partitionLeader{id: partition.ID, leader: leader, partition: partition}
		pls = append(pls, pl)
	}

	return pls, nil
}

// Divides a set of partitions between a set of consumers.
func dividePartitionsBetweenConsumers(consumers kazoo.ConsumergroupInstanceList, partitions partitionLeaders) map[string][]*kazoo.Partition {
	result := make(map[string][]*kazoo.Partition)

	plen := len(partitions)
	clen := len(consumers)
	if clen == 0 {
		return result
	}

	sort.Sort(partitions)
	sort.Sort(consumers)

	n := plen / clen
	m := plen % clen
	p := 0
	for i, consumer := range consumers {
		first := p
		last := first + n
		if m > 0 && i < m {
			last++
		}
		if last > plen {
			last = plen
		}

		for _, pl := range partitions[first:last] {
			result[consumer.ID] = append(result[consumer.ID], pl.partition)
		}
		p = last
	}

	return result
}

type partitionLeader struct {
	id        int32
	leader    int32
	partition *kazoo.Partition
}

// A sortable slice of PartitionLeader structs
type partitionLeaders []partitionLeader

func (pls partitionLeaders) Len() int {
	return len(pls)
}

func (pls partitionLeaders) Less(i, j int) bool {
	return pls[i].leader < pls[j].leader || (pls[i].leader == pls[j].leader && pls[i].id < pls[j].id)
}

func (s partitionLeaders) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func generateUUID() (string, error) {
	uuid := make([]byte, 16)
	n, err := io.ReadFull(rand.Reader, uuid)
	if n != len(uuid) || err != nil {
		return "", err
	}
	// variant bits; see section 4.1.1
	uuid[8] = uuid[8]&^0xc0 | 0x80
	// version 4 (pseudo-random); see section 4.1.3
	uuid[6] = uuid[6]&^0xf0 | 0x40
	return fmt.Sprintf("%x-%x-%x-%x-%x", uuid[0:4], uuid[4:6], uuid[6:8], uuid[8:10], uuid[10:]), nil
}

func generateConsumerID() (consumerID string, err error) {
	var uuid, hostname string

	uuid, err = generateUUID()
	if err != nil {
		return
	}

	hostname, err = os.Hostname()
	if err != nil {
		return
	}

	consumerID = fmt.Sprintf("%s:%s", hostname, uuid)
	return
}
<commit_msg>Default hostname to ENV HOST<commit_after>package consumergroup

import (
	"crypto/rand"
	"fmt"
	"io"
	"os"
	"sort"

	"github.com/wvanbergen/kazoo-go"
)

func retrievePartitionLeaders(partitions kazoo.PartitionList) (partitionLeaders, error) {

	pls := make(partitionLeaders, 0, len(partitions))
	for _, partition := range partitions {
		leader, err := partition.Leader()
		if err != nil {
			return nil, err
		}

		pl := partitionLeader{id: partition.ID, leader: leader, partition: partition}
		pls = append(pls, pl)
	}

	return pls, nil
}

// Divides a set of partitions between a set of consumers.
func dividePartitionsBetweenConsumers(consumers kazoo.ConsumergroupInstanceList, partitions partitionLeaders) map[string][]*kazoo.Partition {
	result := make(map[string][]*kazoo.Partition)

	plen := len(partitions)
	clen := len(consumers)
	if clen == 0 {
		return result
	}

	sort.Sort(partitions)
	sort.Sort(consumers)

	n := plen / clen
	m := plen % clen
	p := 0
	for i, consumer := range consumers {
		first := p
		last := first + n
		if m > 0 && i < m {
			last++
		}
		if last > plen {
			last = plen
		}

		for _, pl := range partitions[first:last] {
			result[consumer.ID] = append(result[consumer.ID], pl.partition)
		}
		p = last
	}

	return result
}

type partitionLeader struct {
	id        int32
	leader    int32
	partition *kazoo.Partition
}

// A sortable slice of PartitionLeader structs
type partitionLeaders []partitionLeader

func (pls partitionLeaders) Len() int {
	return len(pls)
}

func (pls partitionLeaders) Less(i, j int) bool {
	return pls[i].leader < pls[j].leader || (pls[i].leader == pls[j].leader && pls[i].id < pls[j].id)
}

func (s partitionLeaders) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func generateUUID() (string, error) {
	uuid := make([]byte, 16)
	n, err := io.ReadFull(rand.Reader, uuid)
	if n != len(uuid) || err != nil {
		return "", err
	}
	// variant bits; see section 4.1.1
	uuid[8] = uuid[8]&^0xc0 | 0x80
	// version 4 (pseudo-random); see section 4.1.3
	uuid[6] = uuid[6]&^0xf0 | 0x40
	return fmt.Sprintf("%x-%x-%x-%x-%x", uuid[0:4], uuid[4:6], uuid[6:8], uuid[8:10], uuid[10:]), nil
}

func generateConsumerID() (consumerID string, err error) {
	var uuid, hostname string

	uuid, err = generateUUID()
	if err != nil {
		return
	}
	if hostname = os.Getenv("HOST"); hostname == "" {
		hostname, err = os.Hostname()
		if err != nil {
			return
		}
	}

	consumerID = fmt.Sprintf("%s:%s", hostname, uuid)
	return
}
<|endoftext|>


<commit_before>package gorm

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
	"time"

	"reflect"
)

type Scope struct {
	Search          *search
	Value           interface{}
	Sql             string
	SqlVars         []interface{}
	db              *DB
	indirectValue   *reflect.Value
	instanceId      string
	primaryKeyField *Field
	skipLeft        bool
	fields          map[string]*Field
	selectAttrs     *[]string
}

func (scope *Scope) IndirectValue() reflect.Value {
	if scope.indirectValue == nil {
		value := reflect.Indirect(reflect.ValueOf(scope.Value))
		if value.Kind() == reflect.Ptr {
			value = value.Elem()
		}
		scope.indirectValue = &value
	}
	return *scope.indirectValue
}

func (scope *Scope) NeedPtr() *Scope {
	reflectKind := reflect.ValueOf(scope.Value).Kind()
	if !((reflectKind == reflect.Invalid) || (reflectKind == reflect.Ptr)) {
		err := fmt.Errorf("%v %v\n", fileWithLineNum(), "using unaddressable value")
		scope.Err(err)
		fmt.Printf(err.Error())
	}
	return scope
}

// New create a new Scope without search information
func (scope *Scope) New(value interface{}) *Scope {
	return &Scope{db: scope.NewDB(), Search: &search{}, Value: value}
}

// NewDB create a new DB without search information
func (scope *Scope) NewDB() *DB {
	if scope.db != nil {
		db := scope.db.clone()
		db.search = nil
		db.Value = nil
		return db
	}
	return nil
}

func (scope *Scope) DB() *DB {
	return scope.db
}

// SqlDB return *sql.DB
func (scope *Scope) SqlDB() sqlCommon {
	return scope.db.db
}

// SkipLeft skip remaining callbacks
func (scope *Scope) SkipLeft() {
	scope.skipLeft = true
}

// Quote used to quote database column name according to database dialect
func (scope *Scope) Quote(str string) string {
	if strings.Index(str, ".") != -1 {
		newStrs := []string{}
		for _, str := range strings.Split(str, ".") {
			newStrs = append(newStrs, scope.Dialect().Quote(str))
		}
		return strings.Join(newStrs, ".")
	} else {
		return scope.Dialect().Quote(str)
	}
}

func (scope *Scope) QuoteIfPossible(str string) string {
	if regexp.MustCompile("^[a-zA-Z]+(.[a-zA-Z]+)*$").MatchString(str) {
		return scope.Quote(str)
	}
	return str
}

// Dialect get dialect
func (scope *Scope) Dialect() Dialect {
	return scope.db.parent.dialect
}

// Err write error
func (scope *Scope) Err(err error) error {
	if err != nil {
		scope.db.AddError(err)
	}
	return err
}

// Log print log message
func (scope *Scope) Log(v ...interface{}) {
	scope.db.log(v...)
}

// HasError check if there are any error
func (scope *Scope) HasError() bool {
	return scope.db.Error != nil
}

func (scope *Scope) PrimaryFields() []*Field {
	var fields = []*Field{}
	for _, field := range scope.GetModelStruct().PrimaryFields {
		fields = append(fields, scope.Fields()[field.DBName])
	}
	return fields
}

func (scope *Scope) PrimaryField() *Field {
	if primaryFields := scope.GetModelStruct().PrimaryFields; len(primaryFields) > 0 {
		if len(primaryFields) > 1 {
			if field, ok := scope.Fields()["id"]; ok {
				return field
			}
		}
		return scope.Fields()[primaryFields[0].DBName]
	}
	return nil
}

// PrimaryKey get the primary key's column name
func (scope *Scope) PrimaryKey() string {
	if field := scope.PrimaryField(); field != nil {
		return field.DBName
	}
	return ""
}

// PrimaryKeyZero check the primary key is blank or not
func (scope *Scope) PrimaryKeyZero() bool {
	field := scope.PrimaryField()
	return field == nil || field.IsBlank
}

// PrimaryKeyValue get the primary key's value
func (scope *Scope) PrimaryKeyValue() interface{} {
	if field := scope.PrimaryField(); field != nil && field.Field.IsValid() {
		return field.Field.Interface()
	}
	return 0
}

// HasColumn to check if has column
func (scope *Scope) HasColumn(column string) bool {
	for _, field := range scope.GetStructFields() {
		if field.IsNormal && (field.Name == column || field.DBName == column) {
			return true
		}
	}
	return false
}

// SetColumn to set the column's value
func (scope *Scope) SetColumn(column interface{}, value interface{}) error {
	if field, ok := column.(*Field); ok {
		return field.Set(value)
	} else if name, ok := column.(string); ok {

		if field, ok := scope.Fields()[name]; ok {
			return field.Set(value)
		}

		dbName := ToDBName(name)
		if field, ok := scope.Fields()[dbName]; ok {
			return field.Set(value)
		}

		if field, ok := scope.FieldByName(name); ok {
			return field.Set(value)
		}
	}
	return errors.New("could not convert column to field")
}

func (scope *Scope) CallMethod(name string, checkError bool) {
	if scope.Value == nil || (checkError && scope.HasError()) {
		return
	}

	call := func(value interface{}) {
		if fm := reflect.ValueOf(value).MethodByName(name); fm.IsValid() {
			switch f := fm.Interface().(type) {
			case func():
				f()
			case func(s *Scope):
				f(scope)
			case func(s *DB):
				newDB := scope.NewDB()
				f(newDB)
				scope.Err(newDB.Error)
			case func() error:
				scope.Err(f())
			case func(s *Scope) error:
				scope.Err(f(scope))
			case func(s *DB) error:
				newDB := scope.NewDB()
				scope.Err(f(newDB))
				scope.Err(newDB.Error)
			default:
				scope.Err(fmt.Errorf("unsupported function %v", name))
			}
		}
	}

	if values := scope.IndirectValue(); values.Kind() == reflect.Slice {
		for i := 0; i < values.Len(); i++ {
			value := values.Index(i).Addr().Interface()
			if values.Index(i).Kind() == reflect.Ptr {
				value = values.Index(i).Interface()
			}
			call(value)
		}
	} else {
		if scope.IndirectValue().CanAddr() {
			call(scope.IndirectValue().Addr().Interface())
		} else {
			call(scope.IndirectValue().Interface())
		}
	}
}

func (scope *Scope) CallMethodWithErrorCheck(name string) {
	scope.CallMethod(name, true)
}

// AddToVars add value as sql's vars, gorm will escape them
func (scope *Scope) AddToVars(value interface{}) string {
	if expr, ok := value.(*expr); ok {
		exp := expr.expr
		for _, arg := range expr.args {
			exp = strings.Replace(exp, "?", scope.AddToVars(arg), 1)
		}
		return exp
	} else {
		scope.SqlVars = append(scope.SqlVars, value)
		return scope.Dialect().BinVar(len(scope.SqlVars))
	}
}

type tabler interface {
	TableName() string
}

type dbTabler interface {
	TableName(*DB) string
}

// TableName get table name
func (scope *Scope) TableName() string {
	if scope.Search != nil && len(scope.Search.tableName) > 0 {
		return scope.Search.tableName
	}

	if tabler, ok := scope.Value.(tabler); ok {
		return tabler.TableName()
	}

	if tabler, ok := scope.Value.(dbTabler); ok {
		return tabler.TableName(scope.db)
	}

	return scope.GetModelStruct().TableName(scope.db.Model(scope.Value))
}

func (scope *Scope) QuotedTableName() (name string) {
	if scope.Search != nil && len(scope.Search.tableName) > 0 {
		if strings.Index(scope.Search.tableName, " ") != -1 {
			return scope.Search.tableName
		}
		return scope.Quote(scope.Search.tableName)
	} else {
		return scope.Quote(scope.TableName())
	}
}

// CombinedConditionSql get combined condition sql
func (scope *Scope) CombinedConditionSql() string {
	return scope.joinsSql() + scope.whereSql() + scope.groupSql() +
		scope.havingSql() + scope.orderSql() + scope.limitSql() + scope.offsetSql()
}

func (scope *Scope) FieldByName(name string) (field *Field, ok bool) {
	for _, field := range scope.Fields() {
		if field.Name == name || field.DBName == name {
			return field, true
		}
	}
	return nil, false
}

// Raw set sql
func (scope *Scope) Raw(sql string) *Scope {
	scope.Sql = strings.Replace(sql, "$$", "?", -1)
	return scope
}

// Exec invoke sql
func (scope *Scope) Exec() *Scope {
	defer scope.Trace(NowFunc())

	if !scope.HasError() {
		if result, err := scope.SqlDB().Exec(scope.Sql, scope.SqlVars...); scope.Err(err) == nil {
			if count, err := result.RowsAffected(); scope.Err(err) == nil {
				scope.db.RowsAffected = count
			}
		}
	}
	return scope
}

// Set set value by name
func (scope *Scope) Set(name string, value interface{}) *Scope {
	scope.db.InstantSet(name, value)
	return scope
}

// Get get value by name
func (scope *Scope) Get(name string) (interface{}, bool) {
	return scope.db.Get(name)
}

// InstanceId get InstanceId for scope
func (scope *Scope) InstanceId() string {
	if scope.instanceId == "" {
		scope.instanceId = fmt.Sprintf("%v%v", &scope, &scope.db)
	}
	return scope.instanceId
}

func (scope *Scope) InstanceSet(name string, value interface{}) *Scope {
	return scope.Set(name+scope.InstanceId(), value)
}

func (scope *Scope) InstanceGet(name string) (interface{}, bool) {
	return scope.Get(name + scope.InstanceId())
}

// Trace print sql log
func (scope *Scope) Trace(t time.Time) {
	if len(scope.Sql) > 0 {
		scope.db.slog(scope.Sql, t, scope.SqlVars...)
	}
}

// Begin start a transaction
func (scope *Scope) Begin() *Scope {
	if db, ok := scope.SqlDB().(sqlDb); ok {
		if tx, err := db.Begin(); err == nil {
			scope.db.db = interface{}(tx).(sqlCommon)
			scope.InstanceSet("gorm:started_transaction", true)
		}
	}
	return scope
}

// CommitOrRollback commit current transaction if there is no error, otherwise rollback it
func (scope *Scope) CommitOrRollback() *Scope {
	if _, ok := scope.InstanceGet("gorm:started_transaction"); ok {
		if db, ok := scope.db.db.(sqlTx); ok {
			if scope.HasError() {
				db.Rollback()
			} else {
				db.Commit()
			}
			scope.db.db = scope.db.parent.db
		}
	}
	return scope
}

func (scope *Scope) SelectAttrs() []string {
	if scope.selectAttrs == nil {
		attrs := []string{}
		for _, value := range scope.Search.selects {
			if str, ok := value.(string); ok {
				attrs = append(attrs, str)
			} else if strs, ok := value.([]string); ok {
				attrs = append(attrs, strs...)
			} else if strs, ok := value.([]interface{}); ok {
				for _, str := range strs {
					attrs = append(attrs, fmt.Sprintf("%v", str))
				}
			}
		}
		scope.selectAttrs = &attrs
	}
	return *scope.selectAttrs
}

func (scope *Scope) OmitAttrs() []string {
	return scope.Search.omits
}

func (scope *Scope) changeableDBColumn(column string) bool {
	selectAttrs := scope.SelectAttrs()
	omitAttrs := scope.OmitAttrs()

	if len(selectAttrs) > 0 {
		for _, attr := range selectAttrs {
			if column == ToDBName(attr) {
				return true
			}
		}
		return false
	}

	for _, attr := range omitAttrs {
		if column == ToDBName(attr) {
			return false
		}
	}
	return true
}

func (scope *Scope) changeableField(field *Field) bool {
	selectAttrs := scope.SelectAttrs()
	omitAttrs := scope.OmitAttrs()

	if len(selectAttrs) > 0 {
		for _, attr := range selectAttrs {
			if field.Name == attr || field.DBName == attr {
				return true
			}
		}
		return false
	}

	for _, attr := range omitAttrs {
		if field.Name == attr || field.DBName == attr {
			return false
		}
	}

	return !field.IsIgnored
}

func (scope *Scope) shouldSaveAssociations() bool {
	saveAssociations, ok := scope.Get("gorm:save_associations")
	if ok && !saveAssociations.(bool) {
		return false
	}
	return true && !scope.HasError()
}
<commit_msg>scope: Catch error for db.Commit()<commit_after>package gorm

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
	"time"

	"reflect"
)

type Scope struct {
	Search          *search
	Value           interface{}
	Sql             string
	SqlVars         []interface{}
	db              *DB
	indirectValue   *reflect.Value
	instanceId      string
	primaryKeyField *Field
	skipLeft        bool
	fields          map[string]*Field
	selectAttrs     *[]string
}

func (scope *Scope) IndirectValue() reflect.Value {
	if scope.indirectValue == nil {
		value := reflect.Indirect(reflect.ValueOf(scope.Value))
		if value.Kind() == reflect.Ptr {
			value = value.Elem()
		}
		scope.indirectValue = &value
	}
	return *scope.indirectValue
}

func (scope *Scope) NeedPtr() *Scope {
	reflectKind := reflect.ValueOf(scope.Value).Kind()
	if !((reflectKind == reflect.Invalid) || (reflectKind == reflect.Ptr)) {
		err := fmt.Errorf("%v %v\n", fileWithLineNum(), "using unaddressable value")
		scope.Err(err)
		fmt.Printf(err.Error())
	}
	return scope
}

// New create a new Scope without search information
func (scope *Scope) New(value interface{}) *Scope {
	return &Scope{db: scope.NewDB(), Search: &search{}, Value: value}
}

// NewDB create a new DB without search information
func (scope *Scope) NewDB() *DB {
	if scope.db != nil {
		db := scope.db.clone()
		db.search = nil
		db.Value = nil
		return db
	}
	return nil
}

func (scope *Scope) DB() *DB {
	return scope.db
}

// SqlDB return *sql.DB
func (scope *Scope) SqlDB() sqlCommon {
	return scope.db.db
}

// SkipLeft skip remaining callbacks
func (scope *Scope) SkipLeft() {
	scope.skipLeft = true
}

// Quote used to quote database column name according to database dialect
func (scope *Scope) Quote(str string) string {
	if strings.Index(str, ".") != -1 {
		newStrs := []string{}
		for _, str := range strings.Split(str, ".") {
			newStrs = append(newStrs, scope.Dialect().Quote(str))
		}
		return strings.Join(newStrs, ".")
	} else {
		return scope.Dialect().Quote(str)
	}
}

func (scope *Scope) QuoteIfPossible(str string) string {
	if regexp.MustCompile("^[a-zA-Z]+(.[a-zA-Z]+)*$").MatchString(str) {
		return scope.Quote(str)
	}
	return str
}

// Dialect get dialect
func (scope *Scope) Dialect() Dialect {
	return scope.db.parent.dialect
}

// Err write error
func (scope *Scope) Err(err error) error {
	if err != nil {
		scope.db.AddError(err)
	}
	return err
}

// Log print log message
func (scope *Scope) Log(v ...interface{}) {
	scope.db.log(v...)
}

// HasError check if there are any error
func (scope *Scope) HasError() bool {
	return scope.db.Error != nil
}

func (scope *Scope) PrimaryFields() []*Field {
	var fields = []*Field{}
	for _, field := range scope.GetModelStruct().PrimaryFields {
		fields = append(fields, scope.Fields()[field.DBName])
	}
	return fields
}

func (scope *Scope) PrimaryField() *Field {
	if primaryFields := scope.GetModelStruct().PrimaryFields; len(primaryFields) > 0 {
		if len(primaryFields) > 1 {
			if field, ok := scope.Fields()["id"]; ok {
				return field
			}
		}
		return scope.Fields()[primaryFields[0].DBName]
	}
	return nil
}

// PrimaryKey get the primary key's column name
func (scope *Scope) PrimaryKey() string {
	if field := scope.PrimaryField(); field != nil {
		return field.DBName
	}
	return ""
}

// PrimaryKeyZero check the primary key is blank or not
func (scope *Scope) PrimaryKeyZero() bool {
	field := scope.PrimaryField()
	return field == nil || field.IsBlank
}

// PrimaryKeyValue get the primary key's value
func (scope *Scope) PrimaryKeyValue() interface{} {
	if field := scope.PrimaryField(); field != nil && field.Field.IsValid() {
		return field.Field.Interface()
	}
	return 0
}

// HasColumn to check if has column
func (scope *Scope) HasColumn(column string) bool {
	for _, field := range scope.GetStructFields() {
		if field.IsNormal && (field.Name == column || field.DBName == column) {
			return true
		}
	}
	return false
}

// SetColumn to set the column's value
func (scope *Scope) SetColumn(column interface{}, value interface{}) error {
	if field, ok := column.(*Field); ok {
		return field.Set(value)
	} else if name, ok := column.(string); ok {

		if field, ok := scope.Fields()[name]; ok {
			return field.Set(value)
		}

		dbName := ToDBName(name)
		if field, ok := scope.Fields()[dbName]; ok {
			return field.Set(value)
		}

		if field, ok := scope.FieldByName(name); ok {
			return field.Set(value)
		}
	}
	return errors.New("could not convert column to field")
}

func (scope *Scope) CallMethod(name string, checkError bool) {
	if scope.Value == nil || (checkError && scope.HasError()) {
		return
	}

	call := func(value interface{}) {
		if fm := reflect.ValueOf(value).MethodByName(name); fm.IsValid() {
			switch f := fm.Interface().(type) {
			case func():
				f()
			case func(s *Scope):
				f(scope)
			case func(s *DB):
				newDB := scope.NewDB()
				f(newDB)
				scope.Err(newDB.Error)
			case func() error:
				scope.Err(f())
			case func(s *Scope) error:
				scope.Err(f(scope))
			case func(s *DB) error:
				newDB := scope.NewDB()
				scope.Err(f(newDB))
				scope.Err(newDB.Error)
			default:
				scope.Err(fmt.Errorf("unsupported function %v", name))
			}
		}
	}

	if values := scope.IndirectValue(); values.Kind() == reflect.Slice {
		for i := 0; i < values.Len(); i++ {
			value := values.Index(i).Addr().Interface()
			if values.Index(i).Kind() == reflect.Ptr {
				value = values.Index(i).Interface()
			}
			call(value)
		}
	} else {
		if scope.IndirectValue().CanAddr() {
			call(scope.IndirectValue().Addr().Interface())
		} else {
			call(scope.IndirectValue().Interface())
		}
	}
}

func (scope *Scope) CallMethodWithErrorCheck(name string) {
	scope.CallMethod(name, true)
}

// AddToVars add value as sql's vars, gorm will escape them
func (scope *Scope) AddToVars(value interface{}) string {
	if expr, ok := value.(*expr); ok {
		exp := expr.expr
		for _, arg := range expr.args {
			exp = strings.Replace(exp, "?", scope.AddToVars(arg), 1)
		}
		return exp
	} else {
		scope.SqlVars = append(scope.SqlVars, value)
		return scope.Dialect().BinVar(len(scope.SqlVars))
	}
}

type tabler interface {
	TableName() string
}

type dbTabler interface {
	TableName(*DB) string
}

// TableName get table name
func (scope *Scope) TableName() string {
	if scope.Search != nil && len(scope.Search.tableName) > 0 {
		return scope.Search.tableName
	}

	if tabler, ok := scope.Value.(tabler); ok {
		return tabler.TableName()
	}

	if tabler, ok := scope.Value.(dbTabler); ok {
		return tabler.TableName(scope.db)
	}

	return scope.GetModelStruct().TableName(scope.db.Model(scope.Value))
}

func (scope *Scope) QuotedTableName() (name string) {
	if scope.Search != nil && len(scope.Search.tableName) > 0 {
		if strings.Index(scope.Search.tableName, " ") != -1 {
			return scope.Search.tableName
		}
		return scope.Quote(scope.Search.tableName)
	} else {
		return scope.Quote(scope.TableName())
	}
}

// CombinedConditionSql get combined condition sql
func (scope *Scope) CombinedConditionSql() string {
	return scope.joinsSql() + scope.whereSql() + scope.groupSql() +
		scope.havingSql() + scope.orderSql() + scope.limitSql() + scope.offsetSql()
}

func (scope *Scope) FieldByName(name string) (field *Field, ok bool) {
	for _, field := range scope.Fields() {
		if field.Name == name || field.DBName == name {
			return field, true
		}
	}
	return nil, false
}

// Raw set sql
func (scope *Scope) Raw(sql string) *Scope {
	scope.Sql = strings.Replace(sql, "$$", "?", -1)
	return scope
}

// Exec invoke sql
func (scope *Scope) Exec() *Scope {
	defer scope.Trace(NowFunc())

	if !scope.HasError() {
		if result, err := scope.SqlDB().Exec(scope.Sql, scope.SqlVars...); scope.Err(err) == nil {
			if count, err := result.RowsAffected(); scope.Err(err) == nil {
				scope.db.RowsAffected = count
			}
		}
	}
	return scope
}

// Set set value by name
func (scope *Scope) Set(name string, value interface{}) *Scope {
	scope.db.InstantSet(name, value)
	return scope
}

// Get get value by name
func (scope *Scope) Get(name string) (interface{}, bool) {
	return scope.db.Get(name)
}

// InstanceId get InstanceId for scope
func (scope *Scope) InstanceId() string {
	if scope.instanceId == "" {
		scope.instanceId = fmt.Sprintf("%v%v", &scope, &scope.db)
	}
	return scope.instanceId
}

func (scope *Scope) InstanceSet(name string, value interface{}) *Scope {
	return scope.Set(name+scope.InstanceId(), value)
}

func (scope *Scope) InstanceGet(name string) (interface{}, bool) {
	return scope.Get(name + scope.InstanceId())
}

// Trace print sql log
func (scope *Scope) Trace(t time.Time) {
	if len(scope.Sql) > 0 {
		scope.db.slog(scope.Sql, t, scope.SqlVars...)
	}
}

// Begin start a transaction
func (scope *Scope) Begin() *Scope {
	if db, ok := scope.SqlDB().(sqlDb); ok {
		if tx, err := db.Begin(); err == nil {
			scope.db.db = interface{}(tx).(sqlCommon)
			scope.InstanceSet("gorm:started_transaction", true)
		}
	}
	return scope
}

// CommitOrRollback commit current transaction if there is no error, otherwise rollback it
func (scope *Scope) CommitOrRollback() *Scope {
	if _, ok := scope.InstanceGet("gorm:started_transaction"); ok {
		if db, ok := scope.db.db.(sqlTx); ok {
			if scope.HasError() {
				db.Rollback()
			} else {
				scope.Err(db.Commit())
			}
			scope.db.db = scope.db.parent.db
		}
	}
	return scope
}

func (scope *Scope) SelectAttrs() []string {
	if scope.selectAttrs == nil {
		attrs := []string{}
		for _, value := range scope.Search.selects {
			if str, ok := value.(string); ok {
				attrs = append(attrs, str)
			} else if strs, ok := value.([]string); ok {
				attrs = append(attrs, strs...)
			} else if strs, ok := value.([]interface{}); ok {
				for _, str := range strs {
					attrs = append(attrs, fmt.Sprintf("%v", str))
				}
			}
		}
		scope.selectAttrs = &attrs
	}
	return *scope.selectAttrs
}

func (scope *Scope) OmitAttrs() []string {
	return scope.Search.omits
}

func (scope *Scope) changeableDBColumn(column string) bool {
	selectAttrs := scope.SelectAttrs()
	omitAttrs := scope.OmitAttrs()

	if len(selectAttrs) > 0 {
		for _, attr := range selectAttrs {
			if column == ToDBName(attr) {
				return true
			}
		}
		return false
	}

	for _, attr := range omitAttrs {
		if column == ToDBName(attr) {
			return false
		}
	}
	return true
}

func (scope *Scope) changeableField(field *Field) bool {
	selectAttrs := scope.SelectAttrs()
	omitAttrs := scope.OmitAttrs()

	if len(selectAttrs) > 0 {
		for _, attr := range selectAttrs {
			if field.Name == attr || field.DBName == attr {
				return true
			}
		}
		return false
	}

	for _, attr := range omitAttrs {
		if field.Name == attr || field.DBName == attr {
			return false
		}
	}

	return !field.IsIgnored
}

func (scope *Scope) shouldSaveAssociations() bool {
	saveAssociations, ok := scope.Get("gorm:save_associations")
	if ok && !saveAssociations.(bool) {
		return false
	}
	return true && !scope.HasError()
}
<|endoftext|>


<commit_before>/*
Copyright 2015 Oscar Ruckdeschel, Janik Schmidt, Jonathan Kuhse.

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 data

import "testing"

func TestCreateDatabase(t *testing.T) {
	createDatabase("root", "")
}

func TestTitle(t *testing.T) {
	cases := []string{
		"Hello world",
		"This is a News Title",
	}
	for _, c := range cases {
		newsID := postNews(c)
		got := getNewsTitle(newsID)
		if got != c {
			t.Errorf("save and read title: \n saved:  %q,\n got:     %q", c, got)
		}
	}
}

func TestTrigrams(t *testing.T) {
	cases := []struct {
		id       int
		trigrams []string
	}{
		{1, []string{"Hel", "ell", "llo"}},
		{2, []string{"Tes", "est"}},
	}
	for _, c := range cases {
		for _, tri := range c.trigrams {
			putTrigram(tri, c.id)
		}

		for _, tri := range c.trigrams {
			ids := getIdsOfTrigram(tri)
			contains := false
			for _, id := range ids {
				if id == c.id {
					contains = true
				}
			}
			if !contains {
				t.Errorf("%q with id %q is not in %q", tri, string(c.id), ids)
			}
		}

	}
}
<commit_msg>test cases with unicode chars<commit_after>/*
Copyright 2015 Oscar Ruckdeschel, Janik Schmidt, Jonathan Kuhse.

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 data

import (
	"testing"
)

func TestCreateDatabase(t *testing.T) {
	createDatabase("root", "")
}

func TestTitle(t *testing.T) {
	cases := []string{
		"Hello world",
		"This is a News Title",
		"",
		"»",
	}
	for _, c := range cases {
		newsID := postNews(c)
		got := getNewsTitle(newsID)
		if got != c {
			t.Errorf("save and read title: \n saved:  %q,\n got:     %q", c, got)
		}
	}
}

func TestTrigrams(t *testing.T) {
	cases := []struct {
		id       int
		trigrams []string
	}{
		{1, []string{"Hel", "ell", "llo"}},
		{2, []string{"Tes", "est"}},
		{1, []string{"Wie", "ied", "ede", "der"}}, //duplicate id -> no problem (?)
		// {1, []string{"Hel", "ell", "llo"}}, TODO: better error handling (here: duplicate entry)
	}
	for _, c := range cases {
		for _, tri := range c.trigrams {
			putTrigram(tri, c.id)
		}

		for _, tri := range c.trigrams {
			ids := getIdsOfTrigram(tri)
			contains := false
			for _, id := range ids {
				if id == c.id {
					contains = true
				}
			}
			if !contains {
				t.Errorf("%q with id %q is not in %q", tri, string(c.id), ids)
			}
		}

	}
}
<|endoftext|>


<commit_before>// Copyright 2016 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package jssdk

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

	"github.com/issue9/wechat/mp/common/config"
	"github.com/issue9/wechat/mp/common/result"
)

// 授权作用域，供 GetCodeURL 使用。
const (
	SnsapiUserinfo = "snsapi_uesrnfo"
	SnsapiBase     = "snsapi_base"

	// 获取 code 的地址
	codeURL = "https://open.weixin.qq.com/connect/oauth2/authorize?appid=%v&redirect_uri=%vresponse_type=code&scope=%v&state=%v#wechat_redirect"
)

// GetCodeURL 获取 code
func GetCodeURL(conf *config.Config, redirectURI, scope, state string) string {
	return fmt.Sprintf(codeURL, conf.AppID, redirectURI, scope, state)
}

// GetAccessToken 根据 code 获取 access_token
func GetAccessToken(conf *config.Config, code string) (*AccessToken, error) {
	queries := map[string]string{
		"appid":      conf.AppID,
		"secret":     conf.AppSecret,
		"code":       code,
		"grant_type": "authorization_code",
	}
	url := conf.URL("sns/oauth2/access_token", queries)

	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	return parseAccessToken(resp.Body)
}

// RefreshAccessToken 刷新 access_token
func RefreshAccessToken(conf *config.Config, token *AccessToken) (*AccessToken, error) {
	queries := map[string]string{
		"appid":         conf.AppID,
		"refresh_token": token.RefreshToken,
		"grant_type":    "refresh_token",
	}
	url := conf.URL("sns/oauth2/refresh_token", queries)

	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	return parseAccessToken(resp.Body)
}

// GetUserInfo 获取用户基本信息
//
// 若不指定 lang 则使用 zh_CN 作为其默认值。
func GetUserInfo(conf *config.Config, token *AccessToken, lang string) (*UserInfo, error) {
	if len(lang) == 0 {
		lang = "zh_CN"
	}
	queries := map[string]string{
		"openid":       token.OpenID,
		"access_token": token.AccessToken,
		"lang":         lang,
	}
	url := conf.URL("sns/userinfo", queries)

	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	info := &UserInfo{}
	if err = json.Unmarshal(data, info); err != nil {
		return nil, err
	}
	if len(info.OpenID) > 0 {
		return info, nil
	}

	rslt := &result.Result{}
	if err = json.Unmarshal(data, rslt); err != nil {
		return nil, err
	}
	return nil, rslt
}

// AuthAccessToken 验证 access_token 是否有效
func AuthAccessToken(conf *config.Config, token *AccessToken) (bool, error) {
	queries := map[string]string{
		"openid":       token.OpenID,
		"access_token": token.AccessToken,
	}
	url := conf.URL("sns/auth", queries)

	resp, err := http.Get(url)
	if err != nil {
		return false, err
	}
	defer resp.Body.Close()

	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return false, err
	}

	rslt := result.From(data)
	return rslt.Code == 0, rslt
}

// 分析 r 中的数据到 AccessToken 或是 result.Result 对象中。
func parseAccessToken(r io.Reader) (*AccessToken, error) {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, err
	}

	token := &AccessToken{}
	if err := json.Unmarshal(data, token); err != nil {
		return nil, err
	}
	if len(token.AccessToken) > 0 || token.ExpiresIn > 0 {
		token.Created = time.Now()
		return token, nil
	}

	rslt := &result.Result{}
	if err := json.Unmarshal(data, rslt); err != nil {
		return nil, err
	}
	return nil, rslt
}
<commit_msg>修正未对 redirectURI 进行编码的错误<commit_after>// Copyright 2016 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package jssdk

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

	"github.com/issue9/wechat/mp/common/config"
	"github.com/issue9/wechat/mp/common/result"
)

// 授权作用域，供 GetCodeURL 使用。
const (
	SnsapiUserinfo = "snsapi_uesrnfo"
	SnsapiBase     = "snsapi_base"

	// 获取 code 的地址
	codeURL = "https://open.weixin.qq.com/connect/oauth2/authorize?appid=%v&redirect_uri=%vresponse_type=code&scope=%v&state=%v#wechat_redirect"
)

// GetCodeURL 获取 code
func GetCodeURL(conf *config.Config, redirectURI, scope, state string) string {
	redirectURI = url.QueryEscape(redirectURI)
	return fmt.Sprintf(codeURL, conf.AppID, redirectURI, scope, state)
}

// GetAccessToken 根据 code 获取 access_token
func GetAccessToken(conf *config.Config, code string) (*AccessToken, error) {
	queries := map[string]string{
		"appid":      conf.AppID,
		"secret":     conf.AppSecret,
		"code":       code,
		"grant_type": "authorization_code",
	}
	url := conf.URL("sns/oauth2/access_token", queries)

	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	return parseAccessToken(resp.Body)
}

// RefreshAccessToken 刷新 access_token
func RefreshAccessToken(conf *config.Config, token *AccessToken) (*AccessToken, error) {
	queries := map[string]string{
		"appid":         conf.AppID,
		"refresh_token": token.RefreshToken,
		"grant_type":    "refresh_token",
	}
	url := conf.URL("sns/oauth2/refresh_token", queries)

	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	return parseAccessToken(resp.Body)
}

// GetUserInfo 获取用户基本信息
//
// 若不指定 lang 则使用 zh_CN 作为其默认值。
func GetUserInfo(conf *config.Config, token *AccessToken, lang string) (*UserInfo, error) {
	if len(lang) == 0 {
		lang = "zh_CN"
	}
	queries := map[string]string{
		"openid":       token.OpenID,
		"access_token": token.AccessToken,
		"lang":         lang,
	}
	url := conf.URL("sns/userinfo", queries)

	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	info := &UserInfo{}
	if err = json.Unmarshal(data, info); err != nil {
		return nil, err
	}
	if len(info.OpenID) > 0 {
		return info, nil
	}

	rslt := &result.Result{}
	if err = json.Unmarshal(data, rslt); err != nil {
		return nil, err
	}
	return nil, rslt
}

// AuthAccessToken 验证 access_token 是否有效
func AuthAccessToken(conf *config.Config, token *AccessToken) (bool, error) {
	queries := map[string]string{
		"openid":       token.OpenID,
		"access_token": token.AccessToken,
	}
	url := conf.URL("sns/auth", queries)

	resp, err := http.Get(url)
	if err != nil {
		return false, err
	}
	defer resp.Body.Close()

	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return false, err
	}

	rslt := result.From(data)
	return rslt.Code == 0, rslt
}

// 分析 r 中的数据到 AccessToken 或是 result.Result 对象中。
func parseAccessToken(r io.Reader) (*AccessToken, error) {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, err
	}

	token := &AccessToken{}
	if err := json.Unmarshal(data, token); err != nil {
		return nil, err
	}
	if len(token.AccessToken) > 0 || token.ExpiresIn > 0 {
		token.Created = time.Now()
		return token, nil
	}

	rslt := &result.Result{}
	if err := json.Unmarshal(data, rslt); err != nil {
		return nil, err
	}
	return nil, rslt
}
<|endoftext|>


<commit_before>package pipelines_test

import (
	"os/exec"

	"github.com/concourse/testflight/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
)

var _ = Describe("Configuring a resouce with a tag", func() {
	BeforeEach(func() {
		if !hasTaggedWorkers() {
			Skip("this only runs when a worker with the 'tagged' tag is available")
		}
	})

	It("puts the resource check container on the tagged worker", func() {
		configurePipeline(
			"-c", "fixtures/tagged_resource.yml",
		)
		resourceString := pipelineName + "/" + "some-resource"

		fly := exec.Command(flyBin, "-t", targetedConcourse, "check-resource", "-r", resourceString)
		session := helpers.StartFly(fly)
		Eventually(session).Should(gexec.Exit(0))

		workersTable := flyTable("workers")
		taggedWorkerHandles := []string{}
		for _, w := range workersTable {
			if w["tags"] == "tagged" {
				taggedWorkerHandles = append(taggedWorkerHandles, w["name"])
			}
		}

		containerTable := flyTable("containers")
		Expect(containerTable).To(HaveLen(1))
		Expect(containerTable[0]["type"]).To(Equal("check"))
		Expect(taggedWorkerHandles).To(ContainElement(containerTable[0]["worker"]))
	})
})
<commit_msg>get container for current pipeline<commit_after>package pipelines_test

import (
	"os/exec"

	"github.com/concourse/testflight/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
)

var _ = Describe("Configuring a resouce with a tag", func() {
	BeforeEach(func() {
		if !hasTaggedWorkers() {
			Skip("this only runs when a worker with the 'tagged' tag is available")
		}
	})

	It("puts the resource check container on the tagged worker", func() {
		configurePipeline(
			"-c", "fixtures/tagged_resource.yml",
		)
		resourceString := pipelineName + "/" + "some-resource"

		fly := exec.Command(flyBin, "-t", targetedConcourse, "check-resource", "-r", resourceString)
		session := helpers.StartFly(fly)
		Eventually(session).Should(gexec.Exit(0))

		workersTable := flyTable("workers")
		taggedWorkerHandles := []string{}
		for _, w := range workersTable {
			if w["tags"] == "tagged" {
				taggedWorkerHandles = append(taggedWorkerHandles, w["name"])
			}
		}

		containerTable := flyTable("containers")
		currentPipelineContainers := []map[string]string{}
		for _, c := range containerTable {
			if c["pipeline"] == pipelineName {
				currentPipelineContainers = append(currentPipelineContainers, c)
			}
		}
		Expect(currentPipelineContainers).To(HaveLen(1))
		Expect(currentPipelineContainers[0]["type"]).To(Equal("check"))
		Expect(taggedWorkerHandles).To(ContainElement(currentPipelineContainers[0]["worker"]))
	})
})
<|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 associated

// The iOS format is specified by:
//   https://developer.apple.com/documentation/safariservices/supporting_associated_domains

import (
	"fmt"

	"github.com/google/exposure-notifications-verification-server/pkg/database"
)

type IOSData struct {
	Applinks Applinks `json:"applinks"`

	// The following two fields are included for completeness' sake, but are not
	// currently populated/used by the system.
	Webcredentials *Appstrings `json:"webcredentials,omitempty"`
	Appclips       *Appstrings `json:"appclips,omitempty"`
}

type Applinks struct {
	Apps    []string `json:"apps"`
	Details []Detail `json:"details,omitempty"`
}

type Detail struct {
	AppID string   `json:"appID,omitempty"`
	Paths []string `json:"components,omitempty"`
}

type Appstrings struct {
	Apps []string `json:"apps,omitempty"`
}

// getAppIds finds all the iOS app ids we know about.
func (c *Controller) getAppIds(realmID uint) ([]string, error) {
	apps, err := c.db.ListActiveApps(realmID, database.WithAppOS(database.OSTypeIOS))
	if err != nil {
		return nil, err
	}
	ret := make([]string, 0, len(apps))
	for i := range apps {
		ret = append(ret, apps[i].AppID)
	}
	return ret, nil
}

// getIosData gets the iOS app data.
func (c *Controller) getIosData(region string) (*IOSData, error) {
	realm, err := c.db.FindRealmByRegion(region)
	if err != nil {
		return nil, fmt.Errorf("unable to lookup realm: %w", err)
	}

	ids, err := c.getAppIds(realm.ID)
	if err != nil {
		return nil, err
	}

	if len(ids) == 0 {
		return nil, nil
	}

	details := make([]Detail, len(ids))
	for i, id := range ids {
		details[i] = Detail{
			AppID: id,
			Paths: []string{"*"},
		}
	}

	return &IOSData{
		Applinks: Applinks{
			Apps:    []string{}, // expected always empty.
			Details: details,
		},
	}, nil
}
<commit_msg>fix json output (#1374)<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 associated

// The iOS format is specified by:
//   https://developer.apple.com/documentation/safariservices/supporting_associated_domains

import (
	"fmt"

	"github.com/google/exposure-notifications-verification-server/pkg/database"
)

type IOSData struct {
	Applinks Applinks `json:"applinks"`

	// The following two fields are included for completeness' sake, but are not
	// currently populated/used by the system.
	Webcredentials *Appstrings `json:"webcredentials,omitempty"`
	Appclips       *Appstrings `json:"appclips,omitempty"`
}

type Applinks struct {
	Apps    []string `json:"apps"`
	Details []Detail `json:"details,omitempty"`
}

type Detail struct {
	AppID string   `json:"appID,omitempty"`
	Paths []string `json:"paths,omitempty"`
}

type Appstrings struct {
	Apps []string `json:"apps,omitempty"`
}

// getAppIds finds all the iOS app ids we know about.
func (c *Controller) getAppIds(realmID uint) ([]string, error) {
	apps, err := c.db.ListActiveApps(realmID, database.WithAppOS(database.OSTypeIOS))
	if err != nil {
		return nil, err
	}
	ret := make([]string, 0, len(apps))
	for i := range apps {
		ret = append(ret, apps[i].AppID)
	}
	return ret, nil
}

// getIosData gets the iOS app data.
func (c *Controller) getIosData(region string) (*IOSData, error) {
	realm, err := c.db.FindRealmByRegion(region)
	if err != nil {
		return nil, fmt.Errorf("unable to lookup realm: %w", err)
	}

	ids, err := c.getAppIds(realm.ID)
	if err != nil {
		return nil, err
	}

	if len(ids) == 0 {
		return nil, nil
	}

	details := make([]Detail, len(ids))
	for i, id := range ids {
		details[i] = Detail{
			AppID: id,
			Paths: []string{"*"},
		}
	}

	return &IOSData{
		Applinks: Applinks{
			Apps:    []string{}, // expected always empty.
			Details: details,
		},
	}, nil
}
<|endoftext|>


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

package bytematcher

import (
	"fmt"

	"github.com/richardlehane/match/wac"
	"github.com/richardlehane/siegfried/config"
	"github.com/richardlehane/siegfried/pkg/core"
	"github.com/richardlehane/siegfried/pkg/core/siegreader"
)

func (b *Matcher) start(bof bool) {
	b.mu.Lock()
	defer b.mu.Unlock()
	if bof {
		if b.bAho != nil {
			return
		}
		if b.lowmem {
			b.bAho = wac.NewLowMem(b.bofSeq.set)
			return
		}
		b.bAho = wac.New(b.bofSeq.set)
		return
	}
	if b.eAho != nil {
		return
	}
	if b.lowmem {
		b.eAho = wac.NewLowMem(b.eofSeq.set)
		return
	}
	b.eAho = wac.New(b.eofSeq.set)
}

// Identify function - brings a new matcher into existence
func (b *Matcher) identify(buf siegreader.Buffer, quit chan struct{}, r chan core.Result) {
	buf.SetQuit(quit)
	incoming := b.newScorer(buf, quit, r)

	// Test BOF/EOF sequences
	rdr := siegreader.LimitReaderFrom(buf, b.maxBOF)
	// start bof matcher if not yet started
	b.start(true)
	var bchan chan wac.Result
	if rdr != nil {
		bchan = b.bAho.Index(rdr)
		// Do an initial check of BOF sequences
		for br := range bchan {
			if br.Index[0] == -1 {
				incoming <- progressStrike(br.Offset, false)
				if br.Offset > 2048 {
					break
				}
			} else {
				if config.Debug() {
					fmt.Println(strike{b.bofSeq.testTreeIndex[br.Index[0]], br.Index[1], br.Offset, br.Length, false, false, br.Final})
				}
				incoming <- strike{b.bofSeq.testTreeIndex[br.Index[0]], br.Index[1], br.Offset, br.Length, false, false, br.Final}
			}
		}
		select {
		case <-quit:
			// the matcher has called quit
			for _ = range bchan {
			} // drain first
			close(incoming)
			return
		default:
		}
	}
	// Setup BOF/EOF frame tests
	bfchan := b.bofFrames.index(buf, false, quit)
	efchan := b.eofFrames.index(buf, true, quit)

	// Setup EOF sequences test
	b.start(false)
	rrdr := siegreader.LimitReverseReaderFrom(buf, b.maxEOF)
	echan := b.eAho.Index(rrdr)

	// Now enter main search loop
	for {
		select {
		case bf, ok := <-bfchan:
			if !ok {
				bfchan = nil
			} else {
				if config.Debug() {
					fmt.Println(strike{b.bofFrames.testTreeIndex[bf.idx], 0, bf.off, bf.length, false, true, true})
				}
				incoming <- strike{b.bofFrames.testTreeIndex[bf.idx], 0, bf.off, bf.length, false, true, true}
			}
		case ef, ok := <-efchan:
			if !ok {
				efchan = nil
			} else {
				if config.Debug() {
					fmt.Println(strike{b.eofFrames.testTreeIndex[ef.idx], 0, ef.off, ef.length, true, true, true})
				}
				incoming <- strike{b.eofFrames.testTreeIndex[ef.idx], 0, ef.off, ef.length, true, true, true}
			}
		case br, ok := <-bchan:
			if !ok {
				bchan = nil
			} else {
				if br.Index[0] == -1 {
					incoming <- progressStrike(br.Offset, false)
				} else {
					if config.Debug() {
						fmt.Println(strike{b.bofSeq.testTreeIndex[br.Index[0]], br.Index[1], br.Offset, br.Length, false, false, br.Final})
					}
					incoming <- strike{b.bofSeq.testTreeIndex[br.Index[0]], br.Index[1], br.Offset, br.Length, false, false, br.Final}
				}
			}
		case er, ok := <-echan:
			if !ok {
				echan = nil
			} else {
				if er.Index[0] == -1 {
					incoming <- progressStrike(er.Offset, true)
				} else {
					if config.Debug() {
						fmt.Println(strike{b.eofSeq.testTreeIndex[er.Index[0]], er.Index[1], er.Offset, er.Length, true, false, er.Final})
					}
					incoming <- strike{b.eofSeq.testTreeIndex[er.Index[0]], er.Index[1], er.Offset, er.Length, true, false, er.Final}
				}
			}
		}
		if bfchan == nil && efchan == nil && bchan == nil && echan == nil {
			close(incoming)
			return
		}
	}
}
<commit_msg>test frame sets first<commit_after>// Copyright 2014 Richard Lehane. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package bytematcher

import (
	"fmt"

	"github.com/richardlehane/match/wac"
	"github.com/richardlehane/siegfried/config"
	"github.com/richardlehane/siegfried/pkg/core"
	"github.com/richardlehane/siegfried/pkg/core/siegreader"
)

func (b *Matcher) start(bof bool) {
	b.mu.Lock()
	defer b.mu.Unlock()
	if bof {
		if b.bAho != nil {
			return
		}
		if b.lowmem {
			b.bAho = wac.NewLowMem(b.bofSeq.set)
			return
		}
		b.bAho = wac.New(b.bofSeq.set)
		return
	}
	if b.eAho != nil {
		return
	}
	if b.lowmem {
		b.eAho = wac.NewLowMem(b.eofSeq.set)
		return
	}
	b.eAho = wac.New(b.eofSeq.set)
}

// Identify function - brings a new matcher into existence
func (b *Matcher) identify(buf siegreader.Buffer, quit chan struct{}, r chan core.Result) {
	buf.SetQuit(quit)
	incoming := b.newScorer(buf, quit, r)
	rdr := siegreader.LimitReaderFrom(buf, b.maxBOF)

	// First test BOF frameset
	bfchan := b.bofFrames.index(buf, false, quit)
	for bf := range bfchan {
		if config.Debug() {
			fmt.Println(strike{b.bofFrames.testTreeIndex[bf.idx], 0, bf.off, bf.length, false, true, true})
		}
		incoming <- strike{b.bofFrames.testTreeIndex[bf.idx], 0, bf.off, bf.length, false, true, true}
	}
	select {
	case <-quit: // the matcher has called quit
		close(incoming)
		return
	default:
	}

	// Do an initial check of BOF sequences
	b.start(true) // start bof matcher if not yet started
	var bchan chan wac.Result
	bchan = b.bAho.Index(rdr)
	for br := range bchan {
		if br.Index[0] == -1 {
			incoming <- progressStrike(br.Offset, false)
			if br.Offset > 2048 {
				break
			}
		} else {
			if config.Debug() {
				fmt.Println(strike{b.bofSeq.testTreeIndex[br.Index[0]], br.Index[1], br.Offset, br.Length, false, false, br.Final})
			}
			incoming <- strike{b.bofSeq.testTreeIndex[br.Index[0]], br.Index[1], br.Offset, br.Length, false, false, br.Final}
		}
	}
	select {
	case <-quit: // the matcher has called quit
		for _ = range bchan {
		} // drain first
		close(incoming)
		return
	default:
	}

	// Check EOF frame tests
	efchan := b.eofFrames.index(buf, true, quit)
	for ef := range efchan {
		if config.Debug() {
			fmt.Println(strike{b.eofFrames.testTreeIndex[ef.idx], 0, ef.off, ef.length, true, true, true})
		}
		incoming <- strike{b.eofFrames.testTreeIndex[ef.idx], 0, ef.off, ef.length, true, true, true}
	}

	// Setup EOF sequences test
	b.start(false)
	rrdr := siegreader.LimitReverseReaderFrom(buf, b.maxEOF)
	echan := b.eAho.Index(rrdr)

	// Now enter main search loop
	for {
		select {
		case br, ok := <-bchan:
			if !ok {
				bchan = nil
			} else {
				if br.Index[0] == -1 {
					incoming <- progressStrike(br.Offset, false)
				} else {
					if config.Debug() {
						fmt.Println(strike{b.bofSeq.testTreeIndex[br.Index[0]], br.Index[1], br.Offset, br.Length, false, false, br.Final})
					}
					incoming <- strike{b.bofSeq.testTreeIndex[br.Index[0]], br.Index[1], br.Offset, br.Length, false, false, br.Final}
				}
			}
		case er, ok := <-echan:
			if !ok {
				echan = nil
			} else {
				if er.Index[0] == -1 {
					incoming <- progressStrike(er.Offset, true)
				} else {
					if config.Debug() {
						fmt.Println(strike{b.eofSeq.testTreeIndex[er.Index[0]], er.Index[1], er.Offset, er.Length, true, false, er.Final})
					}
					incoming <- strike{b.eofSeq.testTreeIndex[er.Index[0]], er.Index[1], er.Offset, er.Length, true, false, er.Final}
				}
			}
		}
		if bchan == nil && echan == nil {
			close(incoming)
			return
		}
	}
}
<|endoftext|>