<commit_before>/*
Copyright 2020 The Ceph-CSI 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 rbd

import (
	"context"
	"errors"
	"fmt"

	"github.com/ceph/ceph-csi/internal/journal"
	"github.com/ceph/ceph-csi/internal/util"

	librbd "github.com/ceph/go-ceph/rbd"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"k8s.io/klog"
)

// checkCloneImage check the cloned image exists, if the cloned image is not
// found it will check the temporary cloned snapshot exists, and again it will
// check the snapshot exists on the temporary cloned image, if yes it will
// create a new cloned and delete the temporary snapshot and adds a task to
// flatten the temp cloned image and return success.
//
// if the temporary snapshot does not exists it creates a temporary snapshot on
// temporary cloned image and creates  a new cloned with user-provided image
// features and delete the temporary snapshot and adds a task to flatten the
// temp cloned image and return success
//
// if the temporary clone does not exist and if there is a temporary snapshot
// present on the parent image it will delete the temporary snapshot and
// returns.
func (rv *rbdVolume) checkCloneImage(ctx context.Context, parentVol *rbdVolume) (bool, error) {
	// generate temp cloned volume
	tempClone := rv.generateTempClone()
	snap := &rbdSnapshot{
		RbdSnapName: rv.RbdImageName,
		Pool:        rv.Pool,
	}
	var einf ErrImageNotFound
	var esnf ErrSnapNotFound
	// check if cloned image exists
	err := rv.getImageInfo()
	if err == nil {
		// check if do we have temporary snapshot on temporary cloned image
		sErr := tempClone.checkSnapExists(snap)
		if sErr != nil {
			if errors.As(err, &esnf) {
				return true, nil
			}
			return false, err
		}
		err = tempClone.deleteSnapshot(ctx, snap)
		if err == nil {
			return true, nil
		}
		return false, err
	}
	if !errors.As(err, &einf) {
		// return error if its not image not found
		return false, err
	}

	err = tempClone.checkSnapExists(snap)
	if err != nil {
		if errors.As(err, &esnf) {
			// check temporary image needs flatten, if yes add task to flatten the
			// temporary clone
			err = tempClone.flattenRbdImage(ctx, rv.conn.Creds, false, rbdHardMaxCloneDepth, rbdSoftMaxCloneDepth)
			if err != nil {
				return false, err
			}
			// as the snapshot is not present, create new snapshot,clone and
			// delete the temporary snapshot
			err = createRBDClone(ctx, tempClone, rv, snap, rv.conn.Creds)
			if err != nil {
				return false, err
			}
			return true, nil
		} else if !errors.As(err, &einf) {
			// any error other than image not found return error
			return false, err
		}
	} else {
		// snap will be create after we flatten the temporary cloned image,no
		// need to check for flatten here.
		// as the snap exists,create clone image and delete temporary snapshot
		// and add task to flatten temporary cloned image
		err = rv.cloneRbdImageFromSnapshot(ctx, snap)
		if err != nil {
			klog.Errorf(util.Log(ctx, "failed to clone rbd image %s from snapshot %s: %v"), rv.RbdImageName, snap.RbdSnapName, err)
			err = fmt.Errorf("failed to clone rbd image %s from snapshot %s: %v", rv.RbdImageName, snap.RbdSnapName, err)
			return false, err
		}
		err = tempClone.deleteSnapshot(ctx, snap)
		if err != nil {
			klog.Errorf(util.Log(ctx, "failed to delete snapshot: %v"), err)
			return false, err
		}
		return true, nil
	}
	// as the temp clone doesnot exists,check snapshot exists on parent volume
	// snapshot name is same as temporary clone image
	snap.RbdImageName = tempClone.RbdImageName
	err = parentVol.checkSnapExists(snap)
	if err == nil {
		// the temp clone exists, delete it lets reserve a new ID and
		// create new resources for clearner approach
		err = parentVol.deleteSnapshot(ctx, snap)
	}
	if errors.As(err, &esnf) {
		return false, nil
	}
	return false, err
}

func (rv *rbdVolume) generateTempClone() *rbdVolume {
	tempClone := rbdVolume{}
	tempClone.conn = rv.conn
	// The temp clone image need to have deep flatten feature
	f := []string{librbd.FeatureNameLayering, librbd.FeatureNameDeepFlatten}
	tempClone.imageFeatureSet = librbd.FeatureSetFromNames(f)
	tempClone.ClusterID = rv.ClusterID
	tempClone.Monitors = rv.Monitors
	tempClone.Pool = rv.Pool
	// The temp cloned image name will be always (rbd image name + "-temp")
	// this name will be always unique, as cephcsi never creates an image with
	// this format for new rbd images
	tempClone.RbdImageName = rv.RbdImageName + "-temp"
	return &tempClone
}

func (rv *rbdVolume) createCloneFromImage(ctx context.Context, parentVol *rbdVolume) error {
	// generate temp cloned volume
	tempClone := rv.generateTempClone()
	tempSnap := &rbdSnapshot{
		// snapshot name is same as temporary cloned image, This helps to
		// flatten the temporary cloned images as we cannot have more than 510
		// snapshots on an rbd image
		RbdSnapName: tempClone.RbdImageName,
		Pool:        rv.Pool,
	}

	cloneSnap := &rbdSnapshot{
		RbdSnapName: rv.RbdImageName,
		Pool:        rv.Pool,
	}
	var (
		errClone   error
		errFlatten error
		err        error
	)
	var efip ErrFlattenInProgress
	var j = &journal.Connection{}

	j, err = volJournal.Connect(rv.Monitors, rv.conn.Creds)
	if err != nil {
		return status.Error(codes.Internal, err.Error())
	}
	defer j.Destroy()

	// create snapshot and temporary clone and delete snapshot
	err = createRBDClone(ctx, parentVol, tempClone, tempSnap, rv.conn.Creds)
	if err != nil {
		return err
	}

	defer func() {
		if err != nil || errFlatten != nil {
			if !errors.Is(errFlatten, &efip) {
				// cleanup snapshot
				cErr := cleanUpSnapshot(ctx, parentVol, tempSnap, tempClone, rv.conn.Creds)
				if cErr != nil {
					klog.Errorf(util.Log(ctx, "failed to cleanup image %s or snapshot %s: %v"), tempSnap, tempClone, cErr)
				}
			}
		}
		if err != nil || errClone != nil {
			cErr := cleanUpSnapshot(ctx, tempClone, cloneSnap, rv, rv.conn.Creds)
			if cErr != nil {
				klog.Errorf(util.Log(ctx, "failed to cleanup image %s or snapshot %s: %v"), cloneSnap, tempClone, cErr)
			}
		}
	}()
	// flatten clone
	errFlatten = tempClone.flattenRbdImage(ctx, rv.conn.Creds, false, rbdHardMaxCloneDepth, rbdSoftMaxCloneDepth)
	if errFlatten != nil {
		return err
	}
	// create snap of temp clone from temporary cloned image
	// create final clone
	// delete snap of temp clone
	errClone = createRBDClone(ctx, tempClone, rv, cloneSnap, rv.conn.Creds)
	if errClone != nil {
		// set errFlatten error to cleanup temporary snapshot and temporary clone
		errFlatten = errors.New("failed to create user requested cloned image")
		return errClone
	}
	err = rv.getImageID()
	if err != nil {
		klog.Errorf(util.Log(ctx, "failed to get volume id %s: %v"), rv, err)
		return err
	}

	err = j.StoreImageID(ctx, rv.JournalPool, rv.ReservedID, rv.ImageID, rv.conn.Creds)
	if err != nil {
		klog.Errorf(util.Log(ctx, "failed to store volume %s: %v"), rv, err)
		return err
	}
	return nil
}

func (rv *rbdVolume) flattenCloneImage(ctx context.Context) error {
	tempClone := rv.generateTempClone()
	// reducing the limit for cloned images to make sure the limit is in range,
	// If the intermediate clone reaches the depth we may need to return ABORT
	// error message as it need to be flatten before continuing, this may leak
	// omap entries and stale temporary snapshots in corner cases, if we reduce
	// the limit and check for the depth of the parent image clain it self we
	// can flatten the parent images before use to avoid the stale omap entries.
	hardLimit := rbdHardMaxCloneDepth
	softLimit := rbdSoftMaxCloneDepth
	if rbdHardMaxCloneDepth < 2 {
		hardLimit = rbdHardMaxCloneDepth - 2
	}
	if rbdSoftMaxCloneDepth < 2 {
		softLimit = rbdSoftMaxCloneDepth - 2
	}
	err := tempClone.getImageInfo()
	if err == nil {
		return tempClone.flattenRbdImage(ctx, tempClone.conn.Creds, false, hardLimit, softLimit)
	}
	if err != nil {
		var einf ErrImageNotFound
		if !errors.As(err, &einf) {
			return err
		}
	}
	return rv.flattenRbdImage(ctx, rv.conn.Creds, false, hardLimit, softLimit)
}
<commit_msg>cleanup: replace klog with v2<commit_after>/*
Copyright 2020 The Ceph-CSI 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 rbd

import (
	"context"
	"errors"
	"fmt"

	"github.com/ceph/ceph-csi/internal/journal"
	"github.com/ceph/ceph-csi/internal/util"

	librbd "github.com/ceph/go-ceph/rbd"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	klog "k8s.io/klog/v2"
)

// checkCloneImage check the cloned image exists, if the cloned image is not
// found it will check the temporary cloned snapshot exists, and again it will
// check the snapshot exists on the temporary cloned image, if yes it will
// create a new cloned and delete the temporary snapshot and adds a task to
// flatten the temp cloned image and return success.
//
// if the temporary snapshot does not exists it creates a temporary snapshot on
// temporary cloned image and creates  a new cloned with user-provided image
// features and delete the temporary snapshot and adds a task to flatten the
// temp cloned image and return success
//
// if the temporary clone does not exist and if there is a temporary snapshot
// present on the parent image it will delete the temporary snapshot and
// returns.
func (rv *rbdVolume) checkCloneImage(ctx context.Context, parentVol *rbdVolume) (bool, error) {
	// generate temp cloned volume
	tempClone := rv.generateTempClone()
	snap := &rbdSnapshot{
		RbdSnapName: rv.RbdImageName,
		Pool:        rv.Pool,
	}
	var einf ErrImageNotFound
	var esnf ErrSnapNotFound
	// check if cloned image exists
	err := rv.getImageInfo()
	if err == nil {
		// check if do we have temporary snapshot on temporary cloned image
		sErr := tempClone.checkSnapExists(snap)
		if sErr != nil {
			if errors.As(err, &esnf) {
				return true, nil
			}
			return false, err
		}
		err = tempClone.deleteSnapshot(ctx, snap)
		if err == nil {
			return true, nil
		}
		return false, err
	}
	if !errors.As(err, &einf) {
		// return error if its not image not found
		return false, err
	}

	err = tempClone.checkSnapExists(snap)
	if err != nil {
		if errors.As(err, &esnf) {
			// check temporary image needs flatten, if yes add task to flatten the
			// temporary clone
			err = tempClone.flattenRbdImage(ctx, rv.conn.Creds, false, rbdHardMaxCloneDepth, rbdSoftMaxCloneDepth)
			if err != nil {
				return false, err
			}
			// as the snapshot is not present, create new snapshot,clone and
			// delete the temporary snapshot
			err = createRBDClone(ctx, tempClone, rv, snap, rv.conn.Creds)
			if err != nil {
				return false, err
			}
			return true, nil
		} else if !errors.As(err, &einf) {
			// any error other than image not found return error
			return false, err
		}
	} else {
		// snap will be create after we flatten the temporary cloned image,no
		// need to check for flatten here.
		// as the snap exists,create clone image and delete temporary snapshot
		// and add task to flatten temporary cloned image
		err = rv.cloneRbdImageFromSnapshot(ctx, snap)
		if err != nil {
			klog.Errorf(util.Log(ctx, "failed to clone rbd image %s from snapshot %s: %v"), rv.RbdImageName, snap.RbdSnapName, err)
			err = fmt.Errorf("failed to clone rbd image %s from snapshot %s: %v", rv.RbdImageName, snap.RbdSnapName, err)
			return false, err
		}
		err = tempClone.deleteSnapshot(ctx, snap)
		if err != nil {
			klog.Errorf(util.Log(ctx, "failed to delete snapshot: %v"), err)
			return false, err
		}
		return true, nil
	}
	// as the temp clone doesnot exists,check snapshot exists on parent volume
	// snapshot name is same as temporary clone image
	snap.RbdImageName = tempClone.RbdImageName
	err = parentVol.checkSnapExists(snap)
	if err == nil {
		// the temp clone exists, delete it lets reserve a new ID and
		// create new resources for clearner approach
		err = parentVol.deleteSnapshot(ctx, snap)
	}
	if errors.As(err, &esnf) {
		return false, nil
	}
	return false, err
}

func (rv *rbdVolume) generateTempClone() *rbdVolume {
	tempClone := rbdVolume{}
	tempClone.conn = rv.conn
	// The temp clone image need to have deep flatten feature
	f := []string{librbd.FeatureNameLayering, librbd.FeatureNameDeepFlatten}
	tempClone.imageFeatureSet = librbd.FeatureSetFromNames(f)
	tempClone.ClusterID = rv.ClusterID
	tempClone.Monitors = rv.Monitors
	tempClone.Pool = rv.Pool
	// The temp cloned image name will be always (rbd image name + "-temp")
	// this name will be always unique, as cephcsi never creates an image with
	// this format for new rbd images
	tempClone.RbdImageName = rv.RbdImageName + "-temp"
	return &tempClone
}

func (rv *rbdVolume) createCloneFromImage(ctx context.Context, parentVol *rbdVolume) error {
	// generate temp cloned volume
	tempClone := rv.generateTempClone()
	tempSnap := &rbdSnapshot{
		// snapshot name is same as temporary cloned image, This helps to
		// flatten the temporary cloned images as we cannot have more than 510
		// snapshots on an rbd image
		RbdSnapName: tempClone.RbdImageName,
		Pool:        rv.Pool,
	}

	cloneSnap := &rbdSnapshot{
		RbdSnapName: rv.RbdImageName,
		Pool:        rv.Pool,
	}
	var (
		errClone   error
		errFlatten error
		err        error
	)
	var efip ErrFlattenInProgress
	var j = &journal.Connection{}

	j, err = volJournal.Connect(rv.Monitors, rv.conn.Creds)
	if err != nil {
		return status.Error(codes.Internal, err.Error())
	}
	defer j.Destroy()

	// create snapshot and temporary clone and delete snapshot
	err = createRBDClone(ctx, parentVol, tempClone, tempSnap, rv.conn.Creds)
	if err != nil {
		return err
	}

	defer func() {
		if err != nil || errFlatten != nil {
			if !errors.Is(errFlatten, &efip) {
				// cleanup snapshot
				cErr := cleanUpSnapshot(ctx, parentVol, tempSnap, tempClone, rv.conn.Creds)
				if cErr != nil {
					klog.Errorf(util.Log(ctx, "failed to cleanup image %s or snapshot %s: %v"), tempSnap, tempClone, cErr)
				}
			}
		}
		if err != nil || errClone != nil {
			cErr := cleanUpSnapshot(ctx, tempClone, cloneSnap, rv, rv.conn.Creds)
			if cErr != nil {
				klog.Errorf(util.Log(ctx, "failed to cleanup image %s or snapshot %s: %v"), cloneSnap, tempClone, cErr)
			}
		}
	}()
	// flatten clone
	errFlatten = tempClone.flattenRbdImage(ctx, rv.conn.Creds, false, rbdHardMaxCloneDepth, rbdSoftMaxCloneDepth)
	if errFlatten != nil {
		return err
	}
	// create snap of temp clone from temporary cloned image
	// create final clone
	// delete snap of temp clone
	errClone = createRBDClone(ctx, tempClone, rv, cloneSnap, rv.conn.Creds)
	if errClone != nil {
		// set errFlatten error to cleanup temporary snapshot and temporary clone
		errFlatten = errors.New("failed to create user requested cloned image")
		return errClone
	}
	err = rv.getImageID()
	if err != nil {
		klog.Errorf(util.Log(ctx, "failed to get volume id %s: %v"), rv, err)
		return err
	}

	err = j.StoreImageID(ctx, rv.JournalPool, rv.ReservedID, rv.ImageID, rv.conn.Creds)
	if err != nil {
		klog.Errorf(util.Log(ctx, "failed to store volume %s: %v"), rv, err)
		return err
	}
	return nil
}

func (rv *rbdVolume) flattenCloneImage(ctx context.Context) error {
	tempClone := rv.generateTempClone()
	// reducing the limit for cloned images to make sure the limit is in range,
	// If the intermediate clone reaches the depth we may need to return ABORT
	// error message as it need to be flatten before continuing, this may leak
	// omap entries and stale temporary snapshots in corner cases, if we reduce
	// the limit and check for the depth of the parent image clain it self we
	// can flatten the parent images before use to avoid the stale omap entries.
	hardLimit := rbdHardMaxCloneDepth
	softLimit := rbdSoftMaxCloneDepth
	if rbdHardMaxCloneDepth < 2 {
		hardLimit = rbdHardMaxCloneDepth - 2
	}
	if rbdSoftMaxCloneDepth < 2 {
		softLimit = rbdSoftMaxCloneDepth - 2
	}
	err := tempClone.getImageInfo()
	if err == nil {
		return tempClone.flattenRbdImage(ctx, tempClone.conn.Creds, false, hardLimit, softLimit)
	}
	if err != nil {
		var einf ErrImageNotFound
		if !errors.As(err, &einf) {
			return err
		}
	}
	return rv.flattenRbdImage(ctx, rv.conn.Creds, false, hardLimit, softLimit)
}
<|endoftext|>


<commit_before>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 save

import (
	"fmt"
	"io"
	"os"
	"regexp"

	"golang.org/x/net/context"

	"github.com/jacobsa/comeback/internal/graph"
	"github.com/jacobsa/syncutil"
)

// Given a base directory and a set of exclusions, list all file system nodes
// involved, filling in only the RelPath and Parent fields. Output nodes are
// guaranteed to be in reverse topologically sorted order: children appear
// before their parents.
func listNodes(
	ctx context.Context,
	basePath string,
	exclusions []*regexp.Regexp,
	nodes chan<- *fsNode) (err error) {
	b := syncutil.NewBundle(ctx)

	// Find the nodes in the appropriate order, writing them to a channel of
	// graph.Node.
	graphNodes := make(chan graph.Node, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(graphNodes)

		// Set up a root node.
		rootNode := &fsNode{
			RelPath: "",
			Parent:  nil,
		}

		rootNode.Info, err = os.Lstat(basePath)
		if err != nil {
			err = fmt.Errorf("os.Lstat: %v", err)
			return
		}

		if !rootNode.Info.IsDir() {
			err = fmt.Errorf("Not a directory: %q", basePath)
			return
		}

		// Explore the graph.
		sf := newSuccessorFinder(basePath, exclusions)
		err = graph.ReverseTopsortTree(
			ctx,
			sf,
			rootNode,
			graphNodes)

		if err != nil {
			err = fmt.Errorf("ReverseTopsortTree: %v", err)
			return
		}

		return
	})

	// Convert to *fsNode.
	b.Add(func(ctx context.Context) (err error) {
		for graphNode := range graphNodes {
			n, ok := graphNode.(*fsNode)
			if !ok {
				err = fmt.Errorf("Unexpected node type: %T", graphNode)
				return
			}

			select {
			case nodes <- n:
			case <-ctx.Done():
				err = ctx.Err()
				return
			}
		}

		return
	})

	err = b.Join()
	return
}

// Given a base directory and a set of exclusions, list the files and
// directories that would be saved by a backup job with the same info in a
// human-readable format. Write the output to the supplied writer.
func List(
	ctx context.Context,
	w io.Writer,
	basePath string,
	exclusions []*regexp.Regexp) (err error) {
	// TODO(jacobsa): Make this function use listNodes.

	b := syncutil.NewBundle(context.Background())

	// Explore the file system graph, writing all non-excluded nodes into a
	// channel.
	graphNodes := make(chan graph.Node, 100)
	b.Add(func(ctx context.Context) (err error) {
		// Set up a root node.
		rootNode := &fsNode{
			RelPath: "",
			Parent:  nil,
		}

		rootNode.Info, err = os.Lstat(basePath)
		if err != nil {
			err = fmt.Errorf("os.Lstat: %v", err)
			return
		}

		if !rootNode.Info.IsDir() {
			err = fmt.Errorf("Not a directory: %q", basePath)
			return
		}

		// Explore the graph.
		defer close(graphNodes)
		sf := newSuccessorFinder(basePath, exclusions)

		const parallelism = 8
		err = graph.ExploreDirectedGraph(
			ctx,
			sf,
			[]graph.Node{rootNode},
			graphNodes,
			parallelism)

		if err != nil {
			err = fmt.Errorf("ExploreDirectedGraph: %v", err)
			return
		}

		return
	})

	// Print out info about each node.
	b.Add(func(ctx context.Context) (err error) {
		for graphNode := range graphNodes {
			n, ok := graphNode.(*fsNode)
			if !ok {
				err = fmt.Errorf("Unexpected node type: %T", n)
				return
			}

			// Skip the root node.
			_, err = fmt.Fprintf(w, "%q %d\n", n.RelPath, n.Info.Size())
			if err != nil {
				err = fmt.Errorf("Fprintf: %v", err)
				return
			}
		}

		return
	})

	err = b.Join()
	return
}
<commit_msg>Updated save.List.<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 save

import (
	"fmt"
	"io"
	"os"
	"regexp"

	"golang.org/x/net/context"

	"github.com/jacobsa/comeback/internal/graph"
	"github.com/jacobsa/syncutil"
)

// Given a base directory and a set of exclusions, list all file system nodes
// involved, filling in only the RelPath and Parent fields. Output nodes are
// guaranteed to be in reverse topologically sorted order: children appear
// before their parents.
func listNodes(
	ctx context.Context,
	basePath string,
	exclusions []*regexp.Regexp,
	nodes chan<- *fsNode) (err error) {
	b := syncutil.NewBundle(ctx)

	// Find the nodes in the appropriate order, writing them to a channel of
	// graph.Node.
	graphNodes := make(chan graph.Node, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(graphNodes)

		// Set up a root node.
		rootNode := &fsNode{
			RelPath: "",
			Parent:  nil,
		}

		rootNode.Info, err = os.Lstat(basePath)
		if err != nil {
			err = fmt.Errorf("os.Lstat: %v", err)
			return
		}

		if !rootNode.Info.IsDir() {
			err = fmt.Errorf("Not a directory: %q", basePath)
			return
		}

		// Explore the graph.
		sf := newSuccessorFinder(basePath, exclusions)
		err = graph.ReverseTopsortTree(
			ctx,
			sf,
			rootNode,
			graphNodes)

		if err != nil {
			err = fmt.Errorf("ReverseTopsortTree: %v", err)
			return
		}

		return
	})

	// Convert to *fsNode.
	b.Add(func(ctx context.Context) (err error) {
		for graphNode := range graphNodes {
			n, ok := graphNode.(*fsNode)
			if !ok {
				err = fmt.Errorf("Unexpected node type: %T", graphNode)
				return
			}

			select {
			case nodes <- n:
			case <-ctx.Done():
				err = ctx.Err()
				return
			}
		}

		return
	})

	err = b.Join()
	return
}

// Given a base directory and a set of exclusions, list the files and
// directories that would be saved by a backup job with the same info in a
// human-readable format. Write the output to the supplied writer.
func List(
	ctx context.Context,
	w io.Writer,
	basePath string,
	exclusions []*regexp.Regexp) (err error) {
	b := syncutil.NewBundle(ctx)

	// List nodes.
	nodes := make(chan *fsNode, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(nodes)
		err = listNodes(ctx, basePath, exclusions, nodes)
		if err != nil {
			err = fmt.Errorf("listNodes: %v", err)
			return
		}

		return
	})

	// Fill in stat info.
	statted := make(chan *fsNode, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(statted)
		err = statNodes(ctx, basePath, nodes, statted)
		if err != nil {
			err = fmt.Errorf("statNodes: %v", err)
			return
		}

		return
	})

	// Print out info about each node.
	b.Add(func(ctx context.Context) (err error) {
		for n := range statted {
			_, err = fmt.Fprintf(w, "%q %d\n", n.RelPath, n.Info.Size())
			if err != nil {
				err = fmt.Errorf("Fprintf: %v", err)
				return
			}
		}

		return
	})

	err = b.Join()
	return
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: MIT

package tree

import (
	"fmt"
	"sort"
	"strings"

	"github.com/issue9/errwrap"

	"github.com/issue9/mux/v4/internal/handlers"
	"github.com/issue9/mux/v4/internal/syntax"
	"github.com/issue9/mux/v4/params"
)

// node.children 的数量只有达到此值时，才会为其建立 indexes 索引表。
const indexesSize = 5

// 表示路由中的节点
type node struct {
	parent   *node
	handlers *handlers.Handlers
	children []*node
	segment  *syntax.Segment

	// 保存着 *node 实例在 children 中的下标。
	//
	// 所有节点类型为字符串的子节点，其首字符必定是不同的（相同的都提升到父节点中），
	// 根据此特性，可以将所有字符串类型的首字符做个索引，这样字符串类型节点的比较，
	// 可以通过索引排除不必要的比较操作。
	indexes map[byte]int
}

// 构建当前节点的索引表。
func (n *node) buildIndexes() {
	if len(n.children) < indexesSize {
		n.indexes = nil
		return
	}

	if n.indexes == nil {
		n.indexes = make(map[byte]int, indexesSize)
	}

	for index, node := range n.children {
		if node.segment.Type == syntax.String {
			n.indexes[node.segment.Value[0]] = index
		}
	}
}

// 当前节点的优先级
//
// parent.children 根据此值进行排序。
// 不同的节点类型拥有不同的优先级，相同类型的，则有子节点的优先级低。
func (n *node) priority() int {
	// 目前节点类型只有 3 种，10
	// 可以保证在当前类型的节点进行加权时，不会超过其它节点。
	ret := int(n.segment.Type) * 10

	// 有 children 的，endpoint 必然为 false，两者不可能同时为 true
	if len(n.children) > 0 || n.segment.Endpoint {
		return ret + 1
	}

	return ret
}

// 获取指定路径下的节点，若节点不存在，则添加。
// segments 为被 syntax.Split 拆分之后的字符串数组。
func (n *node) getNode(segments []*syntax.Segment) *node {
	child := n.addSegment(segments[0])

	if len(segments) == 1 { // 最后一个节点
		return child
	}

	return child.getNode(segments[1:])
}

// 将 seg 添加到当前节点，并返回新节点，如果找到相同的节点，则直接返回该子节点。
func (n *node) addSegment(seg *syntax.Segment) *node {
	var child *node // 找到的最匹配节点
	var l int       // 最大的匹配字符数量
	for _, c := range n.children {
		l1 := c.segment.Similarity(seg)

		if l1 == -1 { // 找到完全相同的，则直接返回该节点
			return c
		}

		if l1 > l { // 找到相似度更高的，保存该节点的信息
			l = l1
			child = c
		}
	}

	if l <= 0 { // 没有共同前缀，声明一个新的加入到当前节点
		return n.newChild(seg)
	}

	parent := splitNode(child, l)

	// seg 与 parent 重叠
	if len(seg.Value) == l {
		return parent
	}

	// seg.Value[:l] 与 child.segment.Value[:l] 暨 parent.Value 是相同的
	return parent.addSegment(syntax.NewSegment(seg.Value[l:]))
}

// 根据 s 内容为当前节点产生一个子节点，并返回该新节点。
// 由调用方确保 s 的语法正确性，否则可能 panic。
func (n *node) newChild(s *syntax.Segment) *node {
	child := &node{
		parent:  n,
		segment: s,
	}

	n.children = append(n.children, child)
	sort.SliceStable(n.children, func(i, j int) bool {
		return n.children[i].priority() < n.children[j].priority()
	})
	n.buildIndexes()

	return child
}

// 查找路由项，不存在返回 nil
func (n *node) find(pattern string) *node {
	for _, child := range n.children {
		if child.segment.Value == pattern {
			return child
		}

		if strings.HasPrefix(pattern, child.segment.Value) {
			nn := child.find(pattern[len(child.segment.Value):])
			if nn != nil {
				return nn
			}
		}
	} // end for

	return nil
}

// 清除路由项
func (n *node) clean(prefix string) {
	if len(prefix) == 0 {
		n.children = n.children[:0]
		return
	}

	dels := make([]string, 0, len(n.children))
	for _, child := range n.children {
		if len(child.segment.Value) < len(prefix) {
			if strings.HasPrefix(prefix, child.segment.Value) {
				child.clean(prefix[len(child.segment.Value):])
			}
		}

		if strings.HasPrefix(child.segment.Value, prefix) {
			dels = append(dels, child.segment.Value)
		}
	}

	for _, del := range dels {
		n.children = removeNodes(n.children, del)
	}
	n.buildIndexes()
}

// 从子节点中查找与当前路径匹配的节点，若找不到，则返回 nil。
//
// NOTE: 此函数与 node.trace 是一样的，记得同步两边的代码。
func (n *node) match(path string, params params.Params) *node {
	if len(n.indexes) > 0 && len(path) > 0 { // 普通字符串的匹配
		node := n.children[n.indexes[path[0]]]
		if node == nil {
			goto LOOP
		}

		index := node.segment.Match(path, params)
		if index < 0 {
			goto LOOP
		}

		if nn := node.match(path[index:], params); nn != nil {
			return nn
		}
	}

LOOP:
	// 即使 path 为空，也有可能子节点正好可以匹配空的内容。
	// 比如 /posts/{path:\\w*} 后面的 path 即为空节点。所以此处不判断 len(path)
	for i := len(n.indexes); i < len(n.children); i++ {
		node := n.children[i]

		index := node.segment.Match(path, params)
		if index < 0 {
			continue
		}

		if nn := node.match(path[index:], params); nn != nil {
			return nn
		}

		delete(params, n.segment.Name) // 不匹配，则删除写入的参数
	}

	// 没有子节点匹配，len(path)==0，且子节点不为空，可以判定与当前节点匹配。
	if len(path) == 0 && n.handlers != nil && n.handlers.Len() > 0 {
		return n
	}

	return nil
}

// URL 根据参数生成地址
func (n *node) url(params map[string]string) (string, error) {
	nodes := make([]*node, 0, 5)
	for curr := n; curr.parent != nil; curr = curr.parent { // 从尾部向上开始获取节点
		nodes = append(nodes, curr)
	}

	var buf errwrap.StringBuilder
	for i := len(nodes) - 1; i >= 0; i-- {
		node := nodes[i]
		switch node.segment.Type {
		case syntax.String:
			buf.WString(node.segment.Value)
		case syntax.Named, syntax.Regexp:
			param, exists := params[node.segment.Name]
			if !exists {
				return "", fmt.Errorf("未找到参数 %s 的值", node.segment.Name)
			}
			buf.WString(param).
				WString(node.segment.Suffix) // 如果是 endpoint suffix 肯定为空
		} // end switch
	} // end for

	if buf.Err != nil {
		return "", buf.Err
	}
	return buf.String(), nil
}

// 从 nodes 中删除一个 pattern 字段为指定值的元素，
//
// NOTE: 实际应该中，理论上不会出现多个相同的元素，
// 所以此处不作多余的判断。
func removeNodes(nodes []*node, pattern string) []*node {
	for index, n := range nodes {
		if n.segment.Value == pattern {
			return append(nodes[:index], nodes[index+1:]...)
		}
	}

	return nodes
}

// 将节点 n 从 pos 位置进行拆分。后一段作为当前段的子节点，并返回当前节点。
// 若 pos 大于或等于 n.pattern 的长度，则直接返回 n 不会拆分，pos 处的字符作为子节点的内容。
//
// 若 pos 位置是不可拆分的，或是 n.parent 为 nil，都将触发 panic
func splitNode(n *node, pos int) *node {
	if len(n.segment.Value) <= pos { // 不需要拆分
		return n
	}

	p := n.parent
	if p == nil {
		panic("节点必须要有一个有效的父节点，才能进行拆分")
	}

	// 先从父节点中删除老的 n
	p.children = removeNodes(p.children, n.segment.Value)
	p.buildIndexes()

	segs := n.segment.Split(pos)
	ret := p.newChild(segs[0])
	c := ret.newChild(segs[1])
	c.handlers = n.handlers
	c.children = n.children
	c.indexes = n.indexes
	for _, item := range c.children {
		item.parent = c
	}

	return ret
}

// 获取所有的路由地址列表
func (n *node) all(ignoreHead, ignoreOptions bool, parent string, routes map[string][]string) {
	path := parent + n.segment.Value

	if n.handlers != nil && n.handlers.Len() > 0 {
		routes[path] = n.handlers.Methods(ignoreHead, ignoreOptions)
	}

	for _, v := range n.children {
		v.all(ignoreHead, ignoreOptions, path, routes)
	}
}
<commit_msg>fix(internal/tree): 修正 URL 无法处理 Interceptor 类型的错误<commit_after>// SPDX-License-Identifier: MIT

package tree

import (
	"fmt"
	"sort"
	"strings"

	"github.com/issue9/errwrap"

	"github.com/issue9/mux/v4/internal/handlers"
	"github.com/issue9/mux/v4/internal/syntax"
	"github.com/issue9/mux/v4/params"
)

// node.children 的数量只有达到此值时，才会为其建立 indexes 索引表。
const indexesSize = 5

// 表示路由中的节点
type node struct {
	parent   *node
	handlers *handlers.Handlers
	children []*node
	segment  *syntax.Segment

	// 保存着 *node 实例在 children 中的下标。
	//
	// 所有节点类型为字符串的子节点，其首字符必定是不同的（相同的都提升到父节点中），
	// 根据此特性，可以将所有字符串类型的首字符做个索引，这样字符串类型节点的比较，
	// 可以通过索引排除不必要的比较操作。
	indexes map[byte]int
}

// 构建当前节点的索引表。
func (n *node) buildIndexes() {
	if len(n.children) < indexesSize {
		n.indexes = nil
		return
	}

	if n.indexes == nil {
		n.indexes = make(map[byte]int, indexesSize)
	}

	for index, node := range n.children {
		if node.segment.Type == syntax.String {
			n.indexes[node.segment.Value[0]] = index
		}
	}
}

// 当前节点的优先级
//
// parent.children 根据此值进行排序。
// 不同的节点类型拥有不同的优先级，相同类型的，则有子节点的优先级低。
func (n *node) priority() int {
	// 目前节点类型只有 3 种，10
	// 可以保证在当前类型的节点进行加权时，不会超过其它节点。
	ret := int(n.segment.Type) * 10

	// 有 children 的，endpoint 必然为 false，两者不可能同时为 true
	if len(n.children) > 0 || n.segment.Endpoint {
		return ret + 1
	}

	return ret
}

// 获取指定路径下的节点，若节点不存在，则添加。
// segments 为被 syntax.Split 拆分之后的字符串数组。
func (n *node) getNode(segments []*syntax.Segment) *node {
	child := n.addSegment(segments[0])

	if len(segments) == 1 { // 最后一个节点
		return child
	}

	return child.getNode(segments[1:])
}

// 将 seg 添加到当前节点，并返回新节点，如果找到相同的节点，则直接返回该子节点。
func (n *node) addSegment(seg *syntax.Segment) *node {
	var child *node // 找到的最匹配节点
	var l int       // 最大的匹配字符数量
	for _, c := range n.children {
		l1 := c.segment.Similarity(seg)

		if l1 == -1 { // 找到完全相同的，则直接返回该节点
			return c
		}

		if l1 > l { // 找到相似度更高的，保存该节点的信息
			l = l1
			child = c
		}
	}

	if l <= 0 { // 没有共同前缀，声明一个新的加入到当前节点
		return n.newChild(seg)
	}

	parent := splitNode(child, l)

	// seg 与 parent 重叠
	if len(seg.Value) == l {
		return parent
	}

	// seg.Value[:l] 与 child.segment.Value[:l] 暨 parent.Value 是相同的
	return parent.addSegment(syntax.NewSegment(seg.Value[l:]))
}

// 根据 s 内容为当前节点产生一个子节点，并返回该新节点。
// 由调用方确保 s 的语法正确性，否则可能 panic。
func (n *node) newChild(s *syntax.Segment) *node {
	child := &node{
		parent:  n,
		segment: s,
	}

	n.children = append(n.children, child)
	sort.SliceStable(n.children, func(i, j int) bool {
		return n.children[i].priority() < n.children[j].priority()
	})
	n.buildIndexes()

	return child
}

// 查找路由项，不存在返回 nil
func (n *node) find(pattern string) *node {
	for _, child := range n.children {
		if child.segment.Value == pattern {
			return child
		}

		if strings.HasPrefix(pattern, child.segment.Value) {
			nn := child.find(pattern[len(child.segment.Value):])
			if nn != nil {
				return nn
			}
		}
	} // end for

	return nil
}

// 清除路由项
func (n *node) clean(prefix string) {
	if len(prefix) == 0 {
		n.children = n.children[:0]
		return
	}

	dels := make([]string, 0, len(n.children))
	for _, child := range n.children {
		if len(child.segment.Value) < len(prefix) {
			if strings.HasPrefix(prefix, child.segment.Value) {
				child.clean(prefix[len(child.segment.Value):])
			}
		}

		if strings.HasPrefix(child.segment.Value, prefix) {
			dels = append(dels, child.segment.Value)
		}
	}

	for _, del := range dels {
		n.children = removeNodes(n.children, del)
	}
	n.buildIndexes()
}

// 从子节点中查找与当前路径匹配的节点，若找不到，则返回 nil。
//
// NOTE: 此函数与 node.trace 是一样的，记得同步两边的代码。
func (n *node) match(path string, params params.Params) *node {
	if len(n.indexes) > 0 && len(path) > 0 { // 普通字符串的匹配
		node := n.children[n.indexes[path[0]]]
		if node == nil {
			goto LOOP
		}

		index := node.segment.Match(path, params)
		if index < 0 {
			goto LOOP
		}

		if nn := node.match(path[index:], params); nn != nil {
			return nn
		}
	}

LOOP:
	// 即使 path 为空，也有可能子节点正好可以匹配空的内容。
	// 比如 /posts/{path:\\w*} 后面的 path 即为空节点。所以此处不判断 len(path)
	for i := len(n.indexes); i < len(n.children); i++ {
		node := n.children[i]

		index := node.segment.Match(path, params)
		if index < 0 {
			continue
		}

		if nn := node.match(path[index:], params); nn != nil {
			return nn
		}

		delete(params, n.segment.Name) // 不匹配，则删除写入的参数
	}

	// 没有子节点匹配，len(path)==0，且子节点不为空，可以判定与当前节点匹配。
	if len(path) == 0 && n.handlers != nil && n.handlers.Len() > 0 {
		return n
	}

	return nil
}

// URL 根据参数生成地址
func (n *node) url(params map[string]string) (string, error) {
	nodes := make([]*node, 0, 5)
	for curr := n; curr.parent != nil; curr = curr.parent { // 从尾部向上开始获取节点
		nodes = append(nodes, curr)
	}

	var buf errwrap.StringBuilder
	for i := len(nodes) - 1; i >= 0; i-- {
		node := nodes[i]
		switch node.segment.Type {
		case syntax.String:
			buf.WString(node.segment.Value)
		case syntax.Named, syntax.Regexp, syntax.Interceptor:
			param, exists := params[node.segment.Name]
			if !exists {
				return "", fmt.Errorf("未找到参数 %s 的值", node.segment.Name)
			}
			buf.WString(param).
				WString(node.segment.Suffix) // 如果是 endpoint suffix 肯定为空
		} // end switch
	} // end for

	if buf.Err != nil {
		return "", buf.Err
	}
	return buf.String(), nil
}

// 从 nodes 中删除一个 pattern 字段为指定值的元素，
//
// NOTE: 实际应该中，理论上不会出现多个相同的元素，
// 所以此处不作多余的判断。
func removeNodes(nodes []*node, pattern string) []*node {
	for index, n := range nodes {
		if n.segment.Value == pattern {
			return append(nodes[:index], nodes[index+1:]...)
		}
	}

	return nodes
}

// 将节点 n 从 pos 位置进行拆分。后一段作为当前段的子节点，并返回当前节点。
// 若 pos 大于或等于 n.pattern 的长度，则直接返回 n 不会拆分，pos 处的字符作为子节点的内容。
//
// 若 pos 位置是不可拆分的，或是 n.parent 为 nil，都将触发 panic
func splitNode(n *node, pos int) *node {
	if len(n.segment.Value) <= pos { // 不需要拆分
		return n
	}

	p := n.parent
	if p == nil {
		panic("节点必须要有一个有效的父节点，才能进行拆分")
	}

	// 先从父节点中删除老的 n
	p.children = removeNodes(p.children, n.segment.Value)
	p.buildIndexes()

	segs := n.segment.Split(pos)
	ret := p.newChild(segs[0])
	c := ret.newChild(segs[1])
	c.handlers = n.handlers
	c.children = n.children
	c.indexes = n.indexes
	for _, item := range c.children {
		item.parent = c
	}

	return ret
}

// 获取所有的路由地址列表
func (n *node) all(ignoreHead, ignoreOptions bool, parent string, routes map[string][]string) {
	path := parent + n.segment.Value

	if n.handlers != nil && n.handlers.Len() > 0 {
		routes[path] = n.handlers.Methods(ignoreHead, ignoreOptions)
	}

	for _, v := range n.children {
		v.all(ignoreHead, ignoreOptions, path, routes)
	}
}
<|endoftext|>


<commit_before>package uvm

import (
	"fmt"

	"github.com/Microsoft/hcsshim/internal/guestrequest"
	"github.com/Microsoft/hcsshim/internal/logfields"
	"github.com/Microsoft/hcsshim/internal/requesttype"
	"github.com/Microsoft/hcsshim/internal/schema2"
	"github.com/sirupsen/logrus"
)

// AddPlan9 adds a Plan9 share to a utility VM. Each Plan9 share is ref-counted and
// only added if it isn't already.
func (uvm *UtilityVM) AddPlan9(hostPath string, uvmPath string, readOnly bool) error {
	logrus.WithFields(logrus.Fields{
		logfields.UVMID: uvm.id,
		"host-path":     hostPath,
		"uvm-path":      uvmPath,
		"readOnly":      readOnly,
	}).Debug("uvm::AddPlan9")

	if uvm.operatingSystem != "linux" {
		return errNotSupported
	}
	if uvmPath == "" {
		return fmt.Errorf("uvmPath must be passed to AddPlan9")
	}

	// TODO: JTERRY75 - These are marked private in the schema. For now use them
	// but when there are public variants we need to switch to them.
	const (
		shareFlagsReadOnly      int32 = 0x00000001
		shareFlagsLinuxMetadata int32 = 0x00000004
		shareFlagsCaseSensitive int32 = 0x00000008
	)

	flags := shareFlagsLinuxMetadata | shareFlagsCaseSensitive
	if readOnly {
		flags |= shareFlagsReadOnly
	}

	uvm.m.Lock()
	defer uvm.m.Unlock()
	if uvm.plan9Shares == nil {
		uvm.plan9Shares = make(map[string]*plan9Info)
	}
	if _, ok := uvm.plan9Shares[hostPath]; !ok {
		uvm.plan9Counter++

		modification := &hcsschema.ModifySettingRequest{
			RequestType: requesttype.Add,
			Settings: hcsschema.Plan9Share{
				Name:  fmt.Sprintf("%d", uvm.plan9Counter),
				Path:  hostPath,
				Port:  int32(uvm.plan9Counter), // TODO: Temporary. Will all use a single port (9999)
				Flags: flags,
			},
			ResourcePath: fmt.Sprintf("VirtualMachine/Devices/Plan9/Shares"),
			GuestRequest: guestrequest.GuestRequest{
				ResourceType: guestrequest.ResourceTypeMappedDirectory,
				RequestType:  requesttype.Add,
				Settings: guestrequest.LCOWMappedDirectory{
					MountPath: uvmPath,
					Port:      int32(uvm.plan9Counter), // TODO: Temporary. Will all use a single port (9999)
					ReadOnly:  readOnly,
				},
			},
		}

		if err := uvm.Modify(modification); err != nil {
			return err
		}
		uvm.plan9Shares[hostPath] = &plan9Info{
			refCount:  1,
			uvmPath:   uvmPath,
			idCounter: uvm.plan9Counter,
			port:      int32(uvm.plan9Counter), // TODO: Temporary. Will all use a single port (9999)
		}
	} else {
		uvm.plan9Shares[hostPath].refCount++
	}
	logrus.Debugf("hcsshim::AddPlan9 Success %s: refcount=%d %+v", hostPath, uvm.plan9Shares[hostPath].refCount, uvm.plan9Shares[hostPath])
	return nil
}

// RemovePlan9 removes a Plan9 share from a utility VM. Each Plan9 share is ref-counted
// and only actually removed when the ref-count drops to zero.
func (uvm *UtilityVM) RemovePlan9(hostPath string) error {
	if uvm.operatingSystem != "linux" {
		return errNotSupported
	}
	logrus.Debugf("uvm::RemovePlan9 %s id:%s", hostPath, uvm.id)
	uvm.m.Lock()
	defer uvm.m.Unlock()
	if _, ok := uvm.plan9Shares[hostPath]; !ok {
		return fmt.Errorf("%s is not present as a Plan9 share in %s, cannot remove", hostPath, uvm.id)
	}
	return uvm.removePlan9(hostPath, uvm.plan9Shares[hostPath].uvmPath)
}

// removePlan9 is the internally callable "unsafe" version of RemovePlan9. The mutex
// MUST be held when calling this function.
func (uvm *UtilityVM) removePlan9(hostPath, uvmPath string) error {
	uvm.plan9Shares[hostPath].refCount--
	if uvm.plan9Shares[hostPath].refCount > 0 {
		logrus.Debugf("uvm::RemovePlan9 Success %s id:%s Ref-count now %d. It is still present in the utility VM", hostPath, uvm.id, uvm.plan9Shares[hostPath].refCount)
		return nil
	}
	logrus.Debugf("uvm::RemovePlan9 Zero ref-count, removing. %s id:%s", hostPath, uvm.id)
	modification := &hcsschema.ModifySettingRequest{
		RequestType: requesttype.Remove,
		Settings: hcsschema.Plan9Share{
			Name: fmt.Sprintf("%d", uvm.plan9Shares[hostPath].idCounter),
			Port: uvm.plan9Shares[hostPath].port,
		},
		ResourcePath: fmt.Sprintf("VirtualMachine/Devices/Plan9/Shares"),
		GuestRequest: guestrequest.GuestRequest{
			ResourceType: guestrequest.ResourceTypeMappedDirectory,
			RequestType:  requesttype.Remove,
			Settings: guestrequest.LCOWMappedDirectory{
				MountPath: uvm.plan9Shares[hostPath].uvmPath,
				Port:      uvm.plan9Shares[hostPath].port,
			},
		},
	}
	if err := uvm.Modify(modification); err != nil {
		return fmt.Errorf("failed to remove plan9 share %s from %s: %+v: %s", hostPath, uvm.id, modification, err)
	}
	delete(uvm.plan9Shares, hostPath)
	logrus.Debugf("uvm::RemovePlan9 Success %s id:%s successfully removed from utility VM", hostPath, uvm.id)
	return nil
}
<commit_msg>Fix bug in Plan9 case sensitivity<commit_after>package uvm

import (
	"fmt"

	"github.com/Microsoft/hcsshim/internal/guestrequest"
	"github.com/Microsoft/hcsshim/internal/logfields"
	"github.com/Microsoft/hcsshim/internal/requesttype"
	"github.com/Microsoft/hcsshim/internal/schema2"
	"github.com/sirupsen/logrus"
)

// AddPlan9 adds a Plan9 share to a utility VM. Each Plan9 share is ref-counted and
// only added if it isn't already.
func (uvm *UtilityVM) AddPlan9(hostPath string, uvmPath string, readOnly bool) error {
	logrus.WithFields(logrus.Fields{
		logfields.UVMID: uvm.id,
		"host-path":     hostPath,
		"uvm-path":      uvmPath,
		"readOnly":      readOnly,
	}).Debug("uvm::AddPlan9")

	if uvm.operatingSystem != "linux" {
		return errNotSupported
	}
	if uvmPath == "" {
		return fmt.Errorf("uvmPath must be passed to AddPlan9")
	}

	// TODO: JTERRY75 - These are marked private in the schema. For now use them
	// but when there are public variants we need to switch to them.
	const (
		shareFlagsReadOnly      int32 = 0x00000001
		shareFlagsLinuxMetadata int32 = 0x00000004
		shareFlagsCaseSensitive int32 = 0x00000008
	)

	// TODO: JTERRY75 - `shareFlagsCaseSensitive` only works if the Windows
	// `hostPath` supports case sensitivity. We need to detect this case before
	// forwarding this flag in all cases.
	flags := shareFlagsLinuxMetadata // | shareFlagsCaseSensitive
	if readOnly {
		flags |= shareFlagsReadOnly
	}

	uvm.m.Lock()
	defer uvm.m.Unlock()
	if uvm.plan9Shares == nil {
		uvm.plan9Shares = make(map[string]*plan9Info)
	}
	if _, ok := uvm.plan9Shares[hostPath]; !ok {
		uvm.plan9Counter++

		modification := &hcsschema.ModifySettingRequest{
			RequestType: requesttype.Add,
			Settings: hcsschema.Plan9Share{
				Name:  fmt.Sprintf("%d", uvm.plan9Counter),
				Path:  hostPath,
				Port:  int32(uvm.plan9Counter), // TODO: Temporary. Will all use a single port (9999)
				Flags: flags,
			},
			ResourcePath: fmt.Sprintf("VirtualMachine/Devices/Plan9/Shares"),
			GuestRequest: guestrequest.GuestRequest{
				ResourceType: guestrequest.ResourceTypeMappedDirectory,
				RequestType:  requesttype.Add,
				Settings: guestrequest.LCOWMappedDirectory{
					MountPath: uvmPath,
					Port:      int32(uvm.plan9Counter), // TODO: Temporary. Will all use a single port (9999)
					ReadOnly:  readOnly,
				},
			},
		}

		if err := uvm.Modify(modification); err != nil {
			return err
		}
		uvm.plan9Shares[hostPath] = &plan9Info{
			refCount:  1,
			uvmPath:   uvmPath,
			idCounter: uvm.plan9Counter,
			port:      int32(uvm.plan9Counter), // TODO: Temporary. Will all use a single port (9999)
		}
	} else {
		uvm.plan9Shares[hostPath].refCount++
	}
	logrus.Debugf("hcsshim::AddPlan9 Success %s: refcount=%d %+v", hostPath, uvm.plan9Shares[hostPath].refCount, uvm.plan9Shares[hostPath])
	return nil
}

// RemovePlan9 removes a Plan9 share from a utility VM. Each Plan9 share is ref-counted
// and only actually removed when the ref-count drops to zero.
func (uvm *UtilityVM) RemovePlan9(hostPath string) error {
	if uvm.operatingSystem != "linux" {
		return errNotSupported
	}
	logrus.Debugf("uvm::RemovePlan9 %s id:%s", hostPath, uvm.id)
	uvm.m.Lock()
	defer uvm.m.Unlock()
	if _, ok := uvm.plan9Shares[hostPath]; !ok {
		return fmt.Errorf("%s is not present as a Plan9 share in %s, cannot remove", hostPath, uvm.id)
	}
	return uvm.removePlan9(hostPath, uvm.plan9Shares[hostPath].uvmPath)
}

// removePlan9 is the internally callable "unsafe" version of RemovePlan9. The mutex
// MUST be held when calling this function.
func (uvm *UtilityVM) removePlan9(hostPath, uvmPath string) error {
	uvm.plan9Shares[hostPath].refCount--
	if uvm.plan9Shares[hostPath].refCount > 0 {
		logrus.Debugf("uvm::RemovePlan9 Success %s id:%s Ref-count now %d. It is still present in the utility VM", hostPath, uvm.id, uvm.plan9Shares[hostPath].refCount)
		return nil
	}
	logrus.Debugf("uvm::RemovePlan9 Zero ref-count, removing. %s id:%s", hostPath, uvm.id)
	modification := &hcsschema.ModifySettingRequest{
		RequestType: requesttype.Remove,
		Settings: hcsschema.Plan9Share{
			Name: fmt.Sprintf("%d", uvm.plan9Shares[hostPath].idCounter),
			Port: uvm.plan9Shares[hostPath].port,
		},
		ResourcePath: fmt.Sprintf("VirtualMachine/Devices/Plan9/Shares"),
		GuestRequest: guestrequest.GuestRequest{
			ResourceType: guestrequest.ResourceTypeMappedDirectory,
			RequestType:  requesttype.Remove,
			Settings: guestrequest.LCOWMappedDirectory{
				MountPath: uvm.plan9Shares[hostPath].uvmPath,
				Port:      uvm.plan9Shares[hostPath].port,
			},
		},
	}
	if err := uvm.Modify(modification); err != nil {
		return fmt.Errorf("failed to remove plan9 share %s from %s: %+v: %s", hostPath, uvm.id, modification, err)
	}
	delete(uvm.plan9Shares, hostPath)
	logrus.Debugf("uvm::RemovePlan9 Success %s id:%s successfully removed from utility VM", hostPath, uvm.id)
	return nil
}
<|endoftext|>


<commit_before>package symmetric

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"fmt"
        "io"
	"io/ioutil"
        "os"
)

var (
	ErrBadKey        = fmt.Errorf("invalid key")
	ErrPadding       = fmt.Errorf("invalid padding")
	ErrRandomFailure = fmt.Errorf("failed to read enough random data")
)

var (
	devRandom   *os.File
	SecureLevel = 0
)

func init() {
	var err error

	devRandom, err = os.Open("/dev/random")
	if err != nil {
		fmt.Fprintf(os.Stderr, "*** failed to open /dev/random")
		fmt.Fprintf(os.Stderr, "    long-term key generation not recommended")
	} else {
		SecureLevel = 1
	}
}

// Random returns a byte slice containing size random bytes.
func Random(size int) (b []byte, err error) {
	b = make([]byte, size)
	n, err := rand.Read(b)
	if err != nil {
		return
	} else if size != n {
		err = ErrRandomFailure
	}
	return
}

// GenerateKey returns a key suitable for AES-128 cryptography.
func GenerateKey() (key []byte, err error) {
	return Random(KeySize)
}

// Generates a long-term (LT) key. It uses the `/dev/random`
// device, and will typically be much slower than GenerateKey.
func GenerateLTKey() (key []byte, err error) {
	if devRandom == nil {
		err = DegradedError
		return
	}
	key = make([]byte, KeySize)

	n, err := io.ReadFull(devRandom, key)
	if err == nil && n != KeySize {
		err = fmt.Errorf("[key] invalid random read size %d", n)
	}
	return
}

// GenerateIV returns an initialisation vector suitable for
// AES-CBC encryption.
func GenerateIV() (iv []byte, err error) {
	return Random(aes.BlockSize)
}

// Implement the standard padding scheme for block ciphers.
func PadBuffer(m []byte) (p []byte, err error) {
	mLen := len(m)

	p = make([]byte, mLen)
	copy(p, m)

	if len(p) != mLen {
		return p, ErrPadding
	}

	padding := aes.BlockSize - mLen%aes.BlockSize

	p = append(p, 0x80)
	for i := 1; i < padding; i++ {
		p = append(p, 0x0)
	}
	return
}

// Unpad data padded with the standard padding scheme.
func UnpadBuffer(p []byte) (m []byte, err error) {
        m = p
	var pLen int
	origLen := len(m)

	for pLen = origLen - 1; pLen >= 0; pLen-- {
		if m[pLen] == 0x80 {
			break
		}

		if m[pLen] != 0x0 || (origLen-pLen) > aes.BlockSize {
			err = ErrPadding
			return
		}
	}
        m = m[:pLen]
	return
}

// Encrypt encrypts a message, prepending the IV to the beginning.
func Encrypt(key []byte, msg []byte) (ct []byte, err error) {
	c, err := aes.NewCipher(key)
	if err != nil {
		return
	}

	iv, err := GenerateIV()
	if err != nil {
		return
	}

	padded, err := PadBuffer(msg)
	if err != nil {
		return
	}

	cbc := cipher.NewCBCEncrypter(c, iv)
	cbc.CryptBlocks(padded, padded) // encrypt in-place
	ct = iv
	ct = append(ct, padded...)

	return
}

var ErrInvalidIV = fmt.Errorf("invalid IV")

// Decrypt takes an encrypted messages and decrypts it.
func Decrypt(key []byte, ct []byte) (msg []byte, err error) {
	c, err := aes.NewCipher(key)
	if err != nil {
		return
	}

	// Make sure we don't touch the original slice.
	tmp_ct := make([]byte, len(ct))
	copy(tmp_ct, ct)
	iv := tmp_ct[:aes.BlockSize]
	if len(iv) != aes.BlockSize {
		return msg, ErrInvalidIV
	}
	msg = tmp_ct[aes.BlockSize:]

	cbc := cipher.NewCBCDecrypter(c, iv)
	cbc.CryptBlocks(msg, msg)
	msg, err = UnpadBuffer(msg)
	return
}

// Read a key from a file
func ReadKeyFromFile(filename string) (key []byte, err error) {
	key, err = ioutil.ReadFile(filename)
	if err != nil {
		return
	} else if len(key) != KeySize {
		err = ErrBadKey
	}
	return
}

// Zeroise wipes out the data in a slice before deleting the array.
func Zeroise(data *[]byte) (n int) {
	dLen := len(*data)

	for n = 0; n < dLen; n++ {
		(*data)[n] = 0x0
	}

	*data = make([]byte, 0)
	return
}

// Scrub writes random data to the variable the given number of
// rounds, then zeroises it.
func Scrub(data []byte, rounds int) (err error) {
	dLen := len(data)

	var n int
	for r := 0; r < rounds; r++ {
		for i := 0; i < dLen; i++ {
			n, err = io.ReadFull(rand.Reader, data)
			if err != nil {
				return
			} else if n != dLen {
				err = fmt.Errorf("[scrub] invalid random read size %d", n)
				return
			}
		}
	}
	if dLen != Zeroise(&data) {
		err = fmt.Errorf("zeroise failed")
	}
	return
}
<commit_msg>fix bad comment<commit_after>package symmetric

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"fmt"
        "io"
	"io/ioutil"
        "os"
)

var (
	ErrBadKey        = fmt.Errorf("invalid key")
	ErrPadding       = fmt.Errorf("invalid padding")
	ErrRandomFailure = fmt.Errorf("failed to read enough random data")
)

var (
	devRandom   *os.File
	SecureLevel = 0
)

func init() {
	var err error

	devRandom, err = os.Open("/dev/random")
	if err != nil {
		fmt.Fprintf(os.Stderr, "*** failed to open /dev/random")
		fmt.Fprintf(os.Stderr, "    long-term key generation not recommended")
	} else {
		SecureLevel = 1
	}
}

// Random returns a byte slice containing size random bytes.
func Random(size int) (b []byte, err error) {
	b = make([]byte, size)
	n, err := rand.Read(b)
	if err != nil {
		return
	} else if size != n {
		err = ErrRandomFailure
	}
	return
}

// GenerateKey returns a key suitable for AES-256 cryptography.
func GenerateKey() (key []byte, err error) {
	return Random(KeySize)
}

// Generates a long-term (LT) key. It uses the `/dev/random`
// device, and will typically be much slower than GenerateKey.
func GenerateLTKey() (key []byte, err error) {
	if devRandom == nil {
		err = DegradedError
		return
	}
	key = make([]byte, KeySize)

	n, err := io.ReadFull(devRandom, key)
	if err == nil && n != KeySize {
		err = fmt.Errorf("[key] invalid random read size %d", n)
	}
	return
}

// GenerateIV returns an initialisation vector suitable for
// AES-CBC encryption.
func GenerateIV() (iv []byte, err error) {
	return Random(aes.BlockSize)
}

// Implement the standard padding scheme for block ciphers.
func PadBuffer(m []byte) (p []byte, err error) {
	mLen := len(m)

	p = make([]byte, mLen)
	copy(p, m)

	if len(p) != mLen {
		return p, ErrPadding
	}

	padding := aes.BlockSize - mLen%aes.BlockSize

	p = append(p, 0x80)
	for i := 1; i < padding; i++ {
		p = append(p, 0x0)
	}
	return
}

// Unpad data padded with the standard padding scheme.
func UnpadBuffer(p []byte) (m []byte, err error) {
        m = p
	var pLen int
	origLen := len(m)

	for pLen = origLen - 1; pLen >= 0; pLen-- {
		if m[pLen] == 0x80 {
			break
		}

		if m[pLen] != 0x0 || (origLen-pLen) > aes.BlockSize {
			err = ErrPadding
			return
		}
	}
        m = m[:pLen]
	return
}

// Encrypt encrypts a message, prepending the IV to the beginning.
func Encrypt(key []byte, msg []byte) (ct []byte, err error) {
	c, err := aes.NewCipher(key)
	if err != nil {
		return
	}

	iv, err := GenerateIV()
	if err != nil {
		return
	}

	padded, err := PadBuffer(msg)
	if err != nil {
		return
	}

	cbc := cipher.NewCBCEncrypter(c, iv)
	cbc.CryptBlocks(padded, padded) // encrypt in-place
	ct = iv
	ct = append(ct, padded...)

	return
}

var ErrInvalidIV = fmt.Errorf("invalid IV")

// Decrypt takes an encrypted messages and decrypts it.
func Decrypt(key []byte, ct []byte) (msg []byte, err error) {
	c, err := aes.NewCipher(key)
	if err != nil {
		return
	}

	// Make sure we don't touch the original slice.
	tmp_ct := make([]byte, len(ct))
	copy(tmp_ct, ct)
	iv := tmp_ct[:aes.BlockSize]
	if len(iv) != aes.BlockSize {
		return msg, ErrInvalidIV
	}
	msg = tmp_ct[aes.BlockSize:]

	cbc := cipher.NewCBCDecrypter(c, iv)
	cbc.CryptBlocks(msg, msg)
	msg, err = UnpadBuffer(msg)
	return
}

// Read a key from a file
func ReadKeyFromFile(filename string) (key []byte, err error) {
	key, err = ioutil.ReadFile(filename)
	if err != nil {
		return
	} else if len(key) != KeySize {
		err = ErrBadKey
	}
	return
}

// Zeroise wipes out the data in a slice before deleting the array.
func Zeroise(data *[]byte) (n int) {
	dLen := len(*data)

	for n = 0; n < dLen; n++ {
		(*data)[n] = 0x0
	}

	*data = make([]byte, 0)
	return
}

// Scrub writes random data to the variable the given number of
// rounds, then zeroises it.
func Scrub(data []byte, rounds int) (err error) {
	dLen := len(data)

	var n int
	for r := 0; r < rounds; r++ {
		for i := 0; i < dLen; i++ {
			n, err = io.ReadFull(rand.Reader, data)
			if err != nil {
				return
			} else if n != dLen {
				err = fmt.Errorf("[scrub] invalid random read size %d", n)
				return
			}
		}
	}
	if dLen != Zeroise(&data) {
		err = fmt.Errorf("zeroise failed")
	}
	return
}
<|endoftext|>


<commit_before>package index

import (
	"bytes"
	"compress/gzip"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"sync"
	"text/tabwriter"

	"koding/klient/machine/index/node"

	"github.com/djherbis/times"
)

// Index stores a virtual working tree state. It recursively records objects in
// a given root path and allows to efficiently detect changes on it.
type Index struct {
	t *node.Tree
}

var (
	_ json.Marshaler   = (*Index)(nil)
	_ json.Unmarshaler = (*Index)(nil)
)

// NewIndex creates the empty index object.
func NewIndex() *Index {
	return &Index{
		t: node.NewTree(),
	}
}

type fileDesc struct {
	path string      // relative path to the file.
	info os.FileInfo // file LStat result.
}

// NewIndexFiles walks the given file tree roted at root and records file
// states to resulting Index object.
func NewIndexFiles(root string) (*Index, error) {
	idx := NewIndex()

	// Start worker pool.
	var wg sync.WaitGroup
	fC := make(chan *fileDesc)
	for i := 0; i < 2*runtime.NumCPU(); i++ {
		wg.Add(1)
		go idx.addEntryWorker(root, &wg, fC)
	}
	defer func() {
		close(fC)
		wg.Wait()
	}()

	// In order to get as much entries as we can we ignore errors.
	walkFn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}

		fC <- &fileDesc{path: path, info: info}
		return nil
	}

	if err := filepath.Walk(root, walkFn); err != nil {
		return nil, err
	}

	return idx, nil
}

// addEntryWorker asynchronously adds new entry to index. Errors are ignored.
func (idx *Index) addEntryWorker(root string, wg *sync.WaitGroup, fC <-chan *fileDesc) {
	defer wg.Done()

	for f := range fC {
		// Get relative file name.
		name, err := filepath.Rel(root, f.path)
		if err != nil {
			continue
		}

		// Set root path to zero value.
		if name == "." {
			name = ""
		}

		idx.t.DoPath(name, node.Insert(node.NewEntryFileInfo(f.info)))
	}
}

// Clone returns a deep copy of called index.
func (idx *Index) Clone() *Index {
	return &Index{
		t: idx.t.DataClone(),
	}
}

// Tree returns index tree.
func (idx *Index) Tree() *node.Tree {
	return idx.t
}

// Merge calls MergeBranch on all nodes pointed by root path.
func (idx *Index) Merge(root string) ChangeSlice {
	return idx.MergeBranch(root, "")
}

// MergeBranch rereads the given file tree rooted at root and merges its
// entries with index state rooted at branch node. The called index is treated
// as remote part. Scanned directory is considered local part. The following
// rules apply:
//
//  1. If file exists in both remote and local, its fileinfos will be compared
//     with remote entry and if they differs, ChangeMetaUpdate will be created
//     with the direction specified by the result of ctime and mtime comparison.
//
//  2. If file exists in remote but not in local, the entry will have
//     EntryPromiseVirtual property, and ChangeMetaAdd from remote to local
//     will be created.
//
//  3. If file exists in local but not in remote, the entry will be created in
//     remote with EntryPromiseAdd property, and ChangeMetaAdd from local to
//     remote will be produced.
//
// All detected changes will be stored in returned Change slice.
// If branch is empty, the comparison is made against root of the index.
func (idx *Index) MergeBranch(root, branch string) (cs ChangeSlice) {
	rootBranch := filepath.Join(root, branch)
	visited := map[string]struct{}{
		rootBranch: {}, // Skip root file.
	}

	idx.t.DoPath(branch, node.WalkPath(func(name string, n *node.Node) {
		if n.IsShadowed() || n.Entry.File.Mode == 0 {
			return
		}

		nameOS := filepath.FromSlash(name)
		nodePath := filepath.Join(rootBranch, nameOS)
		visited[nodePath] = struct{}{}

		info, err := os.Lstat(nodePath)
		if os.IsNotExist(err) {
			// File exists in remote but not in local.
			n.PromiseVirtual()
			cs = append(cs, NewChange(
				filepath.ToSlash(filepath.Join(branch, nameOS)),
				PriorityLow,
				ChangeMetaAdd|ChangeMetaRemote,
			))
			return
		} else if err != nil {
			return
		}

		mode, size, mtime := n.Entry.File.Mode, n.Entry.File.Size, n.Entry.File.MTime
		// File exists in both remote and local. We compare entry mtime with
		// file mtime and atime. Sometimes synced files may have their mtimes
		// set to source atime. That's why this is necessary.
		if (mtime == info.ModTime().UTC().UnixNano() || mtime == atime(info)) &&
			size == info.Size() &&
			mode == info.Mode() {
			// Files are identical. Allow different ctimes.
			n.UnsetPromises()
			return
		}

		// Merge will not report directory updates because this means that file
		// inside directory was added/removed and this file should be reported.
		// Howewer we want to detect permission changes in all files.
		if mode.IsDir() && mode == info.Mode() {
			n.UnsetPromises()
			return
		}

		// Files differ. However the local file is not virtual.
		n.PromiseUpdate()
		cs = append(cs, NewChange(
			filepath.ToSlash(filepath.Join(branch, nameOS)),
			PriorityMedium,
			ChangeMetaUpdate,
		))
	}))

	// Walk over current root path and check its files.
	walkFn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}

		if _, ok := visited[path]; ok {
			return nil
		}

		// File exists in local but not in remote.
		name, err := filepath.Rel(root, path)
		if err != nil {
			return nil
		}
		name = filepath.ToSlash(name)

		// Add files to tree.
		idx.t.DoPath(name, node.Insert(node.NewEntryFileInfo(info)))
		cs = append(cs, NewChange(name, PriorityLow, ChangeMetaAdd))

		return nil
	}

	if err := filepath.Walk(rootBranch, walkFn); err != nil {
		return nil
	}

	// Put sortest paths to the end.
	sort.Sort(sort.Reverse(cs))
	return cs
}

// Sync modifies index according to provided change path. It checks the file
// on the underlying file system and updates its corresponding index entry.
// This function invalidates all promises set in change entry.
func (idx *Index) Sync(root string, c *Change) {
	if c == nil {
		return
	}

	info, err := os.Lstat(filepath.Join(root, filepath.FromSlash(c.Path())))
	idx.t.DoPath(c.Path(), func(n *node.Node) bool {
		if os.IsNotExist(err) {
			// File doesn't exist. Return false to remove it from tree.
			return false
		} else if err != nil {
			// Nothing much we can do here. Return true to not remove the node.
			return true
		}

		if n.IsShadowed() {
			// File exists on disk but not in tree. Set entry and return true.
			n.Entry = node.NewEntryFileInfo(info)
		} else {
			n.Entry.MergeIn(node.NewEntryFileInfo(info))
			n.UnsetPromises()
		}

		return true
	})
}

// naturalMin returns the minimal value of provided arguments but not less than
// one.
func naturalMin(a, b int) (n int) {
	if n = b; a < b {
		n = a
	}

	if n < 1 {
		return 1
	}

	return n
}

// MarshalJSON satisfies json.Marshaler interface. It safely marshals index
// private data to JSON format.
func (idx *Index) MarshalJSON() ([]byte, error) {
	var b bytes.Buffer
	w := gzip.NewWriter(&b)
	if err := json.NewEncoder(w).Encode(idx.t); err != nil {
		w.Close()
		return nil, err
	}
	w.Close()

	return []byte(`"` + base64.StdEncoding.EncodeToString(b.Bytes()) + `"`), nil
}

// UnmarshalJSON satisfies json.Unmarshaler interface. It is used to unmarshal
// data into private index fields.
func (idx *Index) UnmarshalJSON(data []byte) error {
	dst := make([]byte, base64.StdEncoding.DecodedLen(len(data)-2))
	n, err := base64.StdEncoding.Decode(dst, data[1:len(data)-1])
	if err != nil {
		return err
	}

	r, err := gzip.NewReader(bytes.NewReader(dst[:n]))
	if err != nil {
		return err
	}
	defer r.Close()

	if err = json.NewDecoder(r).Decode(&idx.t); err != nil {
		return err
	}

	return nil
}

// DebugString dumps content of the index as a string, suitable for debugging.
func (idx *Index) DebugString() string {
	m := make(map[string]*node.Entry)
	idx.t.DoPath("", node.WalkPath(func(nodePath string, n *node.Node) {
		m[nodePath] = n.Entry
	}))

	paths := make([]string, 0, len(m))
	for path := range m {
		paths = append(paths, path)
	}
	sort.Strings(paths)

	var buf bytes.Buffer
	tw := tabwriter.NewWriter(&buf, 0, 0, 1, ' ', 0)
	for i, path := range paths {
		fmt.Fprintf(tw, "%5d %s\t%v\n", i+1, path, m[path])
	}
	tw.Flush()

	return buf.String()
}

// atime gets file's access time in UNIX Nano format.
func atime(fi os.FileInfo) int64 {
	return times.Get(fi).AccessTime().UnixNano()
}
<commit_msg>klient/machine: unmarshall fallback to old Nodes structure<commit_after>package index

import (
	"bytes"
	"compress/gzip"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"sync"
	"text/tabwriter"

	"koding/klient/machine/index/node"

	"github.com/djherbis/times"
)

// Index stores a virtual working tree state. It recursively records objects in
// a given root path and allows to efficiently detect changes on it.
type Index struct {
	t *node.Tree
}

var (
	_ json.Marshaler   = (*Index)(nil)
	_ json.Unmarshaler = (*Index)(nil)
)

// NewIndex creates the empty index object.
func NewIndex() *Index {
	return &Index{
		t: node.NewTree(),
	}
}

type fileDesc struct {
	path string      // relative path to the file.
	info os.FileInfo // file LStat result.
}

// NewIndexFiles walks the given file tree roted at root and records file
// states to resulting Index object.
func NewIndexFiles(root string) (*Index, error) {
	idx := NewIndex()

	// Start worker pool.
	var wg sync.WaitGroup
	fC := make(chan *fileDesc)
	for i := 0; i < 2*runtime.NumCPU(); i++ {
		wg.Add(1)
		go idx.addEntryWorker(root, &wg, fC)
	}
	defer func() {
		close(fC)
		wg.Wait()
	}()

	// In order to get as much entries as we can we ignore errors.
	walkFn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}

		fC <- &fileDesc{path: path, info: info}
		return nil
	}

	if err := filepath.Walk(root, walkFn); err != nil {
		return nil, err
	}

	return idx, nil
}

// addEntryWorker asynchronously adds new entry to index. Errors are ignored.
func (idx *Index) addEntryWorker(root string, wg *sync.WaitGroup, fC <-chan *fileDesc) {
	defer wg.Done()

	for f := range fC {
		// Get relative file name.
		name, err := filepath.Rel(root, f.path)
		if err != nil {
			continue
		}

		// Set root path to zero value.
		if name == "." {
			name = ""
		}

		idx.t.DoPath(name, node.Insert(node.NewEntryFileInfo(f.info)))
	}
}

// Clone returns a deep copy of called index.
func (idx *Index) Clone() *Index {
	return &Index{
		t: idx.t.DataClone(),
	}
}

// Tree returns index tree.
func (idx *Index) Tree() *node.Tree {
	return idx.t
}

// Merge calls MergeBranch on all nodes pointed by root path.
func (idx *Index) Merge(root string) ChangeSlice {
	return idx.MergeBranch(root, "")
}

// MergeBranch rereads the given file tree rooted at root and merges its
// entries with index state rooted at branch node. The called index is treated
// as remote part. Scanned directory is considered local part. The following
// rules apply:
//
//  1. If file exists in both remote and local, its fileinfos will be compared
//     with remote entry and if they differs, ChangeMetaUpdate will be created
//     with the direction specified by the result of ctime and mtime comparison.
//
//  2. If file exists in remote but not in local, the entry will have
//     EntryPromiseVirtual property, and ChangeMetaAdd from remote to local
//     will be created.
//
//  3. If file exists in local but not in remote, the entry will be created in
//     remote with EntryPromiseAdd property, and ChangeMetaAdd from local to
//     remote will be produced.
//
// All detected changes will be stored in returned Change slice.
// If branch is empty, the comparison is made against root of the index.
func (idx *Index) MergeBranch(root, branch string) (cs ChangeSlice) {
	rootBranch := filepath.Join(root, branch)
	visited := map[string]struct{}{
		rootBranch: {}, // Skip root file.
	}

	idx.t.DoPath(branch, node.WalkPath(func(name string, n *node.Node) {
		if n.IsShadowed() || n.Entry.File.Mode == 0 {
			return
		}

		nameOS := filepath.FromSlash(name)
		nodePath := filepath.Join(rootBranch, nameOS)
		visited[nodePath] = struct{}{}

		info, err := os.Lstat(nodePath)
		if os.IsNotExist(err) {
			// File exists in remote but not in local.
			n.PromiseVirtual()
			cs = append(cs, NewChange(
				filepath.ToSlash(filepath.Join(branch, nameOS)),
				PriorityLow,
				ChangeMetaAdd|ChangeMetaRemote,
			))
			return
		} else if err != nil {
			return
		}

		mode, size, mtime := n.Entry.File.Mode, n.Entry.File.Size, n.Entry.File.MTime
		// File exists in both remote and local. We compare entry mtime with
		// file mtime and atime. Sometimes synced files may have their mtimes
		// set to source atime. That's why this is necessary.
		if (mtime == info.ModTime().UTC().UnixNano() || mtime == atime(info)) &&
			size == info.Size() &&
			mode == info.Mode() {
			// Files are identical. Allow different ctimes.
			n.UnsetPromises()
			return
		}

		// Merge will not report directory updates because this means that file
		// inside directory was added/removed and this file should be reported.
		// Howewer we want to detect permission changes in all files.
		if mode.IsDir() && mode == info.Mode() {
			n.UnsetPromises()
			return
		}

		// Files differ. However the local file is not virtual.
		n.PromiseUpdate()
		cs = append(cs, NewChange(
			filepath.ToSlash(filepath.Join(branch, nameOS)),
			PriorityMedium,
			ChangeMetaUpdate,
		))
	}))

	// Walk over current root path and check its files.
	walkFn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}

		if _, ok := visited[path]; ok {
			return nil
		}

		// File exists in local but not in remote.
		name, err := filepath.Rel(root, path)
		if err != nil {
			return nil
		}
		name = filepath.ToSlash(name)

		// Add files to tree.
		idx.t.DoPath(name, node.Insert(node.NewEntryFileInfo(info)))
		cs = append(cs, NewChange(name, PriorityLow, ChangeMetaAdd))

		return nil
	}

	if err := filepath.Walk(rootBranch, walkFn); err != nil {
		return nil
	}

	// Put sortest paths to the end.
	sort.Sort(sort.Reverse(cs))
	return cs
}

// Sync modifies index according to provided change path. It checks the file
// on the underlying file system and updates its corresponding index entry.
// This function invalidates all promises set in change entry.
func (idx *Index) Sync(root string, c *Change) {
	if c == nil {
		return
	}

	info, err := os.Lstat(filepath.Join(root, filepath.FromSlash(c.Path())))
	idx.t.DoPath(c.Path(), func(n *node.Node) bool {
		if os.IsNotExist(err) {
			// File doesn't exist. Return false to remove it from tree.
			return false
		} else if err != nil {
			// Nothing much we can do here. Return true to not remove the node.
			return true
		}

		if n.IsShadowed() {
			// File exists on disk but not in tree. Set entry and return true.
			n.Entry = node.NewEntryFileInfo(info)
		} else {
			n.Entry.MergeIn(node.NewEntryFileInfo(info))
			n.UnsetPromises()
		}

		return true
	})
}

// naturalMin returns the minimal value of provided arguments but not less than
// one.
func naturalMin(a, b int) (n int) {
	if n = b; a < b {
		n = a
	}

	if n < 1 {
		return 1
	}

	return n
}

// MarshalJSON satisfies json.Marshaler interface. It safely marshals index
// private data to JSON format.
func (idx *Index) MarshalJSON() ([]byte, error) {
	var b bytes.Buffer
	w := gzip.NewWriter(&b)
	if err := json.NewEncoder(w).Encode(idx.t); err != nil {
		w.Close()
		return nil, err
	}
	w.Close()

	return []byte(`"` + base64.StdEncoding.EncodeToString(b.Bytes()) + `"`), nil
}

// UnmarshalJSON satisfies json.Unmarshaler interface. It is used to unmarshal
// data into private index fields.
func (idx *Index) UnmarshalJSON(data []byte) error {
	dst := make([]byte, base64.StdEncoding.DecodedLen(len(data)-2))
	n, err := base64.StdEncoding.Decode(dst, data[1:len(data)-1])
	if err != nil {
		return err
	}

	r, err := gzip.NewReader(bytes.NewReader(dst[:n]))
	if err != nil {
		return err
	}
	defer r.Close()

	if err = json.NewDecoder(r).Decode(&idx.t); err != nil {
		// Try converting to old nodes.
		rn, err := gzip.NewReader(bytes.NewReader(dst[:n]))
		if err != nil {
			return err
		}
		defer rn.Close()

		var root Node
		if err = json.NewDecoder(r).Decode(&root); err == nil {
			idx.t = root.ToTree()
		} else {
			return err
		}
	}

	return nil
}

// DebugString dumps content of the index as a string, suitable for debugging.
func (idx *Index) DebugString() string {
	m := make(map[string]*node.Entry)
	idx.t.DoPath("", node.WalkPath(func(nodePath string, n *node.Node) {
		m[nodePath] = n.Entry
	}))

	paths := make([]string, 0, len(m))
	for path := range m {
		paths = append(paths, path)
	}
	sort.Strings(paths)

	var buf bytes.Buffer
	tw := tabwriter.NewWriter(&buf, 0, 0, 1, ' ', 0)
	for i, path := range paths {
		fmt.Fprintf(tw, "%5d %s\t%v\n", i+1, path, m[path])
	}
	tw.Flush()

	return buf.String()
}

// atime gets file's access time in UNIX Nano format.
func atime(fi os.FileInfo) int64 {
	return times.Get(fi).AccessTime().UnixNano()
}
<|endoftext|>


<commit_before>package models

type ChannelContainer struct {
	Channel             Channel                  `json:"channel"`
	IsParticipant       bool                     `json:"isParticipant"`
	ParticipantCount    int                      `json:"participantCount"`
	ParticipantsPreview []string                 `json:"participantsPreview"`
	LastMessage         *ChannelMessageContainer `json:"lastMessage"`
}

func NewChannelContainer() *ChannelContainer {
	return &ChannelContainer{}
}

func PopulateChannelContainers(channelList []Channel, accountId int64) ([]*ChannelContainer, error) {
	channelContainers := make([]*ChannelContainer, len(channelList))

	var err error
	for i, channel := range channelList {
		channelContainers[i], err = PopulateChannelContainer(channel, accountId)
		if err != nil {
			return nil, err
		}
	}

	return channelContainers, nil
}

func PopulateChannelContainer(channel Channel, accountId int64) (*ChannelContainer, error) {
	cp := NewChannelParticipant()
	cp.ChannelId = channel.Id

	// add participantCount
	participantCount, err := cp.FetchParticipantCount()
	if err != nil {
		return nil, err
	}

	// add participant preview
	cpList, err := cp.ListAccountIds(5)
	if err != nil {
		return nil, err
	}

	// add participation status
	isParticipant, err := cp.IsParticipant(accountId)
	if err != nil {
		return nil, err
	}

	cc := NewChannelContainer()
	cc.Channel = channel
	cc.IsParticipant = isParticipant
	cc.ParticipantCount = participantCount
	participantOldIds, err := AccountOldsIdByIds(cpList)
	if err != nil {
		return nil, err
	}

	cc.ParticipantsPreview = participantOldIds

	// add last message of the channel
	cm, err := channel.FetchLastMessage()
	if err != nil {
		return nil, err
	}

	if cm != nil {
		cmc, err := cm.BuildEmptyMessageContainer()
		if err != nil {
			return nil, err
		}
		cc.LastMessage = cmc
	}

	return cc, nil
}
<commit_msg>Social: refactor function call<commit_after>package models

type ChannelContainer struct {
	Channel             Channel                  `json:"channel"`
	IsParticipant       bool                     `json:"isParticipant"`
	ParticipantCount    int                      `json:"participantCount"`
	ParticipantsPreview []string                 `json:"participantsPreview"`
	LastMessage         *ChannelMessageContainer `json:"lastMessage"`
}

func NewChannelContainer() *ChannelContainer {
	return &ChannelContainer{}
}

func PopulateChannelContainers(channelList []Channel, accountId int64) ([]*ChannelContainer, error) {
	channelContainers := make([]*ChannelContainer, len(channelList))

	var err error
	for i, channel := range channelList {
		channelContainers[i], err = PopulateChannelContainer(channel, accountId)
		if err != nil {
			return nil, err
		}
	}

	return channelContainers, nil
}

func PopulateChannelContainer(channel Channel, accountId int64) (*ChannelContainer, error) {
	cp := NewChannelParticipant()
	cp.ChannelId = channel.Id

	// add participantCount
	participantCount, err := cp.FetchParticipantCount()
	if err != nil {
		return nil, err
	}

	// add participant preview
	cpList, err := cp.ListAccountIds(5)
	if err != nil {
		return nil, err
	}

	// add participation status
	isParticipant, err := cp.IsParticipant(accountId)
	if err != nil {
		return nil, err
	}

	cc := NewChannelContainer()
	cc.Channel = channel
	cc.IsParticipant = isParticipant
	cc.ParticipantCount = participantCount
	participantOldIds, err := FetchAccountOldsIdByIdsFromCache(cpList)
	if err != nil {
		return nil, err
	}

	cc.ParticipantsPreview = participantOldIds

	// add last message of the channel
	cm, err := channel.FetchLastMessage()
	if err != nil {
		return nil, err
	}

	if cm != nil {
		cmc, err := cm.BuildEmptyMessageContainer()
		if err != nil {
			return nil, err
		}
		cc.LastMessage = cmc
	}

	return cc, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/howeyc/fsnotify"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
)

func init() {
	flag.IntVar(&port, "port", 8080, "The port at which to serve http.")
	watched = make(map[string]bool)

	var err error
	watcher, err = fsnotify.NewWatcher()
	if err != nil {
		panic(err)
	}
}

func main() {
	flag.Parse()
	defer watcher.Close()

	go reactToChanges()

	working, err := os.Getwd()
	if err != nil {
		panic(err)
	}
	updateWatch(working)

	_, file, _, _ := runtime.Caller(0)
	here := filepath.Dir(file)
	static := filepath.Join(here, "../client/")
	http.Handle("/", http.FileServer(http.Dir(static)))
	http.HandleFunc("/watch", watchHandler)
	http.HandleFunc("/latest", reportHandler)
	http.ListenAndServe(fmt.Sprintf(":%d", port), nil)
}

func reportHandler(writer http.ResponseWriter, request *http.Request) {
	writer.Header().Set("Content-Type", "application/json")
	writer.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
	writer.Header().Set("Pragma", "no-cache")
	writer.Header().Set("Expires", "0")
	fmt.Fprint(writer, latestOutput)
}

func watchHandler(writer http.ResponseWriter, request *http.Request) {
	if request.Method == "GET" {
		writer.Write([]byte(rootWatch))
		return
	}

	value := request.URL.Query()["root"]
	if len(value) == 0 {
		http.Error(writer, "No 'root' query string parameter included!", http.StatusBadRequest)
		return
	}
	newRoot := value[0]
	if !exists(newRoot) {
		http.Error(writer, "The 'root' value provided is not an existing directory.", http.StatusNotFound)
	} else {
		updateWatch(newRoot)
	}
}

var (
	port         int
	latestOutput string
	rootWatch    string
	watched      map[string]bool
	watcher      *fsnotify.Watcher
)
<commit_msg>Restored the proper order of operations in the main method.<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/howeyc/fsnotify"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
)

func init() {
	flag.IntVar(&port, "port", 8080, "The port at which to serve http.")
	watched = make(map[string]bool)

	var err error
	watcher, err = fsnotify.NewWatcher()
	if err != nil {
		panic(err)
	}
}

func main() {
	flag.Parse()
	defer watcher.Close()

	working, err := os.Getwd()
	if err != nil {
		panic(err)
	}
	updateWatch(working)

	go reactToChanges()

	_, file, _, _ := runtime.Caller(0)
	here := filepath.Dir(file)
	static := filepath.Join(here, "../client/")
	http.Handle("/", http.FileServer(http.Dir(static)))
	http.HandleFunc("/watch", watchHandler)
	http.HandleFunc("/latest", reportHandler)
	http.ListenAndServe(fmt.Sprintf(":%d", port), nil)
}

func reportHandler(writer http.ResponseWriter, request *http.Request) {
	writer.Header().Set("Content-Type", "application/json")
	writer.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
	writer.Header().Set("Pragma", "no-cache")
	writer.Header().Set("Expires", "0")
	fmt.Fprint(writer, latestOutput)
}

func watchHandler(writer http.ResponseWriter, request *http.Request) {
	if request.Method == "GET" {
		writer.Write([]byte(rootWatch))
		return
	}

	value := request.URL.Query()["root"]
	if len(value) == 0 {
		http.Error(writer, "No 'root' query string parameter included!", http.StatusBadRequest)
		return
	}
	newRoot := value[0]
	if !exists(newRoot) {
		http.Error(writer, "The 'root' value provided is not an existing directory.", http.StatusNotFound)
	} else {
		updateWatch(newRoot)
	}
}

var (
	port         int
	latestOutput string
	rootWatch    string
	watched      map[string]bool
	watcher      *fsnotify.Watcher
)
<|endoftext|>


<commit_before>package storage_test

import (
	"bytes"
	"github.com/stretchr/testify/assert"
	"github.com/viant/toolbox"
	"github.com/viant/toolbox/storage"
	"io/ioutil"
	"os"
	"path"
	"testing"
)

func TestStorageService_List(t *testing.T) {
	service := storage.NewService()
	assert.NotNil(t, service)
	fileName, _, _ := toolbox.CallerInfo(2)
	parent, _ := path.Split(fileName)
	baseUrl := "file://" + parent + "/test"

	if toolbox.FileExists(parent + "/test/file3.txt") {
		os.Remove(parent + "/test/file3.txt")
	}
	defer os.Remove(parent + "/test/file3.txt")

	objects, err := service.List(baseUrl)
	assert.Nil(t, err)
	assert.Equal(t, 3, len(objects))
	var objectByUrl = make(map[string]storage.Object)
	for _, object := range objects {
		objectByUrl[object.URL()] = object
	}
	assert.NotNil(t, objectByUrl[baseUrl+"/dir"])
	assert.NotNil(t, objectByUrl[baseUrl+"/file1.txt"])
	assert.NotNil(t, objectByUrl[baseUrl+"/file2.txt"])
	assert.True(t, objectByUrl[baseUrl+"/dir"].IsFolder())
	assert.True(t, objectByUrl[baseUrl+"/file2.txt"].IsContent())

	{
		reader, err := service.Download(objectByUrl[baseUrl+"/file2.txt"])
		assert.Nil(t, err)
		content, err := ioutil.ReadAll(reader)
		assert.Nil(t, err)
		assert.Equal(t, "line1\nline2", string(content))
	}

	var newFileUrl = baseUrl + "/file3.txt"
	err = service.Upload(baseUrl+"/file3.txt", bytes.NewReader([]byte("abc")))
	assert.Nil(t, err)

	exists, err := service.Exists(baseUrl + "/file3.txt")
	assert.Nil(t, err)
	assert.True(t, exists)

	{
		object, err := service.StorageObject(newFileUrl)
		assert.Nil(t, err)
		reader, err := service.Download(object)
		assert.Nil(t, err)
		content, err := ioutil.ReadAll(reader)
		assert.Nil(t, err)
		assert.Equal(t, "abc", string(content))
	}

}

func TestCopy(t *testing.T) {
	service := storage.NewService()
	assert.NotNil(t, service)

	fileName, _, _ := toolbox.CallerInfo(2)
	parent, _ := path.Split(fileName)
	baseUrl := "file://" + parent + "/test"

	sourceURL := path.Join(baseUrl, "source/")
	targetURL := path.Join(baseUrl, "target/")

	err := storage.Copy(service, sourceURL, service, targetURL, nil)
	assert.Nil(t, err)

	expectedFiles := []string{
		path.Join(parent, "test/target/file1.txt"),
		path.Join(parent, "test/target/file2.txt"),
		path.Join(parent, "test/target/dir/file.json"),
		path.Join(parent, "test/target/dir2/subdir/file1.txt"),

	}
	for _, file := range expectedFiles {
		assert.True(t, toolbox.FileExists(file))
		os.Remove(file)
	}
}<commit_msg>patched test<commit_after>package storage_test

import (
	"bytes"
	"github.com/stretchr/testify/assert"
	"github.com/viant/toolbox"
	"github.com/viant/toolbox/storage"
	"io/ioutil"
	"os"
	"path"
	"testing"
)

func TestStorageService_List(t *testing.T) {
	service := storage.NewService()
	assert.NotNil(t, service)
	fileName, _, _ := toolbox.CallerInfo(2)
	parent, _ := path.Split(fileName)
	baseUrl := "file://" + parent + "/test"

	if toolbox.FileExists(parent + "/test/file3.txt") {
		os.Remove(parent + "/test/file3.txt")
	}
	defer os.Remove(parent + "/test/file3.txt")

	objects, err := service.List(baseUrl)
	assert.Nil(t, err)
	assert.Equal(t, 5, len(objects))
	var objectByUrl = make(map[string]storage.Object)
	for _, object := range objects {
		objectByUrl[object.URL()] = object
	}
	assert.NotNil(t, objectByUrl[baseUrl+"/dir"])
	assert.NotNil(t, objectByUrl[baseUrl+"/file1.txt"])
	assert.NotNil(t, objectByUrl[baseUrl+"/file2.txt"])
	assert.True(t, objectByUrl[baseUrl+"/dir"].IsFolder())
	assert.True(t, objectByUrl[baseUrl+"/file2.txt"].IsContent())

	{
		reader, err := service.Download(objectByUrl[baseUrl+"/file2.txt"])
		assert.Nil(t, err)
		content, err := ioutil.ReadAll(reader)
		assert.Nil(t, err)
		assert.Equal(t, "line1\nline2", string(content))
	}

	var newFileUrl = baseUrl + "/file3.txt"
	err = service.Upload(baseUrl+"/file3.txt", bytes.NewReader([]byte("abc")))
	assert.Nil(t, err)

	exists, err := service.Exists(baseUrl + "/file3.txt")
	assert.Nil(t, err)
	assert.True(t, exists)

	{
		object, err := service.StorageObject(newFileUrl)
		assert.Nil(t, err)
		reader, err := service.Download(object)
		assert.Nil(t, err)
		content, err := ioutil.ReadAll(reader)
		assert.Nil(t, err)
		assert.Equal(t, "abc", string(content))
	}

}

func TestCopy(t *testing.T) {
	service := storage.NewService()
	assert.NotNil(t, service)

	fileName, _, _ := toolbox.CallerInfo(2)
	parent, _ := path.Split(fileName)
	baseUrl := "file://" + parent + "/test"

	sourceURL := path.Join(baseUrl, "source/")
	targetURL := path.Join(baseUrl, "target/")

	err := storage.Copy(service, sourceURL, service, targetURL, nil)
	assert.Nil(t, err)

	expectedFiles := []string{
		path.Join(parent, "test/target/file1.txt"),
		path.Join(parent, "test/target/file2.txt"),
		path.Join(parent, "test/target/dir/file.json"),
		path.Join(parent, "test/target/dir2/subdir/file1.txt"),

	}
	for _, file := range expectedFiles {
		assert.True(t, toolbox.FileExists(file))
		os.Remove(file)
	}
}<|endoftext|>


<commit_before>// Package hammer generates ipfix packets
//: ----------------------------------------------------------------------------
//: Copyright (C) 2017 Verizon.  All Rights Reserved.
//: All Rights Reserved
//:
//: file:    hammer.go
//: details: TODO
//: author:  Mehrdad Arshad Rad
//: date:    02/01/2017
//:
//: 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 hammer

import (
	"fmt"
	"net"
	"sync"
	"time"

	"git.edgecastcdn.net/vflow/mirror"
)

type Packet struct {
	payload []byte
	length  int
}

type IPFIX struct {
	conn mirror.Conn
	pool *sync.Pool
	ch   chan Packet
	srcs []net.IP

	MaxRouter   int
	TplInterval time.Duration
	IPFIXAddr   net.IP
	IPFIXPort   int
	RateLimit   int
}

func NewIPFIX() (*IPFIX, error) {

	raddr := net.ParseIP("127.0.0.1")
	conn, err := mirror.NewRawConn(raddr)
	if err != nil {
		return nil, err
	}

	return &IPFIX{
		conn: conn,
		pool: &sync.Pool{
			New: func() interface{} {
				return make([]byte, 1500)
			},
		},
		ch:          make(chan Packet, 1000),
		MaxRouter:   10,
		TplInterval: 10 * time.Second,
	}, nil
}

func (i *IPFIX) Run() {
	var wg sync.WaitGroup

	for j := 1; j < i.MaxRouter; j++ {
		i.srcs = append(i.srcs, net.ParseIP(fmt.Sprintf("192.168.1.%d", j)))
	}

	wg.Add(1)
	go func() {
		var p Packet
		defer wg.Done()
		for {
			p = <-i.ch
			i.conn.Send(p.payload)
			i.pool.Put(p.payload)
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		i.sendData()
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		i.sendTemplate()
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		i.sendTemplateOpt()
	}()

	wg.Wait()
}

func (i *IPFIX) sendData() {
	var data, payload []byte
	ipHLen := mirror.IPv4HLen

	udp := mirror.UDP{55117, 4739, 0, 0}
	udpHdr := udp.Marshal()

	ip := mirror.NewIPv4HeaderTpl(mirror.UDPProto)
	ipHdr := ip.Marshal()

	for {
		for j := 0; j < len(ipfixDataSamples); j++ {
			for _, src := range i.srcs {
				data = ipfixDataSamples[j]
				payload = i.pool.Get().([]byte)

				udp.SetLen(udpHdr, len(data))

				ip.SetAddrs(ipHdr, src, net.ParseIP("127.0.0.1"))

				copy(payload[0:ipHLen], ipHdr)
				copy(payload[ipHLen:ipHLen+8], udpHdr)
				copy(payload[ipHLen+8:], data)

				i.ch <- Packet{
					payload: payload,
					length:  ipHLen + 8 + len(data),
				}
			}
		}
	}
}

func (i *IPFIX) sendTemplate() {
	var data, payload []byte
	c := time.Tick(i.TplInterval)
	ipHLen := mirror.IPv4HLen

	udp := mirror.UDP{55117, 4739, 0, 0}
	udpHdr := udp.Marshal()

	ip := mirror.NewIPv4HeaderTpl(mirror.UDPProto)
	ipHdr := ip.Marshal()

	for range c {
		for j := 0; j < len(ipfixTemplates); j++ {
			for _, src := range i.srcs {
				data = ipfixTemplates[j]
				payload = i.pool.Get().([]byte)

				udp.SetLen(udpHdr, len(data))

				ip.SetAddrs(ipHdr, src, net.ParseIP("127.0.0.1"))

				copy(payload[0:ipHLen], ipHdr)
				copy(payload[ipHLen:ipHLen+8], udpHdr)
				copy(payload[ipHLen+8:], data)

				i.ch <- Packet{
					payload: payload,
					length:  ipHLen + 8 + len(data),
				}
			}
		}
	}
}

func (i *IPFIX) sendTemplateOpt() {
	var data, payload []byte
	c := time.Tick(i.TplInterval)
	ipHLen := mirror.IPv4HLen

	udp := mirror.UDP{55117, 4739, 0, 0}
	udpHdr := udp.Marshal()

	ip := mirror.NewIPv4HeaderTpl(mirror.UDPProto)
	ipHdr := ip.Marshal()

	for range c {
		for j := 0; j < len(ipfixTemplatesOpt); j++ {
			for _, src := range i.srcs {
				data = ipfixTemplates[j]
				payload = i.pool.Get().([]byte)

				udp.SetLen(udpHdr, len(data))

				ip.SetAddrs(ipHdr, src, net.ParseIP("127.0.0.1"))

				copy(payload[0:ipHLen], ipHdr)
				copy(payload[ipHLen:ipHLen+8], udpHdr)
				copy(payload[ipHLen+8:], data)

				i.ch <- Packet{
					payload: payload,
					length:  ipHLen + 8 + len(data),
				}
			}
		}
	}
}
<commit_msg>optimize code<commit_after>// Package hammer generates ipfix packets
//: ----------------------------------------------------------------------------
//: Copyright (C) 2017 Verizon.  All Rights Reserved.
//: All Rights Reserved
//:
//: file:    hammer.go
//: details: TODO
//: author:  Mehrdad Arshad Rad
//: date:    02/01/2017
//:
//: 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 hammer

import (
	"fmt"
	"net"
	"sync"
	"time"

	"git.edgecastcdn.net/vflow/mirror"
)

const (
	dataType = iota
	templateType
	templateOptType
)

type Packet struct {
	payload []byte
	length  int
}

type IPFIX struct {
	conn mirror.Conn
	ch   chan Packet
	srcs []net.IP

	MaxRouter   int
	TplInterval time.Duration
	IPFIXAddr   net.IP
	IPFIXPort   int
	RateLimit   int
}

func NewIPFIX() (*IPFIX, error) {

	raddr := net.ParseIP("127.0.0.1")
	conn, err := mirror.NewRawConn(raddr)
	if err != nil {
		return nil, err
	}

	return &IPFIX{
		conn:        conn,
		ch:          make(chan Packet, 10000),
		MaxRouter:   10,
		TplInterval: 10 * time.Second,
	}, nil
}

func (i *IPFIX) Run() {
	var wg sync.WaitGroup

	for j := 1; j < i.MaxRouter; j++ {
		i.srcs = append(i.srcs, net.ParseIP(fmt.Sprintf("192.168.1.%d", j)))
	}

	wg.Add(1)
	go func() {
		var p Packet
		defer wg.Done()
		for {
			p = <-i.ch
			i.conn.Send(p.payload[:p.length])
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		i.sendTemplate()
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		i.sendTemplateOpt()
	}()

	time.Sleep(1 * time.Second)

	wg.Add(1)
	go func() {
		defer wg.Done()
		i.sendData()
	}()

	wg.Wait()
}

func (i *IPFIX) sendData() {
	packets := i.genPackets(dataType)

	for {
		for j, _ := range packets {
			i.ch <- packets[j]
		}
	}
}

func (i *IPFIX) sendTemplate() {
	c := time.Tick(i.TplInterval)
	packets := i.genPackets(templateType)

	for j, _ := range packets {
		i.ch <- packets[j]
	}

	for range c {
		for j, _ := range packets {
			i.ch <- packets[j]
		}
	}
}

func (i *IPFIX) sendTemplateOpt() {
	c := time.Tick(i.TplInterval)
	packets := i.genPackets(templateOptType)

	for j, _ := range packets {
		i.ch <- packets[j]
	}

	for range c {
		for j, _ := range packets {
			i.ch <- packets[j]
		}
	}
}

func (i *IPFIX) genPackets(typ int) []Packet {
	var (
		packets []Packet
		samples [][]byte
	)

	ipHLen := mirror.IPv4HLen
	udp := mirror.UDP{55117, 4739, 0, 0}
	udpHdr := udp.Marshal()

	ip := mirror.NewIPv4HeaderTpl(mirror.UDPProto)
	ipHdr := ip.Marshal()

	switch typ {
	case dataType:
		samples = ipfixDataSamples
	case templateType:
		samples = ipfixTemplates
	case templateOptType:
		samples = ipfixTemplatesOpt
	}

	for j := 0; j < len(samples); j++ {
		for _, src := range i.srcs {
			data := samples[j]
			payload := make([]byte, 1500)

			udp.SetLen(udpHdr, len(data))

			ip.SetAddrs(ipHdr, src, net.ParseIP("127.0.0.1"))

			copy(payload[0:ipHLen], ipHdr)
			copy(payload[ipHLen:ipHLen+8], udpHdr)
			copy(payload[ipHLen+8:], data)

			packets = append(packets, Packet{
				payload: payload,
				length:  ipHLen + 8 + len(data),
			})

		}
	}

	return packets
}
<|endoftext|>