<commit_before>package kontrol

import (
	"errors"
	"fmt"
	"net/http"
	"time"

	"github.com/koding/kite"
	"github.com/koding/kite/dnode"
	"github.com/koding/kite/kontrol/onceevery"
	kontrolprotocol "github.com/koding/kite/kontrol/protocol"
	"github.com/koding/kite/protocol"
)

func (k *Kontrol) handleHeartbeat(rw http.ResponseWriter, req *http.Request) {
	id := req.URL.Query().Get("id")

	k.heartbeatsMu.Lock()
	defer k.heartbeatsMu.Unlock()

	if updateTimer, ok := k.heartbeats[id]; ok {
		// try to reset the timer every time the remote kite sends sends us a
		// heartbeat. Because the timer get reset, the timer is never fired, so
		// the value get always updated with the updater in the background
		// according to the write interval. If the kite doesn't send any
		// heartbeat, the timer func is being called, which stops the updater
		// so the key is being deleted automatically via the TTL mechanism.
		updateTimer.Reset(HeartbeatInterval + HeartbeatDelay)
		rw.Write([]byte("pong"))
		return
	}

	// if we reach here than it has several meanings:
	// * kite was registered before, but kontrol is restarted
	// * kite was registered before, but kontrol has lost track
	// * kite was no registered and someone else sends an heartbeat
	// we send back "registeragain" so the caller can be added in to the
	// heartbeats map above.
	rw.Write([]byte("registeragain"))
}

func (k *Kontrol) handleRegisterHTTP(r *kite.Request) (interface{}, error) {
	k.log.Info("Register (via HTTP) request from: %s", r.Client.Kite)

	if r.Args.One().MustMap()["url"].MustString() == "" {
		return nil, errors.New("invalid url")
	}

	var args struct {
		URL string `json:"url"`
	}
	r.Args.One().MustUnmarshal(&args)
	if args.URL == "" {
		return nil, errors.New("empty url")
	}

	// Only accept requests with kiteKey because we need this info
	// for generating tokens for this kite.
	if r.Auth.Type != "kiteKey" {
		return nil, fmt.Errorf("Unexpected authentication type: %s", r.Auth.Type)
	}

	kiteURL := args.URL
	remote := r.Client

	if err := validateKiteKey(&remote.Kite); err != nil {
		return nil, err
	}

	value := &kontrolprotocol.RegisterValue{
		URL: kiteURL,
	}

	// Register first by adding the value to the storage. Return if there is
	// any error.
	if err := k.storage.Upsert(&remote.Kite, value); err != nil {
		k.log.Error("storage add '%s' error: %s", remote.Kite, err)
		return nil, errors.New("internal error - register")
	}

	// we create a new ticker which is going to update the key periodically in
	// the storage so it's always up to date. Instead of updating the key
	// periodically according to the HeartBeatInterval below, we are buffering
	// the write speed here with the UpdateInterval.
	stopped := make(chan struct{})
	updater := time.NewTicker(UpdateInterval)
	updaterFunc := func() {
		for {
			select {
			case <-updater.C:
				k.log.Info("Kite is active (via HTTP), updating the value %s", remote.Kite)
				err := k.storage.Update(&remote.Kite, value)
				if err != nil {
					k.log.Error("storage update '%s' error: %s", remote.Kite, err)
				}
			case <-stopped:
				k.log.Info("Kite is nonactive (via HTTP). Updater is closed %s", remote.Kite)
				return
			}
		}
	}
	go updaterFunc()

	// lostFunc is called when we don't get any heartbeat or we lost
	// connection. In any case it will stop the updater
	lostFunc := func() {
		k.log.Info("Kite didn't get heartbeat (via HTTP). Stopping the updater %s",
			remote.Kite)
		// stop the updater so it doesn't update it in the background
		updater.Stop()

		k.heartbeatsMu.Lock()

		select {
		case <-stopped:
		default:
			close(stopped)
		}

		delete(k.heartbeats, remote.Kite.ID)
		k.heartbeatsMu.Unlock()
	}

	// we are now creating a timer that is going to call the lostFunc, which
	// stops the background updater if it's not resetted. The time is being
	// resetted on a separate HTTP endpoint "/heartbeat"
	k.heartbeatsMu.Lock()
	k.heartbeats[remote.Kite.ID] = time.AfterFunc(HeartbeatInterval+HeartbeatDelay,
		lostFunc)
	k.heartbeatsMu.Unlock()

	k.log.Info("Kite registered (via HTTP): %s", remote.Kite)

	// send response back to the kite, also identify him with the new name
	return &protocol.RegisterResult{
		URL:               args.URL,
		HeartbeatInterval: int64(HeartbeatInterval / time.Second),
	}, nil
}

func (k *Kontrol) handleRegister(r *kite.Request) (interface{}, error) {
	k.log.Info("Register request from: %s", r.Client.Kite)

	if r.Args.One().MustMap()["url"].MustString() == "" {
		return nil, errors.New("invalid url")
	}

	var args struct {
		URL string `json:"url"`
	}
	r.Args.One().MustUnmarshal(&args)
	if args.URL == "" {
		return nil, errors.New("empty url")
	}

	// Only accept requests with kiteKey because we need this info
	// for generating tokens for this kite.
	if r.Auth.Type != "kiteKey" {
		return nil, fmt.Errorf("Unexpected authentication type: %s", r.Auth.Type)
	}

	kiteURL := args.URL
	remote := r.Client

	if err := validateKiteKey(&remote.Kite); err != nil {
		return nil, err
	}

	value := &kontrolprotocol.RegisterValue{
		URL: kiteURL,
	}

	// Register first by adding the value to the storage. Return if there is
	// any error.
	if err := k.storage.Upsert(&remote.Kite, value); err != nil {
		k.log.Error("storage add '%s' error: %s", remote.Kite, err)
		return nil, errors.New("internal error - register")
	}

	every := onceevery.New(UpdateInterval)

	ping := make(chan struct{}, 1)
	closed := false

	updaterFunc := func() {
		for {
			select {
			case <-ping:
				k.log.Debug("Kite is active, got a ping %s", remote.Kite)
				every.Do(func() {
					k.log.Info("Kite is active, updating the value %s", remote.Kite)
					err := k.storage.Update(&remote.Kite, value)
					if err != nil {
						k.log.Error("storage update '%s' error: %s", remote.Kite, err)
					}
				})
			case <-time.After(HeartbeatInterval + HeartbeatDelay):
				k.log.Info("Kite didn't sent any heartbeat %s.", remote.Kite)
				every.Stop()
				closed = true
				return
			}
		}
	}
	go updaterFunc()

	heartbeatArgs := []interface{}{
		HeartbeatInterval / time.Second,
		dnode.Callback(func(args *dnode.Partial) {
			k.log.Debug("Kite send us an heartbeat. %s", remote.Kite)

			k.clientLocks.Get(remote.Kite.ID).Lock()
			defer k.clientLocks.Get(remote.Kite.ID).Unlock()

			select {
			case ping <- struct{}{}:
			default:
			}

			// seems we miss a heartbeat, so start it again!
			if closed {
				closed = false
				k.log.Warning("Updater was closed, but we are still getting heartbeats. Starting again %s",
					remote.Kite)

				// it might be removed because the ttl cleaner would come
				// before us, so try to add it again, the updater will than
				// continue to update it afterwards.
				k.storage.Upsert(&remote.Kite, value)
				go updaterFunc()
			}
		}),
	}

	// now trigger the remote kite so it sends us periodically an heartbeat
	remote.GoWithTimeout("kite.heartbeat", 4*time.Second, heartbeatArgs...)

	k.log.Info("Kite registered: %s", remote.Kite)

	remote.OnDisconnect(func() {
		k.log.Info("Kite disconnected: %s", remote.Kite)
		every.Stop()
	})

	// send response back to the kite, also identify him with the new name
	return &protocol.RegisterResult{URL: args.URL}, nil
}

func (k *Kontrol) handleGetKites(r *kite.Request) (interface{}, error) {
	// This type is here until inversion branch is merged.
	// Reason: We can't use the same struct for marshaling and unmarshaling.
	// TODO use the struct in protocol
	type GetKitesArgs struct {
		Query *protocol.KontrolQuery `json:"query"`
	}

	var args GetKitesArgs
	r.Args.One().MustUnmarshal(&args)

	query := args.Query

	// audience will go into the token as "aud" claim.
	audience := getAudience(query)

	// Generate token once here because we are using the same token for every
	// kite we return and generating many tokens is really slow.
	token, err := generateToken(audience, r.Username,
		k.Kite.Kite().Username, k.privateKey)
	if err != nil {
		return nil, err
	}

	// Get kites from the storage
	kites, err := k.storage.Get(query)
	if err != nil {
		return nil, err
	}

	// Attach tokens to kites
	kites.Attach(token)

	return &protocol.GetKitesResult{
		Kites: kites,
	}, nil
}

func (k *Kontrol) handleGetToken(r *kite.Request) (interface{}, error) {
	var query *protocol.KontrolQuery
	err := r.Args.One().Unmarshal(&query)
	if err != nil {
		return nil, errors.New("Invalid query")
	}

	// check if it's exist
	kites, err := k.storage.Get(query)
	if err != nil {
		return nil, err
	}

	if len(kites) > 1 {
		return nil, errors.New("query matches more than one kite")
	}

	audience := getAudience(query)

	return generateToken(audience, r.Username, k.Kite.Kite().Username, k.privateKey)
}

func (k *Kontrol) handleMachine(r *kite.Request) (interface{}, error) {
	if k.MachineAuthenticate != nil {
		if err := k.MachineAuthenticate(r); err != nil {
			return nil, errors.New("cannot authenticate user")
		}
	}

	username := r.Args.One().MustString() // username should be send as an argument
	return k.registerUser(username)
}
<commit_msg>kontrol: simplify heartbeat timer func<commit_after>package kontrol

import (
	"errors"
	"fmt"
	"net/http"
	"time"

	"github.com/koding/kite"
	"github.com/koding/kite/dnode"
	"github.com/koding/kite/kontrol/onceevery"
	kontrolprotocol "github.com/koding/kite/kontrol/protocol"
	"github.com/koding/kite/protocol"
)

func (k *Kontrol) handleHeartbeat(rw http.ResponseWriter, req *http.Request) {
	id := req.URL.Query().Get("id")

	k.heartbeatsMu.Lock()
	defer k.heartbeatsMu.Unlock()

	if updateTimer, ok := k.heartbeats[id]; ok {
		// try to reset the timer every time the remote kite sends sends us a
		// heartbeat. Because the timer get reset, the timer is never fired, so
		// the value get always updated with the updater in the background
		// according to the write interval. If the kite doesn't send any
		// heartbeat, the timer func is being called, which stops the updater
		// so the key is being deleted automatically via the TTL mechanism.
		updateTimer.Reset(HeartbeatInterval + HeartbeatDelay)
		rw.Write([]byte("pong"))
		return
	}

	// if we reach here than it has several meanings:
	// * kite was registered before, but kontrol is restarted
	// * kite was registered before, but kontrol has lost track
	// * kite was no registered and someone else sends an heartbeat
	// we send back "registeragain" so the caller can be added in to the
	// heartbeats map above.
	rw.Write([]byte("registeragain"))
}

func (k *Kontrol) handleRegisterHTTP(r *kite.Request) (interface{}, error) {
	k.log.Info("Register (via HTTP) request from: %s", r.Client.Kite)

	if r.Args.One().MustMap()["url"].MustString() == "" {
		return nil, errors.New("invalid url")
	}

	var args struct {
		URL string `json:"url"`
	}
	r.Args.One().MustUnmarshal(&args)
	if args.URL == "" {
		return nil, errors.New("empty url")
	}

	// Only accept requests with kiteKey because we need this info
	// for generating tokens for this kite.
	if r.Auth.Type != "kiteKey" {
		return nil, fmt.Errorf("Unexpected authentication type: %s", r.Auth.Type)
	}

	kiteURL := args.URL
	remote := r.Client

	if err := validateKiteKey(&remote.Kite); err != nil {
		return nil, err
	}

	value := &kontrolprotocol.RegisterValue{
		URL: kiteURL,
	}

	// Register first by adding the value to the storage. Return if there is
	// any error.
	if err := k.storage.Upsert(&remote.Kite, value); err != nil {
		k.log.Error("storage add '%s' error: %s", remote.Kite, err)
		return nil, errors.New("internal error - register")
	}

	// we create a new ticker which is going to update the key periodically in
	// the storage so it's always up to date. Instead of updating the key
	// periodically according to the HeartBeatInterval below, we are buffering
	// the write speed here with the UpdateInterval.
	stopped := make(chan struct{})
	updater := time.NewTicker(UpdateInterval)
	updaterFunc := func() {
		for {
			select {
			case <-updater.C:
				k.log.Info("Kite is active (via HTTP), updating the value %s", remote.Kite)
				err := k.storage.Update(&remote.Kite, value)
				if err != nil {
					k.log.Error("storage update '%s' error: %s", remote.Kite, err)
				}
			case <-stopped:
				k.log.Info("Kite is nonactive (via HTTP). Updater is closed %s", remote.Kite)
				return
			}
		}
	}
	go updaterFunc()

	// we are now creating a timer that is going to call the function which
	// stops the background updater if it's not resetted. The time is being
	// resetted on a separate HTTP endpoint "/heartbeat"
	k.heartbeatsMu.Lock()
	k.heartbeats[remote.Kite.ID] = time.AfterFunc(HeartbeatInterval+HeartbeatDelay, func() {
		k.log.Info("Kite didn't get heartbeat (via HTTP). Stopping the updater %s",
			remote.Kite)
		// stop the updater so it doesn't update it in the background
		updater.Stop()

		k.heartbeatsMu.Lock()
		defer k.heartbeatsMu.Unlock()

		select {
		case <-stopped:
		default:
			close(stopped)
		}

		delete(k.heartbeats, remote.Kite.ID)
	})
	k.heartbeatsMu.Unlock()

	k.log.Info("Kite registered (via HTTP): %s", remote.Kite)

	// send response back to the kite, also identify him with the new name
	return &protocol.RegisterResult{
		URL:               args.URL,
		HeartbeatInterval: int64(HeartbeatInterval / time.Second),
	}, nil
}

func (k *Kontrol) handleRegister(r *kite.Request) (interface{}, error) {
	k.log.Info("Register request from: %s", r.Client.Kite)

	if r.Args.One().MustMap()["url"].MustString() == "" {
		return nil, errors.New("invalid url")
	}

	var args struct {
		URL string `json:"url"`
	}
	r.Args.One().MustUnmarshal(&args)
	if args.URL == "" {
		return nil, errors.New("empty url")
	}

	// Only accept requests with kiteKey because we need this info
	// for generating tokens for this kite.
	if r.Auth.Type != "kiteKey" {
		return nil, fmt.Errorf("Unexpected authentication type: %s", r.Auth.Type)
	}

	kiteURL := args.URL
	remote := r.Client

	if err := validateKiteKey(&remote.Kite); err != nil {
		return nil, err
	}

	value := &kontrolprotocol.RegisterValue{
		URL: kiteURL,
	}

	// Register first by adding the value to the storage. Return if there is
	// any error.
	if err := k.storage.Upsert(&remote.Kite, value); err != nil {
		k.log.Error("storage add '%s' error: %s", remote.Kite, err)
		return nil, errors.New("internal error - register")
	}

	every := onceevery.New(UpdateInterval)

	ping := make(chan struct{}, 1)
	closed := false

	updaterFunc := func() {
		for {
			select {
			case <-ping:
				k.log.Debug("Kite is active, got a ping %s", remote.Kite)
				every.Do(func() {
					k.log.Info("Kite is active, updating the value %s", remote.Kite)
					err := k.storage.Update(&remote.Kite, value)
					if err != nil {
						k.log.Error("storage update '%s' error: %s", remote.Kite, err)
					}
				})
			case <-time.After(HeartbeatInterval + HeartbeatDelay):
				k.log.Info("Kite didn't sent any heartbeat %s.", remote.Kite)
				every.Stop()
				closed = true
				return
			}
		}
	}
	go updaterFunc()

	heartbeatArgs := []interface{}{
		HeartbeatInterval / time.Second,
		dnode.Callback(func(args *dnode.Partial) {
			k.log.Debug("Kite send us an heartbeat. %s", remote.Kite)

			k.clientLocks.Get(remote.Kite.ID).Lock()
			defer k.clientLocks.Get(remote.Kite.ID).Unlock()

			select {
			case ping <- struct{}{}:
			default:
			}

			// seems we miss a heartbeat, so start it again!
			if closed {
				closed = false
				k.log.Warning("Updater was closed, but we are still getting heartbeats. Starting again %s",
					remote.Kite)

				// it might be removed because the ttl cleaner would come
				// before us, so try to add it again, the updater will than
				// continue to update it afterwards.
				k.storage.Upsert(&remote.Kite, value)
				go updaterFunc()
			}
		}),
	}

	// now trigger the remote kite so it sends us periodically an heartbeat
	remote.GoWithTimeout("kite.heartbeat", 4*time.Second, heartbeatArgs...)

	k.log.Info("Kite registered: %s", remote.Kite)

	remote.OnDisconnect(func() {
		k.log.Info("Kite disconnected: %s", remote.Kite)
		every.Stop()
	})

	// send response back to the kite, also identify him with the new name
	return &protocol.RegisterResult{URL: args.URL}, nil
}

func (k *Kontrol) handleGetKites(r *kite.Request) (interface{}, error) {
	// This type is here until inversion branch is merged.
	// Reason: We can't use the same struct for marshaling and unmarshaling.
	// TODO use the struct in protocol
	type GetKitesArgs struct {
		Query *protocol.KontrolQuery `json:"query"`
	}

	var args GetKitesArgs
	r.Args.One().MustUnmarshal(&args)

	query := args.Query

	// audience will go into the token as "aud" claim.
	audience := getAudience(query)

	// Generate token once here because we are using the same token for every
	// kite we return and generating many tokens is really slow.
	token, err := generateToken(audience, r.Username,
		k.Kite.Kite().Username, k.privateKey)
	if err != nil {
		return nil, err
	}

	// Get kites from the storage
	kites, err := k.storage.Get(query)
	if err != nil {
		return nil, err
	}

	// Attach tokens to kites
	kites.Attach(token)

	return &protocol.GetKitesResult{
		Kites: kites,
	}, nil
}

func (k *Kontrol) handleGetToken(r *kite.Request) (interface{}, error) {
	var query *protocol.KontrolQuery
	err := r.Args.One().Unmarshal(&query)
	if err != nil {
		return nil, errors.New("Invalid query")
	}

	// check if it's exist
	kites, err := k.storage.Get(query)
	if err != nil {
		return nil, err
	}

	if len(kites) > 1 {
		return nil, errors.New("query matches more than one kite")
	}

	audience := getAudience(query)

	return generateToken(audience, r.Username, k.Kite.Kite().Username, k.privateKey)
}

func (k *Kontrol) handleMachine(r *kite.Request) (interface{}, error) {
	if k.MachineAuthenticate != nil {
		if err := k.MachineAuthenticate(r); err != nil {
			return nil, errors.New("cannot authenticate user")
		}
	}

	username := r.Args.One().MustString() // username should be send as an argument
	return k.registerUser(username)
}
<|endoftext|>


<commit_before>package app

import (
	"sync"
	"testing"
	"time"

	"github.com/moncho/dry/docker"
	"github.com/moncho/dry/mocks"
)

func TestDryRefresh(t *testing.T) {
	dry := newDryForTest()
	firstRefreshMark := dry.lastRefresh
	if !firstRefreshMark.Before(time.Now()) {
		t.Error("dry initial refresh mark is not correct")
	}
	dry.Refresh()
	lastRefresh := dry.lastRefresh
	if !firstRefreshMark.Before(lastRefresh) {
		t.Error("dry was not refreshed")
	}
	//This should not refresh since dry was just refreshed
	dry.tryRefresh()
	if !lastRefresh.Equal(dry.lastRefresh) {
		t.Error("dry was refreshed when it should not")
	}
	//but refresh always works
	dry.Refresh()
	if !lastRefresh.Before(dry.lastRefresh) {
		t.Error("dry was not refreshed when it should had been")
	}

}

func newDryForTest() *Dry {
	dry := &Dry{}
	dry.state = &state{
		changed:              true,
		showingAllContainers: false,
		SortMode:             docker.SortByContainerID,
		viewMode:             Main,
	}
	dry.dockerDaemon = new(mocks.ContainerDaemonMock)
	dry.refreshTimeMutex = &sync.Mutex{}
	dry.resetTimer()
	return dry
}
<commit_msg>Fix compilation error<commit_after>package app

import (
	"sync"
	"testing"
	"time"

	"github.com/moncho/dry/docker"
	"github.com/moncho/dry/mocks"
)

func TestDryRefresh(t *testing.T) {
	dry := newDryForTest()
	firstRefreshMark := dry.lastRefresh
	if !firstRefreshMark.Before(time.Now()) {
		t.Error("dry initial refresh mark is not correct")
	}
	dry.Refresh()
	lastRefresh := dry.lastRefresh
	if !firstRefreshMark.Before(lastRefresh) {
		t.Error("dry was not refreshed")
	}
	//This should not refresh since dry was just refreshed
	dry.tryRefresh()
	if !lastRefresh.Equal(dry.lastRefresh) {
		t.Error("dry was refreshed when it should not")
	}
	//but refresh always works
	dry.Refresh()
	if !lastRefresh.Before(dry.lastRefresh) {
		t.Error("dry was not refreshed when it should had been")
	}

}

func newDryForTest() *Dry {
	dry := &Dry{}
	dry.state = &state{
		changed:              true,
		showingAllContainers: false,
		SortMode:             docker.SortByContainerID,
		viewMode:             Main,
	}
	dry.dockerDaemon = new(mocks.ContainerDaemonMock)
	dry.resetTimer()
	return dry
}
<|endoftext|>


<commit_before>package arrays

// Checks if string exists in a slice
// Returns the key for needle if it is found in the slice, -1 otherwise.
func StringInSlice(haystack []string, needle string) (result int) {
	result = -1

	for k, v := range haystack {
		if v == needle {
			result = k
			break
		}
	}

	return
}
<commit_msg>Added function for checking if the integer is in slice.<commit_after>package arrays

// Checks if string exists in a slice
// Returns the key for needle if it is found in the slice, -1 otherwise.
func StringInSlice(haystack []string, needle string) (result int) {
	result = -1

	for k, v := range haystack {
		if v == needle {
			result = k
			break
		}
	}

	return
}

// Checks if integer exists in a slice
// Returns the index for needle if it is found in the slice, -1 otherwise.
func IntInSlice(haystack []int, needle int) (index int) {
	index = -1

	for k, v := range haystack {
		if v == needle {
			index = k
			break
		}
	}

	return
}
<|endoftext|>


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

package integration_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"runtime"
	"testing"

	"github.com/googlecloudplatform/gcsfuse/internal/canned"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
)

func TestMountHelper(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type MountHelperTest struct {
	// Path to the mount(8) helper binary.
	helperPath string

	// A temporary directory into which a file system may be mounted. Removed in
	// TearDown.
	dir string
}

var _ SetUpInterface = &MountHelperTest{}
var _ TearDownInterface = &MountHelperTest{}

func init() { RegisterTestSuite(&MountHelperTest{}) }

func (t *MountHelperTest) SetUp(_ *TestInfo) {
	var err error

	// Set up the appropriate helper path.
	switch runtime.GOOS {
	case "darwin":
		t.helperPath = path.Join(gBuildDir, "sbin/mount_gcsfuse")

	case "linux":
		t.helperPath = path.Join(gBuildDir, "sbin/mount.gcsfuse")

	default:
		AddFailure("Don't know how to deal with OS: %q", runtime.GOOS)
		AbortTest()
	}

	// Set up the temporary directory.
	t.dir, err = ioutil.TempDir("", "mount_helper_test")
	AssertEq(nil, err)
}

func (t *MountHelperTest) TearDown() {
	err := os.Remove(t.dir)
	AssertEq(nil, err)
}

func (t *MountHelperTest) mountHelperCommand(args []string) (cmd *exec.Cmd) {
	cmd = exec.Command(t.helperPath)
	cmd.Args = append(cmd.Args, args...)
	cmd.Env = []string{
		fmt.Sprintf("PATH=%s", path.Join(gBuildDir, "bin")),
	}

	return
}

func (t *MountHelperTest) mount(args []string) (err error) {
	cmd := t.mountHelperCommand(args)

	output, err := cmd.CombinedOutput()
	if err != nil {
		err = fmt.Errorf("CombinedOutput: %v\nOutput:\n%s", err, output)
		return
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *MountHelperTest) BadUsage() {
	testCases := []struct {
		args           []string
		expectedOutput string
	}{
		// Too few args
		0: {
			[]string{canned.FakeBucketName},
			"two positional arguments",
		},

		// Too many args
		1: {
			[]string{canned.FakeBucketName, "a", "b"},
			"Unexpected arg 3",
		},

		// Trailing -o
		2: {
			[]string{canned.FakeBucketName, "a", "-o"},
			"Unexpected -o",
		},
	}

	// Run each test case.
	for i, tc := range testCases {
		cmd := t.mountHelperCommand(tc.args)

		output, err := cmd.CombinedOutput()
		ExpectThat(err, Error(HasSubstr("exit status")), "case %d", i)
		ExpectThat(string(output), MatchesRegexp(tc.expectedOutput), "case %d", i)
	}
}

func (t *MountHelperTest) SuccessfulMount() {
	var err error
	var fi os.FileInfo

	// Mount.
	args := []string{canned.FakeBucketName, t.dir}

	err = t.mount(args)
	AssertEq(nil, err)
	defer unmount(t.dir)

	// Check that the file system is available.
	fi, err = os.Lstat(path.Join(t.dir, canned.TopLevelFile))
	AssertEq(nil, err)
	ExpectEq(os.FileMode(0644), fi.Mode())
	ExpectEq(len(canned.TopLevelFile_Contents), fi.Size())
}

func (t *MountHelperTest) RelativeMountPoint() {
	var err error
	var fi os.FileInfo

	// Mount with a relative mount point.
	cmd := t.mountHelperCommand([]string{
		canned.FakeBucketName,
		path.Base(t.dir),
	})

	cmd.Dir = path.Dir(t.dir)

	output, err := cmd.CombinedOutput()
	AssertEq(nil, err, "output:\n%s", output)

	defer unmount(t.dir)

	// The file system should be available.
	fi, err = os.Lstat(path.Join(t.dir, canned.TopLevelFile))
	AssertEq(nil, err)
	ExpectEq(os.FileMode(0644), fi.Mode())
	ExpectEq(len(canned.TopLevelFile_Contents), fi.Size())
}

func (t *MountHelperTest) ReadOnlyMode() {
	var err error

	// Mount.
	args := []string{"-o", "ro", canned.FakeBucketName, t.dir}

	err = t.mount(args)
	AssertEq(nil, err)
	defer unmount(t.dir)

	// Writing to the file system should fail.
	err = ioutil.WriteFile(path.Join(t.dir, "blah"), []byte{}, 0400)
	ExpectThat(err, Error(HasSubstr("read-only")))
}

func (t *MountHelperTest) ExtraneousOptions() {
	var err error
	var fi os.FileInfo

	// Mount with extra junk that shouldn't be passed on.
	args := []string{
		"-o", "noauto,nouser,auto,user",
		canned.FakeBucketName,
		t.dir,
	}

	err = t.mount(args)
	AssertEq(nil, err)
	defer unmount(t.dir)

	// Check that the file system is available.
	fi, err = os.Lstat(path.Join(t.dir, canned.TopLevelFile))
	AssertEq(nil, err)
	ExpectEq(os.FileMode(0644), fi.Mode())
	ExpectEq(len(canned.TopLevelFile_Contents), fi.Size())
}

func (t *MountHelperTest) LinuxArgumentOrder() {
	var err error

	// Linux places the options at the end.
	args := []string{canned.FakeBucketName, t.dir, "-o", "ro"}

	err = t.mount(args)
	AssertEq(nil, err)
	defer unmount(t.dir)

	// Writing to the file system should fail.
	err = ioutil.WriteFile(path.Join(t.dir, "blah"), []byte{}, 0400)
	ExpectThat(err, Error(HasSubstr("read-only")))
}

func (t *MountHelperTest) FuseSubtype() {
	var err error
	var fi os.FileInfo

	// This test isn't relevant except on Linux.
	if runtime.GOOS != "linux" {
		return
	}

	// Mount using the tool that would be invoked by ~mount -t fuse.gcsfuse`.
	t.helperPath = path.Join(gBuildDir, "sbin/mount.fuse.gcsfuse")
	args := []string{canned.FakeBucketName, t.dir}

	err = t.mount(args)
	AssertEq(nil, err)
	defer unmount(t.dir)

	// Check that the file system is available.
	fi, err = os.Lstat(path.Join(t.dir, canned.TopLevelFile))
	AssertEq(nil, err)
	ExpectEq(os.FileMode(0644), fi.Mode())
	ExpectEq(len(canned.TopLevelFile_Contents), fi.Size())
}
<commit_msg>MountHelperTest.ImplicitDirs<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package integration_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"runtime"
	"testing"

	"github.com/googlecloudplatform/gcsfuse/internal/canned"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
)

func TestMountHelper(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type MountHelperTest struct {
	// Path to the mount(8) helper binary.
	helperPath string

	// A temporary directory into which a file system may be mounted. Removed in
	// TearDown.
	dir string
}

var _ SetUpInterface = &MountHelperTest{}
var _ TearDownInterface = &MountHelperTest{}

func init() { RegisterTestSuite(&MountHelperTest{}) }

func (t *MountHelperTest) SetUp(_ *TestInfo) {
	var err error

	// Set up the appropriate helper path.
	switch runtime.GOOS {
	case "darwin":
		t.helperPath = path.Join(gBuildDir, "sbin/mount_gcsfuse")

	case "linux":
		t.helperPath = path.Join(gBuildDir, "sbin/mount.gcsfuse")

	default:
		AddFailure("Don't know how to deal with OS: %q", runtime.GOOS)
		AbortTest()
	}

	// Set up the temporary directory.
	t.dir, err = ioutil.TempDir("", "mount_helper_test")
	AssertEq(nil, err)
}

func (t *MountHelperTest) TearDown() {
	err := os.Remove(t.dir)
	AssertEq(nil, err)
}

func (t *MountHelperTest) mountHelperCommand(args []string) (cmd *exec.Cmd) {
	cmd = exec.Command(t.helperPath)
	cmd.Args = append(cmd.Args, args...)
	cmd.Env = []string{
		fmt.Sprintf("PATH=%s", path.Join(gBuildDir, "bin")),
	}

	return
}

func (t *MountHelperTest) mount(args []string) (err error) {
	cmd := t.mountHelperCommand(args)

	output, err := cmd.CombinedOutput()
	if err != nil {
		err = fmt.Errorf("CombinedOutput: %v\nOutput:\n%s", err, output)
		return
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *MountHelperTest) BadUsage() {
	testCases := []struct {
		args           []string
		expectedOutput string
	}{
		// Too few args
		0: {
			[]string{canned.FakeBucketName},
			"two positional arguments",
		},

		// Too many args
		1: {
			[]string{canned.FakeBucketName, "a", "b"},
			"Unexpected arg 3",
		},

		// Trailing -o
		2: {
			[]string{canned.FakeBucketName, "a", "-o"},
			"Unexpected -o",
		},
	}

	// Run each test case.
	for i, tc := range testCases {
		cmd := t.mountHelperCommand(tc.args)

		output, err := cmd.CombinedOutput()
		ExpectThat(err, Error(HasSubstr("exit status")), "case %d", i)
		ExpectThat(string(output), MatchesRegexp(tc.expectedOutput), "case %d", i)
	}
}

func (t *MountHelperTest) SuccessfulMount() {
	var err error
	var fi os.FileInfo

	// Mount.
	args := []string{canned.FakeBucketName, t.dir}

	err = t.mount(args)
	AssertEq(nil, err)
	defer unmount(t.dir)

	// Check that the file system is available.
	fi, err = os.Lstat(path.Join(t.dir, canned.TopLevelFile))
	AssertEq(nil, err)
	ExpectEq(os.FileMode(0644), fi.Mode())
	ExpectEq(len(canned.TopLevelFile_Contents), fi.Size())
}

func (t *MountHelperTest) RelativeMountPoint() {
	var err error
	var fi os.FileInfo

	// Mount with a relative mount point.
	cmd := t.mountHelperCommand([]string{
		canned.FakeBucketName,
		path.Base(t.dir),
	})

	cmd.Dir = path.Dir(t.dir)

	output, err := cmd.CombinedOutput()
	AssertEq(nil, err, "output:\n%s", output)

	defer unmount(t.dir)

	// The file system should be available.
	fi, err = os.Lstat(path.Join(t.dir, canned.TopLevelFile))
	AssertEq(nil, err)
	ExpectEq(os.FileMode(0644), fi.Mode())
	ExpectEq(len(canned.TopLevelFile_Contents), fi.Size())
}

func (t *MountHelperTest) ReadOnlyMode() {
	var err error

	// Mount.
	args := []string{"-o", "ro", canned.FakeBucketName, t.dir}

	err = t.mount(args)
	AssertEq(nil, err)
	defer unmount(t.dir)

	// Writing to the file system should fail.
	err = ioutil.WriteFile(path.Join(t.dir, "blah"), []byte{}, 0400)
	ExpectThat(err, Error(HasSubstr("read-only")))
}

func (t *MountHelperTest) ExtraneousOptions() {
	var err error
	var fi os.FileInfo

	// Mount with extra junk that shouldn't be passed on.
	args := []string{
		"-o", "noauto,nouser,auto,user",
		canned.FakeBucketName,
		t.dir,
	}

	err = t.mount(args)
	AssertEq(nil, err)
	defer unmount(t.dir)

	// Check that the file system is available.
	fi, err = os.Lstat(path.Join(t.dir, canned.TopLevelFile))
	AssertEq(nil, err)
	ExpectEq(os.FileMode(0644), fi.Mode())
	ExpectEq(len(canned.TopLevelFile_Contents), fi.Size())
}

func (t *MountHelperTest) LinuxArgumentOrder() {
	var err error

	// Linux places the options at the end.
	args := []string{canned.FakeBucketName, t.dir, "-o", "ro"}

	err = t.mount(args)
	AssertEq(nil, err)
	defer unmount(t.dir)

	// Writing to the file system should fail.
	err = ioutil.WriteFile(path.Join(t.dir, "blah"), []byte{}, 0400)
	ExpectThat(err, Error(HasSubstr("read-only")))
}

func (t *MountHelperTest) FuseSubtype() {
	var err error
	var fi os.FileInfo

	// This test isn't relevant except on Linux.
	if runtime.GOOS != "linux" {
		return
	}

	// Mount using the tool that would be invoked by ~mount -t fuse.gcsfuse`.
	t.helperPath = path.Join(gBuildDir, "sbin/mount.fuse.gcsfuse")
	args := []string{canned.FakeBucketName, t.dir}

	err = t.mount(args)
	AssertEq(nil, err)
	defer unmount(t.dir)

	// Check that the file system is available.
	fi, err = os.Lstat(path.Join(t.dir, canned.TopLevelFile))
	AssertEq(nil, err)
	ExpectEq(os.FileMode(0644), fi.Mode())
	ExpectEq(len(canned.TopLevelFile_Contents), fi.Size())
}

func (t *MountHelperTest) ImplicitDirs() {
	var err error

	// Mount.
	args := []string{"-o", "implicit_dirs", canned.FakeBucketName, t.dir}

	err = t.mount(args)
	AssertEq(nil, err)
	defer unmount(t.dir)

	// The implicit directory should be visible.
	fi, err := os.Lstat(path.Join(t.dir, canned.ImplicitDirFile))
	AssertEq(nil, err)
	ExpectTrue(fi.IsDir())
}
<|endoftext|>


<commit_before>package gorm

import (
	"errors"
	"fmt"
	"reflect"
	"strings"
)

// JoinTableHandlerInterface is an interface for how to handle many2many relations
type JoinTableHandlerInterface interface {
	// initialize join table handler
	Setup(relationship *Relationship, tableName string, source reflect.Type, destination reflect.Type)
	// Table return join table's table name
	Table(db *DB) string
	// Add create relationship in join table for source and destination
	Add(handler JoinTableHandlerInterface, db *DB, source interface{}, destination interface{}) error
	// Delete delete relationship in join table for sources
	Delete(handler JoinTableHandlerInterface, db *DB, sources ...interface{}) error
	// JoinWith query with `Join` conditions
	JoinWith(handler JoinTableHandlerInterface, db *DB, source interface{}) *DB
	// SourceForeignKeys return source foreign keys
	SourceForeignKeys() []JoinTableForeignKey
	// DestinationForeignKeys return destination foreign keys
	DestinationForeignKeys() []JoinTableForeignKey
}

// JoinTableForeignKey join table foreign key struct
type JoinTableForeignKey struct {
	DBName            string
	AssociationDBName string
}

// JoinTableSource is a struct that contains model type and foreign keys
type JoinTableSource struct {
	ModelType   reflect.Type
	ForeignKeys []JoinTableForeignKey
}

// JoinTableHandler default join table handler
type JoinTableHandler struct {
	TableName   string          `sql:"-"`
	Source      JoinTableSource `sql:"-"`
	Destination JoinTableSource `sql:"-"`
}

// SourceForeignKeys return source foreign keys
func (s *JoinTableHandler) SourceForeignKeys() []JoinTableForeignKey {
	return s.Source.ForeignKeys
}

// DestinationForeignKeys return destination foreign keys
func (s *JoinTableHandler) DestinationForeignKeys() []JoinTableForeignKey {
	return s.Destination.ForeignKeys
}

// Setup initialize a default join table handler
func (s *JoinTableHandler) Setup(relationship *Relationship, tableName string, source reflect.Type, destination reflect.Type) {
	s.TableName = tableName

	s.Source = JoinTableSource{ModelType: source}
	s.Source.ForeignKeys = []JoinTableForeignKey{}
	for idx, dbName := range relationship.ForeignFieldNames {
		s.Source.ForeignKeys = append(s.Source.ForeignKeys, JoinTableForeignKey{
			DBName:            relationship.ForeignDBNames[idx],
			AssociationDBName: dbName,
		})
	}

	s.Destination = JoinTableSource{ModelType: destination}
	s.Destination.ForeignKeys = []JoinTableForeignKey{}
	for idx, dbName := range relationship.AssociationForeignFieldNames {
		s.Destination.ForeignKeys = append(s.Destination.ForeignKeys, JoinTableForeignKey{
			DBName:            relationship.AssociationForeignDBNames[idx],
			AssociationDBName: dbName,
		})
	}
}

// Table return join table's table name
func (s JoinTableHandler) Table(db *DB) string {
	return s.TableName
}

func (s JoinTableHandler) updateConditionMap(conditionMap map[string]interface{}, db *DB, joinTableSources []JoinTableSource, sources ...interface{}) {
	for _, source := range sources {
		scope := db.NewScope(source)
		modelType := scope.GetModelStruct().ModelType

		for _, joinTableSource := range joinTableSources {
			if joinTableSource.ModelType == modelType {
				for _, foreignKey := range joinTableSource.ForeignKeys {
					if field, ok := scope.FieldByName(foreignKey.AssociationDBName); ok {
						conditionMap[foreignKey.DBName] = field.Field.Interface()
					}
				}
				break
			}
		}
	}
}

// Add create relationship in join table for source and destination
func (s JoinTableHandler) Add(handler JoinTableHandlerInterface, db *DB, source interface{}, destination interface{}) error {
	var (
		scope        = db.NewScope("")
		conditionMap = map[string]interface{}{}
	)

	// Update condition map for source
	s.updateConditionMap(conditionMap, db, []JoinTableSource{s.Source}, source)

	// Update condition map for destination
	s.updateConditionMap(conditionMap, db, []JoinTableSource{s.Destination}, destination)

	var assignColumns, binVars, conditions []string
	var values []interface{}
	for key, value := range conditionMap {
		assignColumns = append(assignColumns, scope.Quote(key))
		binVars = append(binVars, `?`)
		conditions = append(conditions, fmt.Sprintf("%v = ?", scope.Quote(key)))
		values = append(values, value)
	}

	for _, value := range values {
		values = append(values, value)
	}

	quotedTable := scope.Quote(handler.Table(db))
	sql := fmt.Sprintf(
		"INSERT INTO %v (%v) SELECT %v %v WHERE NOT EXISTS (SELECT * FROM %v WHERE %v)",
		quotedTable,
		strings.Join(assignColumns, ","),
		strings.Join(binVars, ","),
		scope.Dialect().SelectFromDummyTable(),
		quotedTable,
		strings.Join(conditions, " AND "),
	)

	return db.Exec(sql, values...).Error
}

// Delete delete relationship in join table for sources
func (s JoinTableHandler) Delete(handler JoinTableHandlerInterface, db *DB, sources ...interface{}) error {
	var (
		scope        = db.NewScope(nil)
		conditions   []string
		values       []interface{}
		conditionMap = map[string]interface{}{}
	)

	s.updateConditionMap(conditionMap, db, []JoinTableSource{s.Source, s.Destination}, sources...)

	for key, value := range conditionMap {
		conditions = append(conditions, fmt.Sprintf("%v = ?", scope.Quote(key)))
		values = append(values, value)
	}

	return db.Table(handler.Table(db)).Where(strings.Join(conditions, " AND "), values...).Delete("").Error
}

// JoinWith query with `Join` conditions
func (s JoinTableHandler) JoinWith(handler JoinTableHandlerInterface, db *DB, source interface{}) *DB {
	var (
		scope           = db.NewScope(source)
		tableName       = handler.Table(db)
		quotedTableName = scope.Quote(tableName)
		joinConditions  []string
		values          []interface{}
	)

	if s.Source.ModelType == scope.GetModelStruct().ModelType {
		destinationTableName := db.NewScope(reflect.New(s.Destination.ModelType).Interface()).QuotedTableName()
		for _, foreignKey := range s.Destination.ForeignKeys {
			joinConditions = append(joinConditions, fmt.Sprintf("%v.%v = %v.%v", quotedTableName, scope.Quote(foreignKey.DBName), destinationTableName, scope.Quote(foreignKey.AssociationDBName)))
		}

		var foreignDBNames []string
		var foreignFieldNames []string

		for _, foreignKey := range s.Source.ForeignKeys {
			foreignDBNames = append(foreignDBNames, foreignKey.DBName)
			if field, ok := scope.FieldByName(foreignKey.AssociationDBName); ok {
				foreignFieldNames = append(foreignFieldNames, field.Name)
			}
		}

		foreignFieldValues := scope.getColumnAsArray(foreignFieldNames, scope.Value)

		var condString string
		if len(foreignFieldValues) > 0 {
			var quotedForeignDBNames []string
			for _, dbName := range foreignDBNames {
				quotedForeignDBNames = append(quotedForeignDBNames, tableName+"."+dbName)
			}

			condString = fmt.Sprintf("%v IN (%v)", toQueryCondition(scope, quotedForeignDBNames), toQueryMarks(foreignFieldValues))

			keys := scope.getColumnAsArray(foreignFieldNames, scope.Value)
			values = append(values, toQueryValues(keys))
		} else {
			condString = fmt.Sprintf("1 <> 1")
		}

		return db.Joins(fmt.Sprintf("INNER JOIN %v ON %v", quotedTableName, strings.Join(joinConditions, " AND "))).
			Where(condString, toQueryValues(foreignFieldValues)...)
	}

	db.Error = errors.New("wrong source type for join table handler")
	return db
}
<commit_msg>Call DefaultTableNameHandler for JoinTableHandler's table<commit_after>package gorm

import (
	"errors"
	"fmt"
	"reflect"
	"strings"
)

// JoinTableHandlerInterface is an interface for how to handle many2many relations
type JoinTableHandlerInterface interface {
	// initialize join table handler
	Setup(relationship *Relationship, tableName string, source reflect.Type, destination reflect.Type)
	// Table return join table's table name
	Table(db *DB) string
	// Add create relationship in join table for source and destination
	Add(handler JoinTableHandlerInterface, db *DB, source interface{}, destination interface{}) error
	// Delete delete relationship in join table for sources
	Delete(handler JoinTableHandlerInterface, db *DB, sources ...interface{}) error
	// JoinWith query with `Join` conditions
	JoinWith(handler JoinTableHandlerInterface, db *DB, source interface{}) *DB
	// SourceForeignKeys return source foreign keys
	SourceForeignKeys() []JoinTableForeignKey
	// DestinationForeignKeys return destination foreign keys
	DestinationForeignKeys() []JoinTableForeignKey
}

// JoinTableForeignKey join table foreign key struct
type JoinTableForeignKey struct {
	DBName            string
	AssociationDBName string
}

// JoinTableSource is a struct that contains model type and foreign keys
type JoinTableSource struct {
	ModelType   reflect.Type
	ForeignKeys []JoinTableForeignKey
}

// JoinTableHandler default join table handler
type JoinTableHandler struct {
	TableName   string          `sql:"-"`
	Source      JoinTableSource `sql:"-"`
	Destination JoinTableSource `sql:"-"`
}

// SourceForeignKeys return source foreign keys
func (s *JoinTableHandler) SourceForeignKeys() []JoinTableForeignKey {
	return s.Source.ForeignKeys
}

// DestinationForeignKeys return destination foreign keys
func (s *JoinTableHandler) DestinationForeignKeys() []JoinTableForeignKey {
	return s.Destination.ForeignKeys
}

// Setup initialize a default join table handler
func (s *JoinTableHandler) Setup(relationship *Relationship, tableName string, source reflect.Type, destination reflect.Type) {
	s.TableName = tableName

	s.Source = JoinTableSource{ModelType: source}
	s.Source.ForeignKeys = []JoinTableForeignKey{}
	for idx, dbName := range relationship.ForeignFieldNames {
		s.Source.ForeignKeys = append(s.Source.ForeignKeys, JoinTableForeignKey{
			DBName:            relationship.ForeignDBNames[idx],
			AssociationDBName: dbName,
		})
	}

	s.Destination = JoinTableSource{ModelType: destination}
	s.Destination.ForeignKeys = []JoinTableForeignKey{}
	for idx, dbName := range relationship.AssociationForeignFieldNames {
		s.Destination.ForeignKeys = append(s.Destination.ForeignKeys, JoinTableForeignKey{
			DBName:            relationship.AssociationForeignDBNames[idx],
			AssociationDBName: dbName,
		})
	}
}

// Table return join table's table name
func (s JoinTableHandler) Table(db *DB) string {
	return DefaultTableNameHandler(db, s.TableName)
}

func (s JoinTableHandler) updateConditionMap(conditionMap map[string]interface{}, db *DB, joinTableSources []JoinTableSource, sources ...interface{}) {
	for _, source := range sources {
		scope := db.NewScope(source)
		modelType := scope.GetModelStruct().ModelType

		for _, joinTableSource := range joinTableSources {
			if joinTableSource.ModelType == modelType {
				for _, foreignKey := range joinTableSource.ForeignKeys {
					if field, ok := scope.FieldByName(foreignKey.AssociationDBName); ok {
						conditionMap[foreignKey.DBName] = field.Field.Interface()
					}
				}
				break
			}
		}
	}
}

// Add create relationship in join table for source and destination
func (s JoinTableHandler) Add(handler JoinTableHandlerInterface, db *DB, source interface{}, destination interface{}) error {
	var (
		scope        = db.NewScope("")
		conditionMap = map[string]interface{}{}
	)

	// Update condition map for source
	s.updateConditionMap(conditionMap, db, []JoinTableSource{s.Source}, source)

	// Update condition map for destination
	s.updateConditionMap(conditionMap, db, []JoinTableSource{s.Destination}, destination)

	var assignColumns, binVars, conditions []string
	var values []interface{}
	for key, value := range conditionMap {
		assignColumns = append(assignColumns, scope.Quote(key))
		binVars = append(binVars, `?`)
		conditions = append(conditions, fmt.Sprintf("%v = ?", scope.Quote(key)))
		values = append(values, value)
	}

	for _, value := range values {
		values = append(values, value)
	}

	quotedTable := scope.Quote(handler.Table(db))
	sql := fmt.Sprintf(
		"INSERT INTO %v (%v) SELECT %v %v WHERE NOT EXISTS (SELECT * FROM %v WHERE %v)",
		quotedTable,
		strings.Join(assignColumns, ","),
		strings.Join(binVars, ","),
		scope.Dialect().SelectFromDummyTable(),
		quotedTable,
		strings.Join(conditions, " AND "),
	)

	return db.Exec(sql, values...).Error
}

// Delete delete relationship in join table for sources
func (s JoinTableHandler) Delete(handler JoinTableHandlerInterface, db *DB, sources ...interface{}) error {
	var (
		scope        = db.NewScope(nil)
		conditions   []string
		values       []interface{}
		conditionMap = map[string]interface{}{}
	)

	s.updateConditionMap(conditionMap, db, []JoinTableSource{s.Source, s.Destination}, sources...)

	for key, value := range conditionMap {
		conditions = append(conditions, fmt.Sprintf("%v = ?", scope.Quote(key)))
		values = append(values, value)
	}

	return db.Table(handler.Table(db)).Where(strings.Join(conditions, " AND "), values...).Delete("").Error
}

// JoinWith query with `Join` conditions
func (s JoinTableHandler) JoinWith(handler JoinTableHandlerInterface, db *DB, source interface{}) *DB {
	var (
		scope           = db.NewScope(source)
		tableName       = handler.Table(db)
		quotedTableName = scope.Quote(tableName)
		joinConditions  []string
		values          []interface{}
	)

	if s.Source.ModelType == scope.GetModelStruct().ModelType {
		destinationTableName := db.NewScope(reflect.New(s.Destination.ModelType).Interface()).QuotedTableName()
		for _, foreignKey := range s.Destination.ForeignKeys {
			joinConditions = append(joinConditions, fmt.Sprintf("%v.%v = %v.%v", quotedTableName, scope.Quote(foreignKey.DBName), destinationTableName, scope.Quote(foreignKey.AssociationDBName)))
		}

		var foreignDBNames []string
		var foreignFieldNames []string

		for _, foreignKey := range s.Source.ForeignKeys {
			foreignDBNames = append(foreignDBNames, foreignKey.DBName)
			if field, ok := scope.FieldByName(foreignKey.AssociationDBName); ok {
				foreignFieldNames = append(foreignFieldNames, field.Name)
			}
		}

		foreignFieldValues := scope.getColumnAsArray(foreignFieldNames, scope.Value)

		var condString string
		if len(foreignFieldValues) > 0 {
			var quotedForeignDBNames []string
			for _, dbName := range foreignDBNames {
				quotedForeignDBNames = append(quotedForeignDBNames, tableName+"."+dbName)
			}

			condString = fmt.Sprintf("%v IN (%v)", toQueryCondition(scope, quotedForeignDBNames), toQueryMarks(foreignFieldValues))

			keys := scope.getColumnAsArray(foreignFieldNames, scope.Value)
			values = append(values, toQueryValues(keys))
		} else {
			condString = fmt.Sprintf("1 <> 1")
		}

		return db.Joins(fmt.Sprintf("INNER JOIN %v ON %v", quotedTableName, strings.Join(joinConditions, " AND "))).
			Where(condString, toQueryValues(foreignFieldValues)...)
	}

	db.Error = errors.New("wrong source type for join table handler")
	return db
}
<|endoftext|>


<commit_before>package search

import (
	"errors"
	"math"
	"sort"

	gr "github.com/gonum/graph"
)

// Finds all shortest paths between start and goal
type AllPathFunc func(start, goal gr.Node) (path [][]gr.Node, cost float64, err error)

// Finds one path between start and goal, which it finds is arbitrary
type SinglePathFunc func(start, goal gr.Node) (path []gr.Node, cost float64, err error)

// This function returns two functions: one that will generate all shortest paths between two nodes with ids i and j, and one that will generate just one path.
//
// This algorithm requires the CrunchGraph interface which means it only works on nodes with dense ids since it uses an adjacency matrix.
//
// This algorithm isn't blazingly fast, but is relatively fast for the domain. It runs at O((number of vertices)^3), and successfully computes
// the cost between all pairs of vertices.
//
// Generating a single path should be pretty cheap after FW is done running. The AllPathFunc is likely to be considerably more expensive,
// simply because it has to effectively generate all combinations of known valid paths at each recursive step of the algorithm.
func FloydWarshall(graph gr.CrunchGraph, cost func(gr.Node, gr.Node) float64) (AllPathFunc, SinglePathFunc) {
	graph.Crunch()
	successors, _, _, isSuccessor, _, _, cost, _ := setupFuncs(graph, cost, nil)

	nodes := denseNodeSorter(graph.NodeList())
	sort.Sort(nodes)
	numNodes := len(nodes)

	dist := make([]float64, numNodes*numNodes)
	next := make([][]int, numNodes*numNodes)
	for i := 0; i < numNodes; i++ {
		for j := 0; j < numNodes; j++ {
			if j != i {
				dist[i+j*numNodes] = math.Inf(1)
			}
		}
	}

	for _, node := range nodes {
		for _, succ := range successors(node) {
			dist[node.ID()+succ.ID()*numNodes] = cost(node, succ)
		}
	}

	for k := 0; k < numNodes; k++ {
		for i := 0; i < numNodes; i++ {
			for j := 0; j < numNodes; j++ {
				if dist[i+j*numNodes] > dist[i+k*numNodes]+dist[k+j*numNodes] {
					dist[i+j*numNodes] = dist[i+k*numNodes] + dist[k+j*numNodes]

					// Avoid generating too much garbage by reusing the memory in the list if we've allocated one already
					if next[i+j*numNodes] == nil {
						next[i+j*numNodes] = []int{k}
					} else {
						next[i+j*numNodes] = next[i+j*numNodes][:1]
						next[i+j*numNodes][0] = k
					}
					// If the cost between the nodes happens to be the same cost as what we know, add the approriate
					// intermediary to the list
				} else if math.Abs(dist[i+k*numNodes]+dist[k+j*numNodes]-dist[i+j*numNodes]) < 0.00001 && i != k && i != j && j != k {
					next[i+j*numNodes] = append(next[i+j*numNodes], k)
				}
			}
		}
	}

	return genAllPathsFunc(dist, next, nodes, graph, cost, isSuccessor), genSinglePathFunc(dist, next, nodes)
}

func genAllPathsFunc(dist []float64, next [][]int, nodes []gr.Node, graph gr.Graph, cost func(gr.Node, gr.Node) float64, isSuccessor func(gr.Node, gr.Node) bool) func(start, goal gr.Node) ([][]gr.Node, float64, error) {
	numNodes := len(nodes)

	// A recursive function to reconstruct all possible paths.
	// It's not fast, but it's about as fast as can be reasonably expected
	var allPathFinder func(i, j int) ([][]gr.Node, error)
	allPathFinder = func(i, j int) ([][]gr.Node, error) {
		if dist[i+j*numNodes] == math.Inf(1) {
			return nil, errors.New("No path")
		}
		intermediates := next[i+j*numNodes]
		if intermediates == nil || len(intermediates) == 0 {
			// There is exactly one path
			return [][]gr.Node{[]gr.Node{}}, nil
		}

		toReturn := make([][]gr.Node, 0, len(intermediates))
		// Special case: if intermediates exist we need to explicitly check to see if the i and j is also an optimal path
		if isSuccessor(nodes[i], nodes[j]) && math.Abs(dist[i+j*numNodes]-cost(nodes[i], nodes[j])) < .000001 {
			toReturn = append(toReturn, []gr.Node{nodes[i], nodes[j]})
		}

		// This step is a tad convoluted: we have some list of intermediates.
		// We can think of each intermediate as a path junction
		//
		// At this junction, we can find all the shortest paths back to i,
		// and all the shortest paths down to j. Since this is a junction,
		// any predecessor path that runs through this intermediate may
		// freely choose any successor path to get to j. They'll all be
		// of equivalent length.
		//
		// Thus, for each intermediate, we run through and join each predecessor
		// path with each successor path via its junction.
		for _, intermediate := range intermediates {

			// Find predecessors
			preds, err := allPathFinder(i, intermediate)
			if err != nil {
				return nil, err
			}

			// Join each predecessor with its junction
			for a := range preds {
				preds[a] = append(preds[a], nodes[intermediate])
			}

			// Find successors
			succs, err := allPathFinder(intermediate, j)
			if err != nil {
				return nil, err
			}

			// Join each successor with its predecessor at the junction.
			// (the copying stuff is because slices are reference types)
			for a := range succs {
				for b := range preds {
					path := make([]gr.Node, len(succs[a]), len(succs[a])+len(preds[b]))
					copy(path, succs[a])
					path = append(path, preds[b]...)
					toReturn = append(toReturn, path)
				}
			}

		}

		return toReturn, nil
	}

	return func(start, goal gr.Node) ([][]gr.Node, float64, error) {
		paths, err := allPathFinder(start.ID(), goal.ID())
		if err != nil {
			return nil, math.Inf(1), nil
		}

		for i := range paths {
			// Prepend start and postpend goal, but don't repeat start/goal

			if len(paths[i]) != 0 {
				if paths[i][0].ID() != start.ID() {
					paths[i] = append(paths[i], nil)
					copy(paths[i][1:], paths[i][:len(paths[i])-1])
					paths[i][0] = start
				}

				if paths[i][len(paths[i])-1].ID() != goal.ID() {
					paths[i] = append(paths[i], goal)
				}
			} else {
				paths[i] = append(paths[i], start, goal)
			}
		}

		return paths, dist[start.ID()+goal.ID()*numNodes], nil
	}
}

func genSinglePathFunc(dist []float64, next [][]int, nodes []gr.Node) func(start, goal gr.Node) ([]gr.Node, float64, error) {
	numNodes := len(nodes)

	var singlePathFinder func(i, j int) ([]gr.Node, error)
	singlePathFinder = func(i, j int) ([]gr.Node, error) {
		if dist[i+j*numNodes] == math.Inf(1) {
			return nil, errors.New("No path")
		}

		intermediates := next[i+j*numNodes]
		if intermediates == nil || len(intermediates) == 0 {
			return []gr.Node{}, nil
		}

		intermediate := intermediates[0]
		path, err := singlePathFinder(i, intermediate)
		if err != nil {
			return nil, err
		}

		path = append(path, nodes[intermediate])
		p, err := singlePathFinder(intermediate, j)
		if err != nil {
			return nil, err
		}
		path = append(path, p...)

		return path, nil
	}

	return func(start, goal gr.Node) ([]gr.Node, float64, error) {
		path, err := singlePathFinder(start.ID(), goal.ID())
		if err != nil {
			return nil, math.Inf(1), nil
		}

		path = append(path, nil)
		copy(path[1:], path[:len(path)-1])
		path[0] = start
		path = append(path, goal)

		return path, dist[start.ID()+goal.ID()*numNodes], nil
	}
}
<commit_msg>Made FW comment better<commit_after>package search

import (
	"errors"
	"math"
	"sort"

	gr "github.com/gonum/graph"
)

// Finds all shortest paths between start and goal
type AllPathFunc func(start, goal gr.Node) (path [][]gr.Node, cost float64, err error)

// Finds one path between start and goal, which it finds is arbitrary
type PathFunc func(start, goal gr.Node) (path []gr.Node, cost float64, err error)

// This function returns two functions: one that will generate all shortest paths between two nodes with ids i and j, and one that will generate just one path.
//
// This algorithm requires the CrunchGraph interface which means it only works on graphs with dense node ids since it uses an adjacency matrix.
//
// This algorithm isn't blazingly fast, but is relatively fast for the domain. It runs at O((number of vertices)^3) in best, worst, and average case,
//  and successfully computes the cost between all pairs of vertices.
//
// This function operates slightly differently from the others for convenience -- rather than generating paths and returning them to you,
// it gives you the option of calling one of two functions for each start/goal pair you need info for. One will return the path, cost,
// or an error if no path exists.
//
// The other will return the cost and an error if no path exists, but it will also return ALL possible shortest paths between start and goal.
// This is not too much more expensive than generating one path, but it does obviously increase with the number of paths.
func FloydWarshall(graph gr.CrunchGraph, cost func(gr.Node, gr.Node) float64) (AllPathFunc, PathFunc) {
	graph.Crunch()
	successors, _, _, isSuccessor, _, _, cost, _ := setupFuncs(graph, cost, nil)

	nodes := denseNodeSorter(graph.NodeList())
	sort.Sort(nodes)
	numNodes := len(nodes)

	dist := make([]float64, numNodes*numNodes)
	next := make([][]int, numNodes*numNodes)
	for i := 0; i < numNodes; i++ {
		for j := 0; j < numNodes; j++ {
			if j != i {
				dist[i+j*numNodes] = math.Inf(1)
			}
		}
	}

	for _, node := range nodes {
		for _, succ := range successors(node) {
			dist[node.ID()+succ.ID()*numNodes] = cost(node, succ)
		}
	}

	for k := 0; k < numNodes; k++ {
		for i := 0; i < numNodes; i++ {
			for j := 0; j < numNodes; j++ {
				if dist[i+j*numNodes] > dist[i+k*numNodes]+dist[k+j*numNodes] {
					dist[i+j*numNodes] = dist[i+k*numNodes] + dist[k+j*numNodes]

					// Avoid generating too much garbage by reusing the memory in the list if we've allocated one already
					if next[i+j*numNodes] == nil {
						next[i+j*numNodes] = []int{k}
					} else {
						next[i+j*numNodes] = next[i+j*numNodes][:1]
						next[i+j*numNodes][0] = k
					}
					// If the cost between the nodes happens to be the same cost as what we know, add the approriate
					// intermediary to the list
				} else if math.Abs(dist[i+k*numNodes]+dist[k+j*numNodes]-dist[i+j*numNodes]) < 0.00001 && i != k && i != j && j != k {
					next[i+j*numNodes] = append(next[i+j*numNodes], k)
				}
			}
		}
	}

	return genAllPathsFunc(dist, next, nodes, graph, cost, isSuccessor), genSinglePathFunc(dist, next, nodes)
}

func genAllPathsFunc(dist []float64, next [][]int, nodes []gr.Node, graph gr.Graph, cost func(gr.Node, gr.Node) float64, isSuccessor func(gr.Node, gr.Node) bool) func(start, goal gr.Node) ([][]gr.Node, float64, error) {
	numNodes := len(nodes)

	// A recursive function to reconstruct all possible paths.
	// It's not fast, but it's about as fast as can be reasonably expected
	var allPathFinder func(i, j int) ([][]gr.Node, error)
	allPathFinder = func(i, j int) ([][]gr.Node, error) {
		if dist[i+j*numNodes] == math.Inf(1) {
			return nil, errors.New("No path")
		}
		intermediates := next[i+j*numNodes]
		if intermediates == nil || len(intermediates) == 0 {
			// There is exactly one path
			return [][]gr.Node{[]gr.Node{}}, nil
		}

		toReturn := make([][]gr.Node, 0, len(intermediates))
		// Special case: if intermediates exist we need to explicitly check to see if the i and j is also an optimal path
		if isSuccessor(nodes[i], nodes[j]) && math.Abs(dist[i+j*numNodes]-cost(nodes[i], nodes[j])) < .000001 {
			toReturn = append(toReturn, []gr.Node{nodes[i], nodes[j]})
		}

		// This step is a tad convoluted: we have some list of intermediates.
		// We can think of each intermediate as a path junction
		//
		// At this junction, we can find all the shortest paths back to i,
		// and all the shortest paths down to j. Since this is a junction,
		// any predecessor path that runs through this intermediate may
		// freely choose any successor path to get to j. They'll all be
		// of equivalent length.
		//
		// Thus, for each intermediate, we run through and join each predecessor
		// path with each successor path via its junction.
		for _, intermediate := range intermediates {

			// Find predecessors
			preds, err := allPathFinder(i, intermediate)
			if err != nil {
				return nil, err
			}

			// Join each predecessor with its junction
			for a := range preds {
				preds[a] = append(preds[a], nodes[intermediate])
			}

			// Find successors
			succs, err := allPathFinder(intermediate, j)
			if err != nil {
				return nil, err
			}

			// Join each successor with its predecessor at the junction.
			// (the copying stuff is because slices are reference types)
			for a := range succs {
				for b := range preds {
					path := make([]gr.Node, len(succs[a]), len(succs[a])+len(preds[b]))
					copy(path, succs[a])
					path = append(path, preds[b]...)
					toReturn = append(toReturn, path)
				}
			}

		}

		return toReturn, nil
	}

	return func(start, goal gr.Node) ([][]gr.Node, float64, error) {
		paths, err := allPathFinder(start.ID(), goal.ID())
		if err != nil {
			return nil, math.Inf(1), nil
		}

		for i := range paths {
			// Prepend start and postpend goal, but don't repeat start/goal

			if len(paths[i]) != 0 {
				if paths[i][0].ID() != start.ID() {
					paths[i] = append(paths[i], nil)
					copy(paths[i][1:], paths[i][:len(paths[i])-1])
					paths[i][0] = start
				}

				if paths[i][len(paths[i])-1].ID() != goal.ID() {
					paths[i] = append(paths[i], goal)
				}
			} else {
				paths[i] = append(paths[i], start, goal)
			}
		}

		return paths, dist[start.ID()+goal.ID()*numNodes], nil
	}
}

func genSinglePathFunc(dist []float64, next [][]int, nodes []gr.Node) func(start, goal gr.Node) ([]gr.Node, float64, error) {
	numNodes := len(nodes)

	var singlePathFinder func(i, j int) ([]gr.Node, error)
	singlePathFinder = func(i, j int) ([]gr.Node, error) {
		if dist[i+j*numNodes] == math.Inf(1) {
			return nil, errors.New("No path")
		}

		intermediates := next[i+j*numNodes]
		if intermediates == nil || len(intermediates) == 0 {
			return []gr.Node{}, nil
		}

		intermediate := intermediates[0]
		path, err := singlePathFinder(i, intermediate)
		if err != nil {
			return nil, err
		}

		path = append(path, nodes[intermediate])
		p, err := singlePathFinder(intermediate, j)
		if err != nil {
			return nil, err
		}
		path = append(path, p...)

		return path, nil
	}

	return func(start, goal gr.Node) ([]gr.Node, float64, error) {
		path, err := singlePathFinder(start.ID(), goal.ID())
		if err != nil {
			return nil, math.Inf(1), nil
		}

		path = append(path, nil)
		copy(path[1:], path[:len(path)-1])
		path[0] = start
		path = append(path, goal)

		return path, dist[start.ID()+goal.ID()*numNodes], nil
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package cm

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/types"
	utilerrors "k8s.io/apimachinery/pkg/util/errors"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
	"k8s.io/klog"
	v1qos "k8s.io/kubernetes/pkg/apis/core/v1/helper/qos"
	kubefeatures "k8s.io/kubernetes/pkg/features"
)

const (
	podCgroupNamePrefix = "pod"
)

// podContainerManagerImpl implements podContainerManager interface.
// It is the general implementation which allows pod level container
// management if qos Cgroup is enabled.
type podContainerManagerImpl struct {
	// qosContainersInfo hold absolute paths of the top level qos containers
	qosContainersInfo QOSContainersInfo
	// Stores the mounted cgroup subsystems
	subsystems *CgroupSubsystems
	// cgroupManager is the cgroup Manager Object responsible for managing all
	// pod cgroups.
	cgroupManager CgroupManager
	// Maximum number of pids in a pod
	podPidsLimit int64
	// enforceCPULimits controls whether cfs quota is enforced or not
	enforceCPULimits bool
	// cpuCFSQuotaPeriod is the cfs period value, cfs_period_us, setting per
	// node for all containers in usec
	cpuCFSQuotaPeriod uint64
}

// Make sure that podContainerManagerImpl implements the PodContainerManager interface
var _ PodContainerManager = &podContainerManagerImpl{}

// applyLimits sets pod cgroup resource limits
// It also updates the resource limits on top level qos containers.
func (m *podContainerManagerImpl) applyLimits(pod *v1.Pod) error {
	// This function will house the logic for setting the resource parameters
	// on the pod container config and updating top level qos container configs
	return nil
}

// Exists checks if the pod's cgroup already exists
func (m *podContainerManagerImpl) Exists(pod *v1.Pod) bool {
	podContainerName, _ := m.GetPodContainerName(pod)
	return m.cgroupManager.Exists(podContainerName)
}

// EnsureExists takes a pod as argument and makes sure that
// pod cgroup exists if qos cgroup hierarchy flag is enabled.
// If the pod level container doesn't already exist it is created.
func (m *podContainerManagerImpl) EnsureExists(pod *v1.Pod) error {
	podContainerName, _ := m.GetPodContainerName(pod)
	// check if container already exist
	alreadyExists := m.Exists(pod)
	if !alreadyExists {
		// Create the pod container
		containerConfig := &CgroupConfig{
			Name:               podContainerName,
			ResourceParameters: ResourceConfigForPod(pod, m.enforceCPULimits, m.cpuCFSQuotaPeriod),
		}
		if utilfeature.DefaultFeatureGate.Enabled(kubefeatures.SupportPodPidsLimit) && m.podPidsLimit > 0 {
			containerConfig.ResourceParameters.PidsLimit = &m.podPidsLimit
		}
		if err := m.cgroupManager.Create(containerConfig); err != nil {
			return fmt.Errorf("failed to create container for %v : %v", podContainerName, err)
		}
	}
	// Apply appropriate resource limits on the pod container
	// Top level qos containers limits are not updated
	// until we figure how to maintain the desired state in the kubelet.
	// Because maintaining the desired state is difficult without checkpointing.
	if err := m.applyLimits(pod); err != nil {
		return fmt.Errorf("failed to apply resource limits on container for %v : %v", podContainerName, err)
	}
	return nil
}

// GetPodContainerName returns the CgroupName identifier, and its literal cgroupfs form on the host.
func (m *podContainerManagerImpl) GetPodContainerName(pod *v1.Pod) (CgroupName, string) {
	podQOS := v1qos.GetPodQOS(pod)
	// Get the parent QOS container name
	var parentContainer CgroupName
	switch podQOS {
	case v1.PodQOSGuaranteed:
		parentContainer = m.qosContainersInfo.Guaranteed
	case v1.PodQOSBurstable:
		parentContainer = m.qosContainersInfo.Burstable
	case v1.PodQOSBestEffort:
		parentContainer = m.qosContainersInfo.BestEffort
	}
	podContainer := GetPodCgroupNameSuffix(pod.UID)

	// Get the absolute path of the cgroup
	cgroupName := NewCgroupName(parentContainer, podContainer)
	// Get the literal cgroupfs name
	cgroupfsName := m.cgroupManager.Name(cgroupName)

	return cgroupName, cgroupfsName
}

// Kill one process ID
func (m *podContainerManagerImpl) killOnePid(pid int) error {
	// os.FindProcess never returns an error on POSIX
	// https://go-review.googlesource.com/c/go/+/19093
	p, _ := os.FindProcess(pid)
	if err := p.Kill(); err != nil {
		// If the process already exited, that's fine.
		if strings.Contains(err.Error(), "process already finished") {
			// Hate parsing strings, but
			// vendor/github.com/opencontainers/runc/libcontainer/
			// also does this.
			klog.V(3).Infof("process with pid %v no longer exists", pid)
			return nil
		}
		return err
	}
	return nil
}

// Scan through the whole cgroup directory and kill all processes either
// attached to the pod cgroup or to a container cgroup under the pod cgroup
func (m *podContainerManagerImpl) tryKillingCgroupProcesses(podCgroup CgroupName) error {
	pidsToKill := m.cgroupManager.Pids(podCgroup)
	// No pids charged to the terminated pod cgroup return
	if len(pidsToKill) == 0 {
		return nil
	}

	var errlist []error
	// os.Kill often errors out,
	// We try killing all the pids multiple times
	removed := map[int]bool{}
	for i := 0; i < 5; i++ {
		if i != 0 {
			klog.V(3).Infof("Attempt %v failed to kill all unwanted process. Retyring", i)
		}
		errlist = []error{}
		for _, pid := range pidsToKill {
			if _, ok := removed[pid]; ok {
				continue
			}
			klog.V(3).Infof("Attempt to kill process with pid: %v", pid)
			if err := m.killOnePid(pid); err != nil {
				klog.V(3).Infof("failed to kill process with pid: %v", pid)
				errlist = append(errlist, err)
			} else {
				removed[pid] = true
			}
		}
		if len(errlist) == 0 {
			klog.V(3).Infof("successfully killed all unwanted processes.")
			return nil
		}
	}
	return utilerrors.NewAggregate(errlist)
}

// Destroy destroys the pod container cgroup paths
func (m *podContainerManagerImpl) Destroy(podCgroup CgroupName) error {
	// Try killing all the processes attached to the pod cgroup
	if err := m.tryKillingCgroupProcesses(podCgroup); err != nil {
		klog.V(3).Infof("failed to kill all the processes attached to the %v cgroups", podCgroup)
		return fmt.Errorf("failed to kill all the processes attached to the %v cgroups : %v", podCgroup, err)
	}

	// Now its safe to remove the pod's cgroup
	containerConfig := &CgroupConfig{
		Name:               podCgroup,
		ResourceParameters: &ResourceConfig{},
	}
	if err := m.cgroupManager.Destroy(containerConfig); err != nil {
		return fmt.Errorf("failed to delete cgroup paths for %v : %v", podCgroup, err)
	}
	return nil
}

// ReduceCPULimits reduces the CPU CFS values to the minimum amount of shares.
func (m *podContainerManagerImpl) ReduceCPULimits(podCgroup CgroupName) error {
	return m.cgroupManager.ReduceCPULimits(podCgroup)
}

// IsPodCgroup returns true if the literal cgroupfs name corresponds to a pod
func (m *podContainerManagerImpl) IsPodCgroup(cgroupfs string) (bool, types.UID) {
	// convert the literal cgroupfs form to the driver specific value
	cgroupName := m.cgroupManager.CgroupName(cgroupfs)
	qosContainersList := [3]CgroupName{m.qosContainersInfo.BestEffort, m.qosContainersInfo.Burstable, m.qosContainersInfo.Guaranteed}
	basePath := ""
	for _, qosContainerName := range qosContainersList {
		// a pod cgroup is a direct child of a qos node, so check if its a match
		if len(cgroupName) == len(qosContainerName)+1 {
			basePath = cgroupName[len(qosContainerName)]
		}
	}
	if basePath == "" {
		return false, types.UID("")
	}
	if !strings.HasPrefix(basePath, podCgroupNamePrefix) {
		return false, types.UID("")
	}
	parts := strings.Split(basePath, podCgroupNamePrefix)
	if len(parts) != 2 {
		return false, types.UID("")
	}
	return true, types.UID(parts[1])
}

// GetAllPodsFromCgroups scans through all the subsystems of pod cgroups
// Get list of pods whose cgroup still exist on the cgroup mounts
func (m *podContainerManagerImpl) GetAllPodsFromCgroups() (map[types.UID]CgroupName, error) {
	// Map for storing all the found pods on the disk
	foundPods := make(map[types.UID]CgroupName)
	qosContainersList := [3]CgroupName{m.qosContainersInfo.BestEffort, m.qosContainersInfo.Burstable, m.qosContainersInfo.Guaranteed}
	// Scan through all the subsystem mounts
	// and through each QoS cgroup directory for each subsystem mount
	// If a pod cgroup exists in even a single subsystem mount
	// we will attempt to delete it
	for _, val := range m.subsystems.MountPoints {
		for _, qosContainerName := range qosContainersList {
			// get the subsystems QoS cgroup absolute name
			qcConversion := m.cgroupManager.Name(qosContainerName)
			qc := path.Join(val, qcConversion)
			dirInfo, err := ioutil.ReadDir(qc)
			if err != nil {
				if os.IsNotExist(err) {
					continue
				}
				return nil, fmt.Errorf("failed to read the cgroup directory %v : %v", qc, err)
			}
			for i := range dirInfo {
				// its not a directory, so continue on...
				if !dirInfo[i].IsDir() {
					continue
				}
				// convert the concrete cgroupfs name back to an internal identifier
				// this is needed to handle path conversion for systemd environments.
				// we pass the fully qualified path so decoding can work as expected
				// since systemd encodes the path in each segment.
				cgroupfsPath := path.Join(qcConversion, dirInfo[i].Name())
				internalPath := m.cgroupManager.CgroupName(cgroupfsPath)
				// we only care about base segment of the converted path since that
				// is what we are reading currently to know if it is a pod or not.
				basePath := internalPath[len(internalPath)-1]
				if !strings.Contains(basePath, podCgroupNamePrefix) {
					continue
				}
				// we then split the name on the pod prefix to determine the uid
				parts := strings.Split(basePath, podCgroupNamePrefix)
				// the uid is missing, so we log the unexpected cgroup not of form pod<uid>
				if len(parts) != 2 {
					klog.Errorf("pod cgroup manager ignoring unexpected cgroup %v because it is not a pod", cgroupfsPath)
					continue
				}
				podUID := parts[1]
				foundPods[types.UID(podUID)] = internalPath
			}
		}
	}
	return foundPods, nil
}

// podContainerManagerNoop implements podContainerManager interface.
// It is a no-op implementation and basically does nothing
// podContainerManagerNoop is used in case the QoS cgroup Hierarchy is not
// enabled, so Exists() returns true always as the cgroupRoot
// is expected to always exist.
type podContainerManagerNoop struct {
	cgroupRoot CgroupName
}

// Make sure that podContainerManagerStub implements the PodContainerManager interface
var _ PodContainerManager = &podContainerManagerNoop{}

func (m *podContainerManagerNoop) Exists(_ *v1.Pod) bool {
	return true
}

func (m *podContainerManagerNoop) EnsureExists(_ *v1.Pod) error {
	return nil
}

func (m *podContainerManagerNoop) GetPodContainerName(_ *v1.Pod) (CgroupName, string) {
	return m.cgroupRoot, ""
}

func (m *podContainerManagerNoop) GetPodContainerNameForDriver(_ *v1.Pod) string {
	return ""
}

// Destroy destroys the pod container cgroup paths
func (m *podContainerManagerNoop) Destroy(_ CgroupName) error {
	return nil
}

func (m *podContainerManagerNoop) ReduceCPULimits(_ CgroupName) error {
	return nil
}

func (m *podContainerManagerNoop) GetAllPodsFromCgroups() (map[types.UID]CgroupName, error) {
	return nil, nil
}

func (m *podContainerManagerNoop) IsPodCgroup(cgroupfs string) (bool, types.UID) {
	return false, types.UID("")
}
<commit_msg>Add podCgroup to process kill events to allow for correlation<commit_after>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package cm

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/types"
	utilerrors "k8s.io/apimachinery/pkg/util/errors"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
	"k8s.io/klog"
	v1qos "k8s.io/kubernetes/pkg/apis/core/v1/helper/qos"
	kubefeatures "k8s.io/kubernetes/pkg/features"
)

const (
	podCgroupNamePrefix = "pod"
)

// podContainerManagerImpl implements podContainerManager interface.
// It is the general implementation which allows pod level container
// management if qos Cgroup is enabled.
type podContainerManagerImpl struct {
	// qosContainersInfo hold absolute paths of the top level qos containers
	qosContainersInfo QOSContainersInfo
	// Stores the mounted cgroup subsystems
	subsystems *CgroupSubsystems
	// cgroupManager is the cgroup Manager Object responsible for managing all
	// pod cgroups.
	cgroupManager CgroupManager
	// Maximum number of pids in a pod
	podPidsLimit int64
	// enforceCPULimits controls whether cfs quota is enforced or not
	enforceCPULimits bool
	// cpuCFSQuotaPeriod is the cfs period value, cfs_period_us, setting per
	// node for all containers in usec
	cpuCFSQuotaPeriod uint64
}

// Make sure that podContainerManagerImpl implements the PodContainerManager interface
var _ PodContainerManager = &podContainerManagerImpl{}

// applyLimits sets pod cgroup resource limits
// It also updates the resource limits on top level qos containers.
func (m *podContainerManagerImpl) applyLimits(pod *v1.Pod) error {
	// This function will house the logic for setting the resource parameters
	// on the pod container config and updating top level qos container configs
	return nil
}

// Exists checks if the pod's cgroup already exists
func (m *podContainerManagerImpl) Exists(pod *v1.Pod) bool {
	podContainerName, _ := m.GetPodContainerName(pod)
	return m.cgroupManager.Exists(podContainerName)
}

// EnsureExists takes a pod as argument and makes sure that
// pod cgroup exists if qos cgroup hierarchy flag is enabled.
// If the pod level container doesn't already exist it is created.
func (m *podContainerManagerImpl) EnsureExists(pod *v1.Pod) error {
	podContainerName, _ := m.GetPodContainerName(pod)
	// check if container already exist
	alreadyExists := m.Exists(pod)
	if !alreadyExists {
		// Create the pod container
		containerConfig := &CgroupConfig{
			Name:               podContainerName,
			ResourceParameters: ResourceConfigForPod(pod, m.enforceCPULimits, m.cpuCFSQuotaPeriod),
		}
		if utilfeature.DefaultFeatureGate.Enabled(kubefeatures.SupportPodPidsLimit) && m.podPidsLimit > 0 {
			containerConfig.ResourceParameters.PidsLimit = &m.podPidsLimit
		}
		if err := m.cgroupManager.Create(containerConfig); err != nil {
			return fmt.Errorf("failed to create container for %v : %v", podContainerName, err)
		}
	}
	// Apply appropriate resource limits on the pod container
	// Top level qos containers limits are not updated
	// until we figure how to maintain the desired state in the kubelet.
	// Because maintaining the desired state is difficult without checkpointing.
	if err := m.applyLimits(pod); err != nil {
		return fmt.Errorf("failed to apply resource limits on container for %v : %v", podContainerName, err)
	}
	return nil
}

// GetPodContainerName returns the CgroupName identifier, and its literal cgroupfs form on the host.
func (m *podContainerManagerImpl) GetPodContainerName(pod *v1.Pod) (CgroupName, string) {
	podQOS := v1qos.GetPodQOS(pod)
	// Get the parent QOS container name
	var parentContainer CgroupName
	switch podQOS {
	case v1.PodQOSGuaranteed:
		parentContainer = m.qosContainersInfo.Guaranteed
	case v1.PodQOSBurstable:
		parentContainer = m.qosContainersInfo.Burstable
	case v1.PodQOSBestEffort:
		parentContainer = m.qosContainersInfo.BestEffort
	}
	podContainer := GetPodCgroupNameSuffix(pod.UID)

	// Get the absolute path of the cgroup
	cgroupName := NewCgroupName(parentContainer, podContainer)
	// Get the literal cgroupfs name
	cgroupfsName := m.cgroupManager.Name(cgroupName)

	return cgroupName, cgroupfsName
}

// Kill one process ID
func (m *podContainerManagerImpl) killOnePid(pid int) error {
	// os.FindProcess never returns an error on POSIX
	// https://go-review.googlesource.com/c/go/+/19093
	p, _ := os.FindProcess(pid)
	if err := p.Kill(); err != nil {
		// If the process already exited, that's fine.
		if strings.Contains(err.Error(), "process already finished") {
			// Hate parsing strings, but
			// vendor/github.com/opencontainers/runc/libcontainer/
			// also does this.
			klog.V(3).Infof("process with pid %v no longer exists", pid)
			return nil
		}
		return err
	}
	return nil
}

// Scan through the whole cgroup directory and kill all processes either
// attached to the pod cgroup or to a container cgroup under the pod cgroup
func (m *podContainerManagerImpl) tryKillingCgroupProcesses(podCgroup CgroupName) error {
	pidsToKill := m.cgroupManager.Pids(podCgroup)
	// No pids charged to the terminated pod cgroup return
	if len(pidsToKill) == 0 {
		return nil
	}

	var errlist []error
	// os.Kill often errors out,
	// We try killing all the pids multiple times
	removed := map[int]bool{}
	for i := 0; i < 5; i++ {
		if i != 0 {
			klog.V(3).Infof("Attempt %v failed to kill all unwanted process from cgroup: %v. Retyring", i, podCgroup)
		}
		errlist = []error{}
		for _, pid := range pidsToKill {
			if _, ok := removed[pid]; ok {
				continue
			}
			klog.V(3).Infof("Attempt to kill process with pid: %v from cgroup: %v", pid, podCgroup)
			if err := m.killOnePid(pid); err != nil {
				klog.V(3).Infof("failed to kill process with pid: %v from cgroup: %v", pid, podCgroup)
				errlist = append(errlist, err)
			} else {
				removed[pid] = true
			}
		}
		if len(errlist) == 0 {
			klog.V(3).Infof("successfully killed all unwanted processes from cgroup: %v", podCgroup)
			return nil
		}
	}
	return utilerrors.NewAggregate(errlist)
}

// Destroy destroys the pod container cgroup paths
func (m *podContainerManagerImpl) Destroy(podCgroup CgroupName) error {
	// Try killing all the processes attached to the pod cgroup
	if err := m.tryKillingCgroupProcesses(podCgroup); err != nil {
		klog.V(3).Infof("failed to kill all the processes attached to the %v cgroups", podCgroup)
		return fmt.Errorf("failed to kill all the processes attached to the %v cgroups : %v", podCgroup, err)
	}

	// Now its safe to remove the pod's cgroup
	containerConfig := &CgroupConfig{
		Name:               podCgroup,
		ResourceParameters: &ResourceConfig{},
	}
	if err := m.cgroupManager.Destroy(containerConfig); err != nil {
		return fmt.Errorf("failed to delete cgroup paths for %v : %v", podCgroup, err)
	}
	return nil
}

// ReduceCPULimits reduces the CPU CFS values to the minimum amount of shares.
func (m *podContainerManagerImpl) ReduceCPULimits(podCgroup CgroupName) error {
	return m.cgroupManager.ReduceCPULimits(podCgroup)
}

// IsPodCgroup returns true if the literal cgroupfs name corresponds to a pod
func (m *podContainerManagerImpl) IsPodCgroup(cgroupfs string) (bool, types.UID) {
	// convert the literal cgroupfs form to the driver specific value
	cgroupName := m.cgroupManager.CgroupName(cgroupfs)
	qosContainersList := [3]CgroupName{m.qosContainersInfo.BestEffort, m.qosContainersInfo.Burstable, m.qosContainersInfo.Guaranteed}
	basePath := ""
	for _, qosContainerName := range qosContainersList {
		// a pod cgroup is a direct child of a qos node, so check if its a match
		if len(cgroupName) == len(qosContainerName)+1 {
			basePath = cgroupName[len(qosContainerName)]
		}
	}
	if basePath == "" {
		return false, types.UID("")
	}
	if !strings.HasPrefix(basePath, podCgroupNamePrefix) {
		return false, types.UID("")
	}
	parts := strings.Split(basePath, podCgroupNamePrefix)
	if len(parts) != 2 {
		return false, types.UID("")
	}
	return true, types.UID(parts[1])
}

// GetAllPodsFromCgroups scans through all the subsystems of pod cgroups
// Get list of pods whose cgroup still exist on the cgroup mounts
func (m *podContainerManagerImpl) GetAllPodsFromCgroups() (map[types.UID]CgroupName, error) {
	// Map for storing all the found pods on the disk
	foundPods := make(map[types.UID]CgroupName)
	qosContainersList := [3]CgroupName{m.qosContainersInfo.BestEffort, m.qosContainersInfo.Burstable, m.qosContainersInfo.Guaranteed}
	// Scan through all the subsystem mounts
	// and through each QoS cgroup directory for each subsystem mount
	// If a pod cgroup exists in even a single subsystem mount
	// we will attempt to delete it
	for _, val := range m.subsystems.MountPoints {
		for _, qosContainerName := range qosContainersList {
			// get the subsystems QoS cgroup absolute name
			qcConversion := m.cgroupManager.Name(qosContainerName)
			qc := path.Join(val, qcConversion)
			dirInfo, err := ioutil.ReadDir(qc)
			if err != nil {
				if os.IsNotExist(err) {
					continue
				}
				return nil, fmt.Errorf("failed to read the cgroup directory %v : %v", qc, err)
			}
			for i := range dirInfo {
				// its not a directory, so continue on...
				if !dirInfo[i].IsDir() {
					continue
				}
				// convert the concrete cgroupfs name back to an internal identifier
				// this is needed to handle path conversion for systemd environments.
				// we pass the fully qualified path so decoding can work as expected
				// since systemd encodes the path in each segment.
				cgroupfsPath := path.Join(qcConversion, dirInfo[i].Name())
				internalPath := m.cgroupManager.CgroupName(cgroupfsPath)
				// we only care about base segment of the converted path since that
				// is what we are reading currently to know if it is a pod or not.
				basePath := internalPath[len(internalPath)-1]
				if !strings.Contains(basePath, podCgroupNamePrefix) {
					continue
				}
				// we then split the name on the pod prefix to determine the uid
				parts := strings.Split(basePath, podCgroupNamePrefix)
				// the uid is missing, so we log the unexpected cgroup not of form pod<uid>
				if len(parts) != 2 {
					klog.Errorf("pod cgroup manager ignoring unexpected cgroup %v because it is not a pod", cgroupfsPath)
					continue
				}
				podUID := parts[1]
				foundPods[types.UID(podUID)] = internalPath
			}
		}
	}
	return foundPods, nil
}

// podContainerManagerNoop implements podContainerManager interface.
// It is a no-op implementation and basically does nothing
// podContainerManagerNoop is used in case the QoS cgroup Hierarchy is not
// enabled, so Exists() returns true always as the cgroupRoot
// is expected to always exist.
type podContainerManagerNoop struct {
	cgroupRoot CgroupName
}

// Make sure that podContainerManagerStub implements the PodContainerManager interface
var _ PodContainerManager = &podContainerManagerNoop{}

func (m *podContainerManagerNoop) Exists(_ *v1.Pod) bool {
	return true
}

func (m *podContainerManagerNoop) EnsureExists(_ *v1.Pod) error {
	return nil
}

func (m *podContainerManagerNoop) GetPodContainerName(_ *v1.Pod) (CgroupName, string) {
	return m.cgroupRoot, ""
}

func (m *podContainerManagerNoop) GetPodContainerNameForDriver(_ *v1.Pod) string {
	return ""
}

// Destroy destroys the pod container cgroup paths
func (m *podContainerManagerNoop) Destroy(_ CgroupName) error {
	return nil
}

func (m *podContainerManagerNoop) ReduceCPULimits(_ CgroupName) error {
	return nil
}

func (m *podContainerManagerNoop) GetAllPodsFromCgroups() (map[types.UID]CgroupName, error) {
	return nil, nil
}

func (m *podContainerManagerNoop) IsPodCgroup(cgroupfs string) (bool, types.UID) {
	return false, types.UID("")
}
<|endoftext|>


<commit_before>//go:build gofuzz
// +build gofuzz

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

package fuzz

import (
	"bytes"
	"context"
	"os"

	fuzz "github.com/AdaLogics/go-fuzz-headers"
	"github.com/sirupsen/logrus"

	"github.com/containerd/containerd/archive"
	"github.com/containerd/containerd/content/local"
	imageArchive "github.com/containerd/containerd/images/archive"
	"github.com/containerd/containerd/log"
)

// FuzzApply implements a fuzzer that applies
// a fuzzed tar archive on a directory
func FuzzApply(data []byte) int {
	ctx := context.Background()

	// Apply() is logging the below message, which is too noisy and not really useful
	// if the input is random.
	//
	// level=warning msg="ignored xattr ... in archive" error="operation not supported"
	log.G(ctx).Logger.SetLevel(logrus.PanicLevel)

	f := fuzz.NewConsumer(data)
	iters, err := f.GetInt()
	if err != nil {
		return 0
	}
	maxIters := 20
	tmpDir, err := os.MkdirTemp("", "prefix-test")
	if err != nil {
		return 0
	}
	defer os.RemoveAll(tmpDir)
	for i := 0; i < iters%maxIters; i++ {
		rBytes, err := f.TarBytes()
		if err != nil {
			return 0
		}
		r := bytes.NewReader(rBytes)
		_, _ = archive.Apply(ctx, tmpDir, r)
	}
	return 1
}

// FuzzImportIndex implements a fuzzer
// that targets archive.ImportIndex()
func FuzzImportIndex(data []byte) int {
	f := fuzz.NewConsumer(data)
	tarBytes, err := f.TarBytes()
	if err != nil {
		return 0
	}
	ctx := context.Background()
	r := bytes.NewReader(tarBytes)
	tmpdir, err := os.MkdirTemp("", "fuzzing-")
	if err != nil {
		return 0
	}
	cs, err := local.NewStore(tmpdir)
	if err != nil {
		return 0
	}
	_, _ = imageArchive.ImportIndex(ctx, cs, r)
	return 1
}
<commit_msg>fuzzing: improve archive fuzzer<commit_after>//go:build gofuzz
// +build gofuzz

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

package fuzz

import (
	"archive/tar"
	"bytes"
	"context"
	"io"
	"os"
	"path"

	fuzz "github.com/AdaLogics/go-fuzz-headers"
	"github.com/sirupsen/logrus"

	"github.com/containerd/containerd/archive"
	"github.com/containerd/containerd/content/local"
	imageArchive "github.com/containerd/containerd/images/archive"
	"github.com/containerd/containerd/log"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
)

// FuzzApply implements a fuzzer that applies
// a fuzzed tar archive on a directory
func FuzzApply(data []byte) int {
	ctx := context.Background()

	// Apply() is logging the below message, which is too noisy and not really useful
	// if the input is random.
	//
	// level=warning msg="ignored xattr ... in archive" error="operation not supported"
	log.G(ctx).Logger.SetLevel(logrus.PanicLevel)

	f := fuzz.NewConsumer(data)
	iters, err := f.GetInt()
	if err != nil {
		return 0
	}
	maxIters := 20
	tmpDir, err := os.MkdirTemp("", "prefix-test")
	if err != nil {
		return 0
	}
	defer os.RemoveAll(tmpDir)
	for i := 0; i < iters%maxIters; i++ {
		rBytes, err := f.TarBytes()
		if err != nil {
			return 0
		}
		r := bytes.NewReader(rBytes)
		_, _ = archive.Apply(ctx, tmpDir, r)
	}
	return 1
}

// FuzzImportIndex implements a fuzzer
// that targets archive.ImportIndex()
func FuzzImportIndex(data []byte) int {
	f := fuzz.NewConsumer(data)
	tarBytes, err := f.TarBytes()
	if err != nil {
		return 0
	}
	ctx := context.Background()
	r := bytes.NewReader(tarBytes)
	shouldRequireLayoutOrManifest, err := f.GetBool()
	if err != nil {
		return 0
	}
	if shouldRequireLayoutOrManifest {
		tr := tar.NewReader(r)
		for {
			hdr, err := tr.Next()
			if err == io.EOF {
				break
			}
			if err != nil {
				return 0
			}
			hdrName := path.Clean(hdr.Name)
			switch hdrName {
			case ocispec.ImageLayoutFile:
				break
			case "manifest.json":
				break
			default:
				return 0
			}
		}
	}
	tmpdir, err := os.MkdirTemp("", "fuzzing-")
	if err != nil {
		return 0
	}
	cs, err := local.NewStore(tmpdir)
	if err != nil {
		return 0
	}
	_, _ = imageArchive.ImportIndex(ctx, cs, r)
	return 1
}
<|endoftext|>