<commit_before>// Copyright 2019 Authors of Cilium
//
// 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 cmd

import (
	"context"
	"errors"
	"path"
	"time"

	"github.com/cilium/cilium/pkg/allocator"
	"github.com/cilium/cilium/pkg/identity"
	"github.com/cilium/cilium/pkg/identity/cache"
	"github.com/cilium/cilium/pkg/idpool"
	"github.com/cilium/cilium/pkg/k8s"
	"github.com/cilium/cilium/pkg/k8s/identitybackend"
	kvstoreallocator "github.com/cilium/cilium/pkg/kvstore/allocator"
	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"
	"github.com/cilium/cilium/pkg/option"

	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	k8serrors "k8s.io/apimachinery/pkg/api/errors"
)

// opTimeout is the time allowed for each operation to complete. This includes
// listing, allocating and getting identities.
const opTimeout = 30 * time.Second

var migrateIdentityCmd = &cobra.Command{
	Use:   "migrate-identity",
	Short: "Migrate KVStore-backed identities to kubernetes CRD-backed identities",
	Long: `migrate-identity allows migrating to CRD-backed identities while
	minimizing connection interruptions. It will allocate a CRD-backed identity,
	with the same numeric security identity, for each cilium security identity
	defined in the kvstore. When cilium-agents are restarted with
	identity-allocation-mode set to CRD the numeric identities will then be
	equivalent between new instances and not-upgraded ones. In cases where the
	numeric identity is already in-use by a different set of labels, a new
	numeric identity is created.`,
	Run: func(cmd *cobra.Command, args []string) {
		migrateIdentities()
	},
}

// migrateIdentities attempts to mirror the security identities in the kvstore
// into k8s CRD-backed identities. The identities are snapshotted on startup
// and new identities created during migrations will not be seen.
// It is a little odd because it violates the cilium-agent assumption that only
// 1 Backend is active at a time.
// The steps are:
// 1- Connect to the kvstore via a pkg/allocatore.Backend
// 2- Connect to k8s
//   a- Create the ciliumidentity CRD if it is missing.
// 3- Iterate over each identity in the kvstore
//   a- Attempt to allocate the same numeric ID to this key
//   b- Already allocated identies that match ID->key are skipped
//   c- kvstore IDs with conflicting CRDs are allocated with a different ID
//
// NOTE: It is assumed that the migration is from k8s to k8s installations. The
// key labels different when running in non-k8s mode.
func migrateIdentities() {
	// The internal packages log things. Make sure they follow the setup of of
	// the CLI tool.
	logging.DefaultLogger.SetFormatter(log.Formatter)

	// Setup global configuration
	// These are defined in cilium/cmd/kvstore.go
	option.Config.KVStore = kvStore
	option.Config.KVStoreOpt = kvStoreOpts

	// This allows us to initialize a CRD allocator
	option.Config.IdentityAllocationMode = option.IdentityAllocationModeCRD // force CRD mode to make ciliumid

	// Init Identity backends
	initCtx, initCancel := context.WithTimeout(context.Background(), opTimeout)
	kvstoreBackend := initKVStore()

	crdBackend, crdAllocator := initK8s(initCtx)
	initCancel()

	log.Info("Listing identities in kvstore")
	listCtx, listCancel := context.WithTimeout(context.Background(), opTimeout)
	kvstoreIDs, err := getKVStoreIdentities(listCtx, kvstoreBackend)
	if err != nil {
		log.WithError(err).Fatal("Unable to initialize Identity Allocator with CRD backend to allocate identities with already allocated IDs")
	}
	listCancel()

	log.Info("Migrating identities to CRD")
	badKeys := make([]allocator.AllocatorKey, 0)                       // keys that have real errors
	alreadyAllocatedKeys := make(map[idpool.ID]allocator.AllocatorKey) // IDs that are already allocated, maybe with different labels

	for id, key := range kvstoreIDs {
		scopedLog := log.WithFields(logrus.Fields{
			logfields.Identity:       id,
			logfields.IdentityLabels: key.GetKey(),
		})

		ctx, cancel := context.WithTimeout(context.Background(), opTimeout)
		err := crdBackend.AllocateID(ctx, id, key)
		switch {
		case err != nil && k8serrors.IsAlreadyExists(err):
			alreadyAllocatedKeys[id] = key

		case err != nil:
			scopedLog.WithError(err).Error("Cannot allocate CRD ID. This key will be allocated with a new numeric identity")
			badKeys = append(badKeys, key)

		default:
			scopedLog.Info("Migrated identity")
		}
		cancel()
	}

	// Handle IDs that have conflicts. These can be:
	// 1- The same ID -> key (from a previous run). This is a no-op
	// 2- The same ID but with different labels. This is not ideal. A new ID is
	// allocated as a fallback.
	for id, key := range alreadyAllocatedKeys {
		scopedLog := log.WithFields(logrus.Fields{
			logfields.Identity:       id,
			logfields.IdentityLabels: key.GetKey(),
		})

		upstreamKey, err := crdBackend.GetByID(id)
		scopedLog.Debugf("Looking at upstream key with this ID: %+v", upstreamKey)
		switch {
		case err != nil:
			log.WithError(err).Error("ID already allocated but we cannot verify whether it is the same key. It may not be migrated")
			continue

		// nil returns mean the key doesn't exist. This shouldn't happen, but treat
		// it like a mismatch and allocate it. The allocator will find it if it has
		// been re-allocated via master key protection.
		case upstreamKey == nil && err == nil:
			// fallthrough

		case key.GetKey() == upstreamKey.GetKey():
			scopedLog.Info("ID was already allocated to this key. It is already migrated")
			continue
		}

		scopedLog = log.WithFields(logrus.Fields{
			logfields.OldIdentity:    id,
			logfields.IdentityLabels: key.GetKey(),
		})
		scopedLog.Warn("ID is allocated to a different key in CRD. A new ID will be allocated for the this key")

		ctx, cancel := context.WithTimeout(context.Background(), opTimeout)
		defer cancel()
		newID, actuallyAllocated, err := crdAllocator.Allocate(ctx, key)
		switch {
		case err != nil:
			log.WithError(err).Errorf("Cannot allocate new CRD ID for %v", key)
			continue

		case !actuallyAllocated:
			scopedLog.Debug("Expected to allocate ID but this ID->key mapping re-existed")
		}

		log.WithFields(logrus.Fields{
			logfields.OldIdentity:    id,
			logfields.Identity:       newID,
			logfields.IdentityLabels: key.GetKey(),
		}).Info("New ID allocated for key in CRD")
	}
}

// initK8s connects to k8s with a allocator.Backend and an initialized
// allocator.Allocator, using the k8s config passed into the command.
func initK8s(ctx context.Context) (crdBackend allocator.Backend, crdAllocator *allocator.Allocator) {
	log.Info("Setting up kubernetes client")

	k8sClientQPSLimit := viper.GetFloat64(option.K8sClientQPSLimit)
	k8sClientBurst := viper.GetInt(option.K8sClientBurst)

	k8s.Configure(k8sAPIServer, k8sKubeConfigPath, float32(k8sClientQPSLimit), k8sClientBurst)

	if err := k8s.Init(); err != nil {
		log.WithError(err).Fatal("Unable to connect to Kubernetes apiserver")
	}

	// Update CRDs to ensure ciliumIdentity is present
	k8s.RegisterCRDs()

	// Create a CRD Backend
	crdBackend, err := identitybackend.NewCRDBackend(identitybackend.CRDBackendConfiguration{
		NodeName: "cilium-preflight",
		Store:    nil,
		Client:   k8s.CiliumClient(),
		KeyType:  cache.GlobalIdentity{},
	})
	if err != nil {
		log.WithError(err).Fatal("Cannot create CRD identity backend")
	}

	// Create a real allocator with CRD as the backend. This mimics the setup in
	// pkg/allocator/cache
	//
	// FIXME: add options to handle clustermesh with this constructor parameter:
	//    allocator.WithPrefixMask(idpool.ID(option.Config.ClusterID<<identity.ClusterIDShift)))
	minID := idpool.ID(identity.MinimalAllocationIdentity)
	maxID := idpool.ID(identity.MaximumAllocationIdentity)
	crdAllocator, err = allocator.NewAllocator(cache.GlobalIdentity{}, crdBackend,
		allocator.WithMax(maxID), allocator.WithMin(minID))
	if err != nil {
		log.WithError(err).Fatal("Unable to initialize Identity Allocator with CRD backend to allocate identities with already allocated IDs")
	}

	// Wait for the initial sync to complete
	if err := crdAllocator.WaitForInitialSync(ctx); err != nil {
		log.WithError(err).Fatal("Error waiting for k8s identity allocator to sync. No identities have been migrated.")
	}

	return crdBackend, crdAllocator
}

// initKVStore connects to the kvstore with a allocator.Backend, initialised to
// find identities at the default cilium paths.
func initKVStore() (kvstoreBackend allocator.Backend) {
	log.Info("Setting up kvstore client")
	setupKvstore()

	idPath := path.Join(cache.IdentitiesPath, "id")
	kvstoreBackend, err := kvstoreallocator.NewKVStoreBackend(cache.IdentitiesPath, idPath, cache.GlobalIdentity{})
	if err != nil {
		log.WithError(err).Fatal("Cannot create kvstore identity backend")
	}

	return kvstoreBackend
}

// getKVStoreIdentities lists all identities in the kvstore. It will wait for
// the listing to complete.
func getKVStoreIdentities(ctx context.Context, kvstoreBackend allocator.Backend) (identities map[idpool.ID]allocator.AllocatorKey, err error) {
	identities = make(map[idpool.ID]allocator.AllocatorKey)
	stopChan := make(chan struct{})

	kvstoreBackend.ListAndWatch(kvstoreListHandler{
		onAdd: func(id idpool.ID, key allocator.AllocatorKey) {
			log.Debugf("kvstore listed ID: %+v -> %+v", id, key)
			identities[id] = key
		},
		onListDone: func() {
			close(stopChan) // This makes the ListAndWatch exit after the initial listing
		},
	}, stopChan)

	// Wait for the listing to complete
	select {
	case <-stopChan:
		log.Debug("kvstore ID list complete")

	case <-ctx.Done():
		return nil, errors.New("Timeout while listing identities")
	}

	return identities, nil
}

// kvstoreListHandler is a dummy type to receive callbacks from the kvstore subsystem
type kvstoreListHandler struct {
	onAdd      func(id idpool.ID, key allocator.AllocatorKey)
	onListDone func()
}

func (h kvstoreListHandler) OnListDone()                                       { h.onListDone() }
func (h kvstoreListHandler) OnAdd(id idpool.ID, key allocator.AllocatorKey)    { h.onAdd(id, key) }
func (h kvstoreListHandler) OnModify(id idpool.ID, key allocator.AllocatorKey) {}
func (h kvstoreListHandler) OnDelete(id idpool.ID, key allocator.AllocatorKey) {}
<commit_msg>docs: Fix deadlock in cilium preflight on etcd timeout<commit_after>// Copyright 2019 Authors of Cilium
//
// 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 cmd

import (
	"context"
	"errors"
	"path"
	"time"

	"github.com/cilium/cilium/pkg/allocator"
	"github.com/cilium/cilium/pkg/identity"
	"github.com/cilium/cilium/pkg/identity/cache"
	"github.com/cilium/cilium/pkg/idpool"
	"github.com/cilium/cilium/pkg/k8s"
	"github.com/cilium/cilium/pkg/k8s/identitybackend"
	kvstoreallocator "github.com/cilium/cilium/pkg/kvstore/allocator"
	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"
	"github.com/cilium/cilium/pkg/option"

	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	k8serrors "k8s.io/apimachinery/pkg/api/errors"
)

// opTimeout is the time allowed for each operation to complete. This includes
// listing, allocating and getting identities.
const opTimeout = 30 * time.Second

var migrateIdentityCmd = &cobra.Command{
	Use:   "migrate-identity",
	Short: "Migrate KVStore-backed identities to kubernetes CRD-backed identities",
	Long: `migrate-identity allows migrating to CRD-backed identities while
	minimizing connection interruptions. It will allocate a CRD-backed identity,
	with the same numeric security identity, for each cilium security identity
	defined in the kvstore. When cilium-agents are restarted with
	identity-allocation-mode set to CRD the numeric identities will then be
	equivalent between new instances and not-upgraded ones. In cases where the
	numeric identity is already in-use by a different set of labels, a new
	numeric identity is created.`,
	Run: func(cmd *cobra.Command, args []string) {
		migrateIdentities()
	},
}

// migrateIdentities attempts to mirror the security identities in the kvstore
// into k8s CRD-backed identities. The identities are snapshotted on startup
// and new identities created during migrations will not be seen.
// It is a little odd because it violates the cilium-agent assumption that only
// 1 Backend is active at a time.
// The steps are:
// 1- Connect to the kvstore via a pkg/allocatore.Backend
// 2- Connect to k8s
//   a- Create the ciliumidentity CRD if it is missing.
// 3- Iterate over each identity in the kvstore
//   a- Attempt to allocate the same numeric ID to this key
//   b- Already allocated identies that match ID->key are skipped
//   c- kvstore IDs with conflicting CRDs are allocated with a different ID
//
// NOTE: It is assumed that the migration is from k8s to k8s installations. The
// key labels different when running in non-k8s mode.
func migrateIdentities() {
	// The internal packages log things. Make sure they follow the setup of of
	// the CLI tool.
	logging.DefaultLogger.SetFormatter(log.Formatter)

	// Setup global configuration
	// These are defined in cilium/cmd/kvstore.go
	option.Config.KVStore = kvStore
	option.Config.KVStoreOpt = kvStoreOpts

	// This allows us to initialize a CRD allocator
	option.Config.IdentityAllocationMode = option.IdentityAllocationModeCRD // force CRD mode to make ciliumid

	// Init Identity backends
	initCtx, initCancel := context.WithTimeout(context.Background(), opTimeout)
	kvstoreBackend := initKVStore()

	crdBackend, crdAllocator := initK8s(initCtx)
	initCancel()

	log.Info("Listing identities in kvstore")
	listCtx, listCancel := context.WithTimeout(context.Background(), opTimeout)
	kvstoreIDs, err := getKVStoreIdentities(listCtx, kvstoreBackend)
	if err != nil {
		log.WithError(err).Fatal("Unable to initialize Identity Allocator with CRD backend to allocate identities with already allocated IDs")
	}
	listCancel()

	log.Info("Migrating identities to CRD")
	badKeys := make([]allocator.AllocatorKey, 0)                       // keys that have real errors
	alreadyAllocatedKeys := make(map[idpool.ID]allocator.AllocatorKey) // IDs that are already allocated, maybe with different labels

	for id, key := range kvstoreIDs {
		scopedLog := log.WithFields(logrus.Fields{
			logfields.Identity:       id,
			logfields.IdentityLabels: key.GetKey(),
		})

		ctx, cancel := context.WithTimeout(context.Background(), opTimeout)
		err := crdBackend.AllocateID(ctx, id, key)
		switch {
		case err != nil && k8serrors.IsAlreadyExists(err):
			alreadyAllocatedKeys[id] = key

		case err != nil:
			scopedLog.WithError(err).Error("Cannot allocate CRD ID. This key will be allocated with a new numeric identity")
			badKeys = append(badKeys, key)

		default:
			scopedLog.Info("Migrated identity")
		}
		cancel()
	}

	// Handle IDs that have conflicts. These can be:
	// 1- The same ID -> key (from a previous run). This is a no-op
	// 2- The same ID but with different labels. This is not ideal. A new ID is
	// allocated as a fallback.
	for id, key := range alreadyAllocatedKeys {
		scopedLog := log.WithFields(logrus.Fields{
			logfields.Identity:       id,
			logfields.IdentityLabels: key.GetKey(),
		})

		upstreamKey, err := crdBackend.GetByID(id)
		scopedLog.Debugf("Looking at upstream key with this ID: %+v", upstreamKey)
		switch {
		case err != nil:
			log.WithError(err).Error("ID already allocated but we cannot verify whether it is the same key. It may not be migrated")
			continue

		// nil returns mean the key doesn't exist. This shouldn't happen, but treat
		// it like a mismatch and allocate it. The allocator will find it if it has
		// been re-allocated via master key protection.
		case upstreamKey == nil && err == nil:
			// fallthrough

		case key.GetKey() == upstreamKey.GetKey():
			scopedLog.Info("ID was already allocated to this key. It is already migrated")
			continue
		}

		scopedLog = log.WithFields(logrus.Fields{
			logfields.OldIdentity:    id,
			logfields.IdentityLabels: key.GetKey(),
		})
		scopedLog.Warn("ID is allocated to a different key in CRD. A new ID will be allocated for the this key")

		ctx, cancel := context.WithTimeout(context.Background(), opTimeout)
		defer cancel()
		newID, actuallyAllocated, err := crdAllocator.Allocate(ctx, key)
		switch {
		case err != nil:
			log.WithError(err).Errorf("Cannot allocate new CRD ID for %v", key)
			continue

		case !actuallyAllocated:
			scopedLog.Debug("Expected to allocate ID but this ID->key mapping re-existed")
		}

		log.WithFields(logrus.Fields{
			logfields.OldIdentity:    id,
			logfields.Identity:       newID,
			logfields.IdentityLabels: key.GetKey(),
		}).Info("New ID allocated for key in CRD")
	}
}

// initK8s connects to k8s with a allocator.Backend and an initialized
// allocator.Allocator, using the k8s config passed into the command.
func initK8s(ctx context.Context) (crdBackend allocator.Backend, crdAllocator *allocator.Allocator) {
	log.Info("Setting up kubernetes client")

	k8sClientQPSLimit := viper.GetFloat64(option.K8sClientQPSLimit)
	k8sClientBurst := viper.GetInt(option.K8sClientBurst)

	k8s.Configure(k8sAPIServer, k8sKubeConfigPath, float32(k8sClientQPSLimit), k8sClientBurst)

	if err := k8s.Init(); err != nil {
		log.WithError(err).Fatal("Unable to connect to Kubernetes apiserver")
	}

	// Update CRDs to ensure ciliumIdentity is present
	k8s.RegisterCRDs()

	// Create a CRD Backend
	crdBackend, err := identitybackend.NewCRDBackend(identitybackend.CRDBackendConfiguration{
		NodeName: "cilium-preflight",
		Store:    nil,
		Client:   k8s.CiliumClient(),
		KeyType:  cache.GlobalIdentity{},
	})
	if err != nil {
		log.WithError(err).Fatal("Cannot create CRD identity backend")
	}

	// Create a real allocator with CRD as the backend. This mimics the setup in
	// pkg/allocator/cache
	//
	// FIXME: add options to handle clustermesh with this constructor parameter:
	//    allocator.WithPrefixMask(idpool.ID(option.Config.ClusterID<<identity.ClusterIDShift)))
	minID := idpool.ID(identity.MinimalAllocationIdentity)
	maxID := idpool.ID(identity.MaximumAllocationIdentity)
	crdAllocator, err = allocator.NewAllocator(cache.GlobalIdentity{}, crdBackend,
		allocator.WithMax(maxID), allocator.WithMin(minID))
	if err != nil {
		log.WithError(err).Fatal("Unable to initialize Identity Allocator with CRD backend to allocate identities with already allocated IDs")
	}

	// Wait for the initial sync to complete
	if err := crdAllocator.WaitForInitialSync(ctx); err != nil {
		log.WithError(err).Fatal("Error waiting for k8s identity allocator to sync. No identities have been migrated.")
	}

	return crdBackend, crdAllocator
}

// initKVStore connects to the kvstore with a allocator.Backend, initialised to
// find identities at the default cilium paths.
func initKVStore() (kvstoreBackend allocator.Backend) {
	log.Info("Setting up kvstore client")
	setupKvstore()

	idPath := path.Join(cache.IdentitiesPath, "id")
	kvstoreBackend, err := kvstoreallocator.NewKVStoreBackend(cache.IdentitiesPath, idPath, cache.GlobalIdentity{})
	if err != nil {
		log.WithError(err).Fatal("Cannot create kvstore identity backend")
	}

	return kvstoreBackend
}

// getKVStoreIdentities lists all identities in the kvstore. It will wait for
// the listing to complete.
func getKVStoreIdentities(ctx context.Context, kvstoreBackend allocator.Backend) (identities map[idpool.ID]allocator.AllocatorKey, err error) {
	identities = make(map[idpool.ID]allocator.AllocatorKey)
	stopChan := make(chan struct{})

	go kvstoreBackend.ListAndWatch(kvstoreListHandler{
		onAdd: func(id idpool.ID, key allocator.AllocatorKey) {
			log.Debugf("kvstore listed ID: %+v -> %+v", id, key)
			identities[id] = key
		},
		onListDone: func() {
			close(stopChan)
		},
	}, stopChan)
	// This makes the ListAndWatch exit after the initial listing or on a timeout
	// that exits this function

	// Wait for the listing to complete
	select {
	case <-stopChan:
		log.Debug("kvstore ID list complete")

	case <-ctx.Done():
		return nil, errors.New("Timeout while listing identities")
	}

	return identities, nil
}

// kvstoreListHandler is a dummy type to receive callbacks from the kvstore subsystem
type kvstoreListHandler struct {
	onAdd      func(id idpool.ID, key allocator.AllocatorKey)
	onListDone func()
}

func (h kvstoreListHandler) OnListDone()                                       { h.onListDone() }
func (h kvstoreListHandler) OnAdd(id idpool.ID, key allocator.AllocatorKey)    { h.onAdd(id, key) }
func (h kvstoreListHandler) OnModify(id idpool.ID, key allocator.AllocatorKey) {}
func (h kvstoreListHandler) OnDelete(id idpool.ID, key allocator.AllocatorKey) {}
<|endoftext|>


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

// The clnt package go9provides definitions and functions used to implement
// a 9P2000 file client.
package go9p

import (
	"fmt"
	"log"
	"net"
	"sync"
	"sync/atomic"
)

// The Clnt type represents a 9P2000 client. The client is connected to
// a 9P2000 file server and its methods can be used to access and manipulate
// the files exported by the server.
type Clnt struct {
	sync.Mutex
	Debuglevel int    // =0 don't print anything, >0 print Fcalls, >1 print raw packets
	Msize      uint32 // Maximum size of the 9P messages
	Dotu       bool   // If true, 9P2000.u protocol is spoken
	Root       *Fid   // Fid that points to the rood directory
	Id         string // Used when printing debug messages
	Log        *Logger

	conn     net.Conn
	tagpool  *pool
	fidpool  *pool
	reqout   chan *Req
	done     chan bool
	reqfirst *Req
	reqlast  *Req
	err      error

	reqchan chan *Req
	tchan   chan *Fcall

	next, prev *Clnt
}

// A Fid type represents a file on the server. Fids are used for the
// low level methods that correspond directly to the 9P2000 message requests
type Fid struct {
	sync.Mutex
	Clnt   *Clnt // Client the fid belongs to
	Iounit uint32
	Qid           // The Qid description for the file
	Mode   uint8  // Open mode (one of O* values) (if file is open)
	Fid    uint32 // Fid number
	User          // The user the fid belongs to
	walked bool   // true if the fid points to a walked file on the server
}

// The file is similar to the Fid, but is used in the high-level client
// interface. We expose the Fid so that client code can use Remove
// on a fid, the same way a kernel can.
type File struct {
	Fid    *Fid
	offset uint64
}

type Req struct {
	sync.Mutex
	Clnt       *Clnt
	Tc         *Fcall
	Rc         *Fcall
	Err        error
	Done       chan *Req
	tag        uint16
	prev, next *Req
	fid        *Fid
}

type ClntList struct {
	sync.Mutex
	clntList, clntLast *Clnt
}

var clnts *ClntList
var DefaultDebuglevel int
var DefaultLogger *Logger

func (clnt *Clnt) Rpcnb(r *Req) error {
	var tag uint16

	if r.Tc.Type == Tversion {
		tag = NOTAG
	} else {
		tag = r.tag
	}

	SetTag(r.Tc, tag)
	clnt.Lock()
	if clnt.err != nil {
		clnt.Unlock()
		return clnt.err
	}

	if clnt.reqlast != nil {
		clnt.reqlast.next = r
	} else {
		clnt.reqfirst = r
	}

	r.prev = clnt.reqlast
	clnt.reqlast = r
	clnt.Unlock()

	clnt.reqout <- r
	return nil
}

func (clnt *Clnt) Rpc(tc *Fcall) (rc *Fcall, err error) {
	r := clnt.ReqAlloc()
	r.Tc = tc
	r.Done = make(chan *Req)
	err = clnt.Rpcnb(r)
	if err != nil {
		return
	}

	<-r.Done
	rc = r.Rc
	err = r.Err
	clnt.ReqFree(r)
	return
}

func (clnt *Clnt) recv() {
	var err error
	var buf []byte

	err = nil
	pos := 0
	for {
		// Connect can change the client Msize.
		clntmsize := int(atomic.LoadUint32(&clnt.Msize))
		if len(buf) < clntmsize {
			b := make([]byte, clntmsize*8)
			copy(b, buf[0:pos])
			buf = b
			b = nil
		}

		n, oerr := clnt.conn.Read(buf[pos:])
		if oerr != nil || n == 0 {
			err = &Error{oerr.Error(), EIO}
			clnt.Lock()
			clnt.err = err
			clnt.Unlock()
			goto closed
		}

		pos += n
		for pos > 4 {
			sz, _ := Gint32(buf)
			if pos < int(sz) {
				if len(buf) < int(sz) {
					b := make([]byte, atomic.LoadUint32(&clnt.Msize)*8)
					copy(b, buf[0:pos])
					buf = b
					b = nil
				}

				break
			}

			fc, err, fcsize := Unpack(buf, clnt.Dotu)
			clnt.Lock()
			if err != nil {
				clnt.err = err
				clnt.conn.Close()
				clnt.Unlock()
				goto closed
			}

			if clnt.Debuglevel > 0 {
				clnt.logFcall(fc)
				if clnt.Debuglevel&DbgPrintPackets != 0 {
					log.Println("}-}", clnt.Id, fmt.Sprint(fc.Pkt))
				}

				if clnt.Debuglevel&DbgPrintFcalls != 0 {
					log.Println("}}}", clnt.Id, fc.String())
				}
			}

			var r *Req = nil
			for r = clnt.reqfirst; r != nil; r = r.next {
				if r.Tc.Tag == fc.Tag {
					break
				}
			}

			if r == nil {
				clnt.err = &Error{"unexpected response", EINVAL}
				clnt.conn.Close()
				clnt.Unlock()
				goto closed
			}

			r.Rc = fc
			if r.prev != nil {
				r.prev.next = r.next
			} else {
				clnt.reqfirst = r.next
			}

			if r.next != nil {
				r.next.prev = r.prev
			} else {
				clnt.reqlast = r.prev
			}
			clnt.Unlock()

			if r.Tc.Type != r.Rc.Type-1 {
				if r.Rc.Type != Rerror {
					r.Err = &Error{"invalid response", EINVAL}
					log.Println(fmt.Sprintf("TTT %v", r.Tc))
					log.Println(fmt.Sprintf("RRR %v", r.Rc))
				} else {
					if r.Err == nil {
						r.Err = &Error{r.Rc.Error, r.Rc.Errornum}
					}
				}
			}

			if r.Done != nil {
				r.Done <- r
			}

			pos -= fcsize
			buf = buf[fcsize:]
		}
	}

closed:
	clnt.done <- true

	/* send error to all pending requests */
	clnt.Lock()
	r := clnt.reqfirst
	clnt.reqfirst = nil
	clnt.reqlast = nil
	if err == nil {
		err = clnt.err
	}
	clnt.Unlock()
	for ; r != nil; r = r.next {
		r.Err = err
		if r.Done != nil {
			r.Done <- r
		}
	}

	clnts.Lock()
	if clnt.prev != nil {
		clnt.prev.next = clnt.next
	} else {
		clnts.clntList = clnt.next
	}

	if clnt.next != nil {
		clnt.next.prev = clnt.prev
	} else {
		clnts.clntLast = clnt.prev
	}
	clnts.Unlock()

	if sop, ok := (interface{}(clnt)).(StatsOps); ok {
		sop.statsUnregister()
	}
}

func (clnt *Clnt) send() {
	for {
		select {
		case <-clnt.done:
			return

		case req := <-clnt.reqout:
			if clnt.Debuglevel > 0 {
				clnt.logFcall(req.Tc)
				if clnt.Debuglevel&DbgPrintPackets != 0 {
					log.Println("{-{", clnt.Id, fmt.Sprint(req.Tc.Pkt))
				}

				if clnt.Debuglevel&DbgPrintFcalls != 0 {
					log.Println("{{{", clnt.Id, req.Tc.String())
				}
			}

			for buf := req.Tc.Pkt; len(buf) > 0; {
				n, err := clnt.conn.Write(buf)
				if err != nil {
					/* just close the socket, will get signal on clnt.done */
					clnt.conn.Close()
					break
				}

				buf = buf[n:]
			}
		}
	}
}

// Creates and initializes a new Clnt object. Doesn't send any data
// on the wire.
func NewClnt(c net.Conn, msize uint32, dotu bool) *Clnt {
	clnt := new(Clnt)
	clnt.conn = c
	clnt.Msize = msize
	clnt.Dotu = dotu
	clnt.Debuglevel = DefaultDebuglevel
	clnt.Log = DefaultLogger
	clnt.Id = c.RemoteAddr().String() + ":"
	clnt.tagpool = newPool(uint32(NOTAG))
	clnt.fidpool = newPool(NOFID)
	clnt.reqout = make(chan *Req)
	clnt.done = make(chan bool)
	clnt.reqchan = make(chan *Req, 16)
	clnt.tchan = make(chan *Fcall, 16)

	go clnt.recv()
	go clnt.send()

	clnts.Lock()
	if clnts.clntLast != nil {
		clnts.clntLast.next = clnt
	} else {
		clnts.clntList = clnt
	}

	clnt.prev = clnts.clntLast
	clnts.clntLast = clnt
	clnts.Unlock()

	if sop, ok := (interface{}(clnt)).(StatsOps); ok {
		sop.statsRegister()
	}

	return clnt
}

// Establishes a new socket connection to the 9P server and creates
// a client object for it. Negotiates the dialect and msize for the
// connection. Returns a Clnt object, or Error.
func Connect(c net.Conn, msize uint32, dotu bool) (*Clnt, error) {
	clnt := NewClnt(c, msize, dotu)
	ver := "9P2000"
	if clnt.Dotu {
		ver = "9P2000.u"
	}

	clntmsize := atomic.LoadUint32(&clnt.Msize)
	tc := NewFcall(clntmsize)
	err := PackTversion(tc, clntmsize, ver)
	if err != nil {
		return nil, err
	}

	rc, err := clnt.Rpc(tc)
	if err != nil {
		return nil, err
	}

	if rc.Msize < atomic.LoadUint32(&clnt.Msize) {
		atomic.StoreUint32(&clnt.Msize, rc.Msize)
	}

	clnt.Dotu = rc.Version == "9P2000.u" && clnt.Dotu
	return clnt, nil
}

// Creates a new Fid object for the client
func (clnt *Clnt) FidAlloc() *Fid {
	fid := new(Fid)
	fid.Fid = clnt.fidpool.getId()
	fid.Clnt = clnt

	return fid
}

func (clnt *Clnt) NewFcall() *Fcall {
	select {
	case tc := <-clnt.tchan:
		return tc
	default:
	}
	return NewFcall(atomic.LoadUint32(&clnt.Msize))
}

func (clnt *Clnt) FreeFcall(fc *Fcall) {
	if fc != nil && len(fc.Buf) >= int(atomic.LoadUint32(&clnt.Msize)) {
		select {
		case clnt.tchan <- fc:
			break
		default:
		}
	}
}

func (clnt *Clnt) ReqAlloc() *Req {
	var req *Req
	select {
	case req = <-clnt.reqchan:
		break
	default:
		req = new(Req)
		req.Clnt = clnt
		req.tag = uint16(clnt.tagpool.getId())
	}
	return req
}

func (clnt *Clnt) ReqFree(req *Req) {
	clnt.FreeFcall(req.Tc)
	req.Tc = nil
	req.Rc = nil
	req.Err = nil
	req.Done = nil
	req.next = nil
	req.prev = nil

	select {
	case clnt.reqchan <- req:
		break
	default:
		clnt.tagpool.putId(uint32(req.tag))
	}
}

func (clnt *Clnt) logFcall(fc *Fcall) {
	if clnt.Debuglevel&DbgLogPackets != 0 {
		pkt := make([]byte, len(fc.Pkt))
		copy(pkt, fc.Pkt)
		clnt.Log.Log(pkt, clnt, DbgLogPackets)
	}

	if clnt.Debuglevel&DbgLogFcalls != 0 {
		f := new(Fcall)
		*f = *fc
		f.Pkt = nil
		clnt.Log.Log(f, clnt, DbgLogFcalls)
	}
}

func init() {
	clnts = new(ClntList)
	if sop, ok := (interface{}(clnts)).(StatsOps); ok {
		sop.statsRegister()
	}
}
<commit_msg>Add go9p.FidFile.<commit_after>// Copyright 2009 The Go9p Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The clnt package go9provides definitions and functions used to implement
// a 9P2000 file client.
package go9p

import (
	"fmt"
	"log"
	"net"
	"sync"
	"sync/atomic"
)

// The Clnt type represents a 9P2000 client. The client is connected to
// a 9P2000 file server and its methods can be used to access and manipulate
// the files exported by the server.
type Clnt struct {
	sync.Mutex
	Debuglevel int    // =0 don't print anything, >0 print Fcalls, >1 print raw packets
	Msize      uint32 // Maximum size of the 9P messages
	Dotu       bool   // If true, 9P2000.u protocol is spoken
	Root       *Fid   // Fid that points to the rood directory
	Id         string // Used when printing debug messages
	Log        *Logger

	conn     net.Conn
	tagpool  *pool
	fidpool  *pool
	reqout   chan *Req
	done     chan bool
	reqfirst *Req
	reqlast  *Req
	err      error

	reqchan chan *Req
	tchan   chan *Fcall

	next, prev *Clnt
}

// A Fid type represents a file on the server. Fids are used for the
// low level methods that correspond directly to the 9P2000 message requests
type Fid struct {
	sync.Mutex
	Clnt   *Clnt // Client the fid belongs to
	Iounit uint32
	Qid           // The Qid description for the file
	Mode   uint8  // Open mode (one of O* values) (if file is open)
	Fid    uint32 // Fid number
	User          // The user the fid belongs to
	walked bool   // true if the fid points to a walked file on the server
}

// The file is similar to the Fid, but is used in the high-level client
// interface. We expose the Fid so that client code can use Remove
// on a fid, the same way a kernel can.
type File struct {
	Fid    *Fid
	offset uint64
}

type Req struct {
	sync.Mutex
	Clnt       *Clnt
	Tc         *Fcall
	Rc         *Fcall
	Err        error
	Done       chan *Req
	tag        uint16
	prev, next *Req
	fid        *Fid
}

type ClntList struct {
	sync.Mutex
	clntList, clntLast *Clnt
}

var clnts *ClntList
var DefaultDebuglevel int
var DefaultLogger *Logger

func (clnt *Clnt) Rpcnb(r *Req) error {
	var tag uint16

	if r.Tc.Type == Tversion {
		tag = NOTAG
	} else {
		tag = r.tag
	}

	SetTag(r.Tc, tag)
	clnt.Lock()
	if clnt.err != nil {
		clnt.Unlock()
		return clnt.err
	}

	if clnt.reqlast != nil {
		clnt.reqlast.next = r
	} else {
		clnt.reqfirst = r
	}

	r.prev = clnt.reqlast
	clnt.reqlast = r
	clnt.Unlock()

	clnt.reqout <- r
	return nil
}

func (clnt *Clnt) Rpc(tc *Fcall) (rc *Fcall, err error) {
	r := clnt.ReqAlloc()
	r.Tc = tc
	r.Done = make(chan *Req)
	err = clnt.Rpcnb(r)
	if err != nil {
		return
	}

	<-r.Done
	rc = r.Rc
	err = r.Err
	clnt.ReqFree(r)
	return
}

func (clnt *Clnt) recv() {
	var err error
	var buf []byte

	err = nil
	pos := 0
	for {
		// Connect can change the client Msize.
		clntmsize := int(atomic.LoadUint32(&clnt.Msize))
		if len(buf) < clntmsize {
			b := make([]byte, clntmsize*8)
			copy(b, buf[0:pos])
			buf = b
			b = nil
		}

		n, oerr := clnt.conn.Read(buf[pos:])
		if oerr != nil || n == 0 {
			err = &Error{oerr.Error(), EIO}
			clnt.Lock()
			clnt.err = err
			clnt.Unlock()
			goto closed
		}

		pos += n
		for pos > 4 {
			sz, _ := Gint32(buf)
			if pos < int(sz) {
				if len(buf) < int(sz) {
					b := make([]byte, atomic.LoadUint32(&clnt.Msize)*8)
					copy(b, buf[0:pos])
					buf = b
					b = nil
				}

				break
			}

			fc, err, fcsize := Unpack(buf, clnt.Dotu)
			clnt.Lock()
			if err != nil {
				clnt.err = err
				clnt.conn.Close()
				clnt.Unlock()
				goto closed
			}

			if clnt.Debuglevel > 0 {
				clnt.logFcall(fc)
				if clnt.Debuglevel&DbgPrintPackets != 0 {
					log.Println("}-}", clnt.Id, fmt.Sprint(fc.Pkt))
				}

				if clnt.Debuglevel&DbgPrintFcalls != 0 {
					log.Println("}}}", clnt.Id, fc.String())
				}
			}

			var r *Req = nil
			for r = clnt.reqfirst; r != nil; r = r.next {
				if r.Tc.Tag == fc.Tag {
					break
				}
			}

			if r == nil {
				clnt.err = &Error{"unexpected response", EINVAL}
				clnt.conn.Close()
				clnt.Unlock()
				goto closed
			}

			r.Rc = fc
			if r.prev != nil {
				r.prev.next = r.next
			} else {
				clnt.reqfirst = r.next
			}

			if r.next != nil {
				r.next.prev = r.prev
			} else {
				clnt.reqlast = r.prev
			}
			clnt.Unlock()

			if r.Tc.Type != r.Rc.Type-1 {
				if r.Rc.Type != Rerror {
					r.Err = &Error{"invalid response", EINVAL}
					log.Println(fmt.Sprintf("TTT %v", r.Tc))
					log.Println(fmt.Sprintf("RRR %v", r.Rc))
				} else {
					if r.Err == nil {
						r.Err = &Error{r.Rc.Error, r.Rc.Errornum}
					}
				}
			}

			if r.Done != nil {
				r.Done <- r
			}

			pos -= fcsize
			buf = buf[fcsize:]
		}
	}

closed:
	clnt.done <- true

	/* send error to all pending requests */
	clnt.Lock()
	r := clnt.reqfirst
	clnt.reqfirst = nil
	clnt.reqlast = nil
	if err == nil {
		err = clnt.err
	}
	clnt.Unlock()
	for ; r != nil; r = r.next {
		r.Err = err
		if r.Done != nil {
			r.Done <- r
		}
	}

	clnts.Lock()
	if clnt.prev != nil {
		clnt.prev.next = clnt.next
	} else {
		clnts.clntList = clnt.next
	}

	if clnt.next != nil {
		clnt.next.prev = clnt.prev
	} else {
		clnts.clntLast = clnt.prev
	}
	clnts.Unlock()

	if sop, ok := (interface{}(clnt)).(StatsOps); ok {
		sop.statsUnregister()
	}
}

func (clnt *Clnt) send() {
	for {
		select {
		case <-clnt.done:
			return

		case req := <-clnt.reqout:
			if clnt.Debuglevel > 0 {
				clnt.logFcall(req.Tc)
				if clnt.Debuglevel&DbgPrintPackets != 0 {
					log.Println("{-{", clnt.Id, fmt.Sprint(req.Tc.Pkt))
				}

				if clnt.Debuglevel&DbgPrintFcalls != 0 {
					log.Println("{{{", clnt.Id, req.Tc.String())
				}
			}

			for buf := req.Tc.Pkt; len(buf) > 0; {
				n, err := clnt.conn.Write(buf)
				if err != nil {
					/* just close the socket, will get signal on clnt.done */
					clnt.conn.Close()
					break
				}

				buf = buf[n:]
			}
		}
	}
}

// Creates and initializes a new Clnt object. Doesn't send any data
// on the wire.
func NewClnt(c net.Conn, msize uint32, dotu bool) *Clnt {
	clnt := new(Clnt)
	clnt.conn = c
	clnt.Msize = msize
	clnt.Dotu = dotu
	clnt.Debuglevel = DefaultDebuglevel
	clnt.Log = DefaultLogger
	clnt.Id = c.RemoteAddr().String() + ":"
	clnt.tagpool = newPool(uint32(NOTAG))
	clnt.fidpool = newPool(NOFID)
	clnt.reqout = make(chan *Req)
	clnt.done = make(chan bool)
	clnt.reqchan = make(chan *Req, 16)
	clnt.tchan = make(chan *Fcall, 16)

	go clnt.recv()
	go clnt.send()

	clnts.Lock()
	if clnts.clntLast != nil {
		clnts.clntLast.next = clnt
	} else {
		clnts.clntList = clnt
	}

	clnt.prev = clnts.clntLast
	clnts.clntLast = clnt
	clnts.Unlock()

	if sop, ok := (interface{}(clnt)).(StatsOps); ok {
		sop.statsRegister()
	}

	return clnt
}

// Establishes a new socket connection to the 9P server and creates
// a client object for it. Negotiates the dialect and msize for the
// connection. Returns a Clnt object, or Error.
func Connect(c net.Conn, msize uint32, dotu bool) (*Clnt, error) {
	clnt := NewClnt(c, msize, dotu)
	ver := "9P2000"
	if clnt.Dotu {
		ver = "9P2000.u"
	}

	clntmsize := atomic.LoadUint32(&clnt.Msize)
	tc := NewFcall(clntmsize)
	err := PackTversion(tc, clntmsize, ver)
	if err != nil {
		return nil, err
	}

	rc, err := clnt.Rpc(tc)
	if err != nil {
		return nil, err
	}

	if rc.Msize < atomic.LoadUint32(&clnt.Msize) {
		atomic.StoreUint32(&clnt.Msize, rc.Msize)
	}

	clnt.Dotu = rc.Version == "9P2000.u" && clnt.Dotu
	return clnt, nil
}

// Creates a new Fid object for the client
func (clnt *Clnt) FidAlloc() *Fid {
	fid := new(Fid)
	fid.Fid = clnt.fidpool.getId()
	fid.Clnt = clnt

	return fid
}

func (clnt *Clnt) NewFcall() *Fcall {
	select {
	case tc := <-clnt.tchan:
		return tc
	default:
	}
	return NewFcall(atomic.LoadUint32(&clnt.Msize))
}

func (clnt *Clnt) FreeFcall(fc *Fcall) {
	if fc != nil && len(fc.Buf) >= int(atomic.LoadUint32(&clnt.Msize)) {
		select {
		case clnt.tchan <- fc:
			break
		default:
		}
	}
}

func (clnt *Clnt) ReqAlloc() *Req {
	var req *Req
	select {
	case req = <-clnt.reqchan:
		break
	default:
		req = new(Req)
		req.Clnt = clnt
		req.tag = uint16(clnt.tagpool.getId())
	}
	return req
}

func (clnt *Clnt) ReqFree(req *Req) {
	clnt.FreeFcall(req.Tc)
	req.Tc = nil
	req.Rc = nil
	req.Err = nil
	req.Done = nil
	req.next = nil
	req.prev = nil

	select {
	case clnt.reqchan <- req:
		break
	default:
		clnt.tagpool.putId(uint32(req.tag))
	}
}

func (clnt *Clnt) logFcall(fc *Fcall) {
	if clnt.Debuglevel&DbgLogPackets != 0 {
		pkt := make([]byte, len(fc.Pkt))
		copy(pkt, fc.Pkt)
		clnt.Log.Log(pkt, clnt, DbgLogPackets)
	}

	if clnt.Debuglevel&DbgLogFcalls != 0 {
		f := new(Fcall)
		*f = *fc
		f.Pkt = nil
		clnt.Log.Log(f, clnt, DbgLogFcalls)
	}
}

// FidFile returns a File that represents the given Fid, initially at the given
// offset.
func FidFile(fid *Fid, offset uint64) *File {
	return &File{fid, offset}
}

func init() {
	clnts = new(ClntList)
	if sop, ok := (interface{}(clnts)).(StatsOps); ok {
		sop.statsRegister()
	}
}
<|endoftext|>


<commit_before>/*
 * Copyright 2016 Frank Wessels <fwessels@xs4all.nl>
 *
 * 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 cmd

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

	"github.com/s3git/s3git-go"
	"github.com/cheggaaa/pb"
	"github.com/spf13/cobra"
	"github.com/dustin/go-humanize"
)

// cloneCmd represents the clone command
var cloneCmd = &cobra.Command{
	Use:   "clone [resource]",
	Short: "Clone a repository into a new directory",
	Long: "Clone a repository into a new directory",
	Run: func(cmd *cobra.Command, args []string) {

		if len(args) == 0 {
			er("Missing resource to clone from")
		}

		parts := strings.Split(args[0], "//")
		if len(parts) != 2 {
			er(fmt.Sprintf("Bad resource for cloning (missing '//' separator): %s", args[0]))
		}

		dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
		if err != nil {
			er(err)
		}

		dir += "/" + parts[1]

		// Check whether directory to clone into does not yet exist -- abort otherwise
		if _, err := os.Stat(dir); err == nil {
			er(fmt.Sprintf("Cannot clone into existing directory: %s", dir))
		}

		// Output directory and create it
		fmt.Println("Cloning into", dir)
		err = os.MkdirAll(dir, os.ModePerm)
		if err != nil {
			er(err)
		}

		var barDownloading, barProcessing *pb.ProgressBar

		progressDownload := func(total int64) {
			if barDownloading == nil {
				barDownloading = pb.New64(total).Start()
				barDownloading.Prefix("Downloading ")
			}
			if barDownloading.Increment() == int(total) {
				barDownloading.Finish()
			}
		}

		progressProcessing := func(total int64) {
			if barProcessing == nil {
				barProcessing = pb.New64(total).Start()
				barProcessing.Prefix("Processing  ")
			}
			if barProcessing.Increment() == int(total) {
				barProcessing.Finish()
			}
		}

		options := []s3git.CloneOptions{}
		options = append(options, s3git.CloneOptionSetAccessKey(accessKey))
		options = append(options, s3git.CloneOptionSetSecretKey(secretKey))

		repo, err := s3git.Clone(args[0], dir, progressDownload, progressProcessing, options...)
		if err != nil {
			er(err)
		}

		stats, err := repo.Statistics()
		fmt.Printf("Done. Totaling %s objects.\n", humanize.Comma(int64(stats.Objects)))
	},
}

func init() {
	RootCmd.AddCommand(cloneCmd)

	// Add local message flags
	cloneCmd.Flags().StringVarP(&accessKey, "access", "a", "", "Access key for S3 remote")
	cloneCmd.Flags().StringVarP(&secretKey, "secret", "s", "", "Secret key for S3 remote")
}<commit_msg>Clone in current dir instead of directory of executable<commit_after>/*
 * Copyright 2016 Frank Wessels <fwessels@xs4all.nl>
 *
 * 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 cmd

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

	"github.com/s3git/s3git-go"
	"github.com/cheggaaa/pb"
	"github.com/spf13/cobra"
	"github.com/dustin/go-humanize"
)

// cloneCmd represents the clone command
var cloneCmd = &cobra.Command{
	Use:   "clone [resource]",
	Short: "Clone a repository into a new directory",
	Long: "Clone a repository into a new directory",
	Run: func(cmd *cobra.Command, args []string) {

		if len(args) == 0 {
			er("Missing resource to clone from")
		}

		parts := strings.Split(args[0], "//")
		if len(parts) != 2 {
			er(fmt.Sprintf("Bad resource for cloning (missing '//' separator): %s", args[0]))
		}

		dir, err := filepath.Abs(filepath.Dir("."))
		if err != nil {
			er(err)
		}

		dir += "/" + parts[1]

		// Check whether directory to clone into does not yet exist -- abort otherwise
		if _, err := os.Stat(dir); err == nil {
			er(fmt.Sprintf("Cannot clone into existing directory: %s", dir))
		}

		// Output directory and create it
		fmt.Println("Cloning into", dir)
		err = os.MkdirAll(dir, os.ModePerm)
		if err != nil {
			er(err)
		}

		var barDownloading, barProcessing *pb.ProgressBar

		progressDownload := func(total int64) {
			if barDownloading == nil {
				barDownloading = pb.New64(total).Start()
				barDownloading.Prefix("Downloading ")
			}
			if barDownloading.Increment() == int(total) {
				barDownloading.Finish()
			}
		}

		progressProcessing := func(total int64) {
			if barProcessing == nil {
				barProcessing = pb.New64(total).Start()
				barProcessing.Prefix("Processing  ")
			}
			if barProcessing.Increment() == int(total) {
				barProcessing.Finish()
			}
		}

		options := []s3git.CloneOptions{}
		options = append(options, s3git.CloneOptionSetAccessKey(accessKey))
		options = append(options, s3git.CloneOptionSetSecretKey(secretKey))

		repo, err := s3git.Clone(args[0], dir, progressDownload, progressProcessing, options...)
		if err != nil {
			er(err)
		}

		stats, err := repo.Statistics()
		fmt.Printf("Done. Totaling %s objects.\n", humanize.Comma(int64(stats.Objects)))
	},
}

func init() {
	RootCmd.AddCommand(cloneCmd)

	// Add local message flags
	cloneCmd.Flags().StringVarP(&accessKey, "access", "a", "", "Access key for S3 remote")
	cloneCmd.Flags().StringVarP(&secretKey, "secret", "s", "", "Secret key for S3 remote")
}<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package cmd

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"

	"github.com/kelseyhightower/envconfig"
	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/lib/consts"
	"github.com/loadimpact/k6/loader"
	"github.com/loadimpact/k6/stats/cloud"
	"github.com/loadimpact/k6/ui"
	"github.com/loadimpact/k6/ui/pb"
	"github.com/pkg/errors"
	"github.com/spf13/afero"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"

	"github.com/sirupsen/logrus"
)

var (
	exitOnRunning = os.Getenv("K6_EXIT_ON_RUNNING") != ""
)

var cloudCmd = &cobra.Command{
	Use:   "cloud",
	Short: "Run a test on the cloud",
	Long: `Run a test on the cloud.

This will execute the test on the Load Impact cloud service. Use "k6 login cloud" to authenticate.`,
	Example: `
        k6 cloud script.js`[1:],
	Args: exactArgsWithMsg(1, "arg should either be \"-\", if reading script from stdin, or a path to a script file"),
	RunE: func(cmd *cobra.Command, args []string) error {
		//TODO: disable in quiet mode?
		_, _ = BannerColor.Fprintf(stdout, "\n%s\n\n", consts.Banner)

		progressBar := pb.New(pb.WithConstLeft(" Init"))
		printBar(progressBar, "Parsing script")

		// Runner
		pwd, err := os.Getwd()
		if err != nil {
			return err
		}

		filename := args[0]
		filesystems := loader.CreateFilesystems()
		src, err := loader.ReadSource(filename, pwd, filesystems, os.Stdin)
		if err != nil {
			return err
		}

		runtimeOptions, err := getRuntimeOptions(cmd.Flags())
		if err != nil {
			return err
		}

		printBar(progressBar, "Getting script options")
		r, err := newRunner(src, runType, filesystems, runtimeOptions)
		if err != nil {
			return err
		}

		printBar(progressBar, "Consolidating options")
		cliOpts, err := getOptions(cmd.Flags())
		if err != nil {
			return err
		}
		conf, err := getConsolidatedConfig(afero.NewOsFs(), Config{Options: cliOpts}, r)
		if err != nil {
			return err
		}

		derivedConf, cerr := deriveAndValidateConfig(conf)
		if cerr != nil {
			return ExitCode{cerr, invalidConfigErrorCode}
		}

		//TODO: warn about lack of support for --no-setup and --no-teardown in the cloud?
		//TODO: validate for usage of execution segment
		//TODO: validate for externally controlled executor (i.e. executors that aren't distributable)
		//TODO: move those validations to a separate function and reuse validateConfig()?

		err = r.SetOptions(conf.Options)
		if err != nil {
			return err
		}

		// Cloud config
		cloudConfig := cloud.NewConfig().Apply(derivedConf.Collectors.Cloud)
		if err := envconfig.Process("k6", &cloudConfig); err != nil {
			return err
		}
		if !cloudConfig.Token.Valid {
			return errors.New("Not logged in, please use `k6 login cloud`.")
		}

		printBar(progressBar, "Building the archive")
		arc := r.MakeArchive()
		// TODO: Fix this
		// We reuse cloud.Config for parsing options.ext.loadimpact, but this probably shouldn't be
		// done as the idea of options.ext is that they are extensible without touching k6. But in
		// order for this to happen we shouldn't actually marshall cloud.Config on top of it because
		// it will be missing some fields that aren't actually mentioned in the struct.
		// So in order for use to copy the fields that we need for loadimpact's api we unmarshal in
		// map[string]interface{} and copy what we need if it isn't set already
		var tmpCloudConfig map[string]interface{}
		if val, ok := arc.Options.External["loadimpact"]; ok {
			var dec = json.NewDecoder(bytes.NewReader(val))
			dec.UseNumber() // otherwise float64 are used
			if err := dec.Decode(&tmpCloudConfig); err != nil {
				return err
			}
		}

		if err := cloud.MergeFromExternal(arc.Options.External, &cloudConfig); err != nil {
			return err
		}
		if tmpCloudConfig == nil {
			tmpCloudConfig = make(map[string]interface{}, 3)
		}

		if _, ok := tmpCloudConfig["token"]; !ok && cloudConfig.Token.Valid {
			tmpCloudConfig["token"] = cloudConfig.Token
		}
		if _, ok := tmpCloudConfig["name"]; !ok && cloudConfig.Name.Valid {
			tmpCloudConfig["name"] = cloudConfig.Name
		}
		if _, ok := tmpCloudConfig["projectID"]; !ok && cloudConfig.ProjectID.Valid {
			tmpCloudConfig["projectID"] = cloudConfig.ProjectID
		}

		if arc.Options.External == nil {
			arc.Options.External = make(map[string]json.RawMessage)
		}
		arc.Options.External["loadimpact"], err = json.Marshal(tmpCloudConfig)
		if err != nil {
			return err
		}

		name := cloudConfig.Name.String
		if !cloudConfig.Name.Valid || cloudConfig.Name.String == "" {
			name = filepath.Base(filename)
		}

		// Start cloud test run
		printBar(progressBar, "Validating script options")
		client := cloud.NewClient(cloudConfig.Token.String, cloudConfig.Host.String, consts.Version)
		if err := client.ValidateOptions(arc.Options); err != nil {
			return err
		}

		printBar(progressBar, "Uploading archive")
		refID, err := client.StartCloudTestRun(name, cloudConfig.ProjectID.Int64, arc)
		if err != nil {
			return err
		}
		progressBar.Modify(pb.WithConstLeft("   Run"))
		printBar(progressBar, "Initializing the cloud test")

		testURL := cloud.URLForResults(refID, cloudConfig)
		fprintf(stdout, "\n\n")
		fprintf(stdout, "   executor: %s\n", ui.ValueColor.Sprint("cloud"))
		fprintf(stdout, "     script: %s\n", ui.ValueColor.Sprint(filename))
		fprintf(stdout, "     output: %s\n", ui.ValueColor.Sprint(testURL))
		//TODO: print executors information
		fprintf(stdout, "\n")
		printBar(progressBar, "Initializing the cloud test")

		// The quiet option hides the progress bar and disallow aborting the test
		if quiet {
			return nil
		}

		// Trap Interrupts, SIGINTs and SIGTERMs.
		sigC := make(chan os.Signal, 1)
		signal.Notify(sigC, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
		defer signal.Stop(sigC)

		var progressErr error
		testProgress := &cloud.TestProgressResponse{}
		percentageFmt := "[" + pb.GetFixedLengthFloatFormat(100, 2) + "%%] %s"
		progressBar.Modify(
			pb.WithProgress(func() (float64, string) {
				if testProgress.RunStatus < lib.RunStatusRunning {
					return 0, testProgress.RunStatusText
				}
				return testProgress.Progress, fmt.Sprintf(percentageFmt, testProgress.Progress*100, testProgress.RunStatusText)
			}),
		)

		ticker := time.NewTicker(time.Millisecond * 2000)
		shouldExitLoop := false

	runningLoop:
		for {
			select {
			case <-ticker.C:
				testProgress, progressErr = client.GetTestProgress(refID)
				if progressErr == nil {
					if (testProgress.RunStatus > lib.RunStatusRunning) || (exitOnRunning && testProgress.RunStatus == lib.RunStatusRunning) {
						shouldExitLoop = true
					}
					printBar(progressBar, "")
				} else {
					logrus.WithError(progressErr).Error("Test progress error")
				}
				if shouldExitLoop {
					break runningLoop
				}
			case sig := <-sigC:
				logrus.WithField("sig", sig).Print("Exiting in response to signal...")
				err := client.StopCloudTestRun(refID)
				if err != nil {
					logrus.WithError(err).Error("Stop cloud test error")
				}
				shouldExitLoop = true // Exit after the next GetTestProgress call
			}
		}

		if testProgress == nil {
			return ExitCode{errors.New("Test progress error"), 98}
		}

		fprintf(stdout, "     test status: %s\n", ui.ValueColor.Sprint(testProgress.RunStatusText))

		if testProgress.ResultStatus == cloud.ResultStatusFailed {
			return ExitCode{errors.New("The test has failed"), 99}
		}

		return nil
	},
}

func cloudCmdFlagSet() *pflag.FlagSet {
	flags := pflag.NewFlagSet("", pflag.ContinueOnError)
	flags.SortFlags = false
	flags.AddFlagSet(optionFlagSet())
	flags.AddFlagSet(runtimeOptionFlagSet(false))

	//TODO: Figure out a better way to handle the CLI flags:
	// - the default value is specified in this way so we don't overwrire whatever
	//   was specified via the environment variable
	// - global variables are not very testable... :/
	flags.BoolVar(&exitOnRunning, "exit-on-running", exitOnRunning, "exits when test reaches the running status")
	// We also need to explicitly set the default value for the usage message here, so setting
	// K6_EXIT_ON_RUNNING=true won't affect the usage message
	flags.Lookup("exit-on-running").DefValue = "false"

	return flags
}

func init() {
	RootCmd.AddCommand(cloudCmd)
	cloudCmd.Flags().SortFlags = false
	cloudCmd.Flags().AddFlagSet(cloudCmdFlagSet())
}
<commit_msg>Remove an obsolete TODO<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package cmd

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"

	"github.com/kelseyhightower/envconfig"
	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/lib/consts"
	"github.com/loadimpact/k6/loader"
	"github.com/loadimpact/k6/stats/cloud"
	"github.com/loadimpact/k6/ui"
	"github.com/loadimpact/k6/ui/pb"
	"github.com/pkg/errors"
	"github.com/spf13/afero"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"

	"github.com/sirupsen/logrus"
)

var (
	exitOnRunning = os.Getenv("K6_EXIT_ON_RUNNING") != ""
)

var cloudCmd = &cobra.Command{
	Use:   "cloud",
	Short: "Run a test on the cloud",
	Long: `Run a test on the cloud.

This will execute the test on the Load Impact cloud service. Use "k6 login cloud" to authenticate.`,
	Example: `
        k6 cloud script.js`[1:],
	Args: exactArgsWithMsg(1, "arg should either be \"-\", if reading script from stdin, or a path to a script file"),
	RunE: func(cmd *cobra.Command, args []string) error {
		//TODO: disable in quiet mode?
		_, _ = BannerColor.Fprintf(stdout, "\n%s\n\n", consts.Banner)

		progressBar := pb.New(pb.WithConstLeft(" Init"))
		printBar(progressBar, "Parsing script")

		// Runner
		pwd, err := os.Getwd()
		if err != nil {
			return err
		}

		filename := args[0]
		filesystems := loader.CreateFilesystems()
		src, err := loader.ReadSource(filename, pwd, filesystems, os.Stdin)
		if err != nil {
			return err
		}

		runtimeOptions, err := getRuntimeOptions(cmd.Flags())
		if err != nil {
			return err
		}

		printBar(progressBar, "Getting script options")
		r, err := newRunner(src, runType, filesystems, runtimeOptions)
		if err != nil {
			return err
		}

		printBar(progressBar, "Consolidating options")
		cliOpts, err := getOptions(cmd.Flags())
		if err != nil {
			return err
		}
		conf, err := getConsolidatedConfig(afero.NewOsFs(), Config{Options: cliOpts}, r)
		if err != nil {
			return err
		}

		derivedConf, cerr := deriveAndValidateConfig(conf)
		if cerr != nil {
			return ExitCode{cerr, invalidConfigErrorCode}
		}

		//TODO: validate for usage of execution segment
		//TODO: validate for externally controlled executor (i.e. executors that aren't distributable)
		//TODO: move those validations to a separate function and reuse validateConfig()?

		err = r.SetOptions(conf.Options)
		if err != nil {
			return err
		}

		// Cloud config
		cloudConfig := cloud.NewConfig().Apply(derivedConf.Collectors.Cloud)
		if err := envconfig.Process("k6", &cloudConfig); err != nil {
			return err
		}
		if !cloudConfig.Token.Valid {
			return errors.New("Not logged in, please use `k6 login cloud`.")
		}

		printBar(progressBar, "Building the archive")
		arc := r.MakeArchive()
		// TODO: Fix this
		// We reuse cloud.Config for parsing options.ext.loadimpact, but this probably shouldn't be
		// done as the idea of options.ext is that they are extensible without touching k6. But in
		// order for this to happen we shouldn't actually marshall cloud.Config on top of it because
		// it will be missing some fields that aren't actually mentioned in the struct.
		// So in order for use to copy the fields that we need for loadimpact's api we unmarshal in
		// map[string]interface{} and copy what we need if it isn't set already
		var tmpCloudConfig map[string]interface{}
		if val, ok := arc.Options.External["loadimpact"]; ok {
			var dec = json.NewDecoder(bytes.NewReader(val))
			dec.UseNumber() // otherwise float64 are used
			if err := dec.Decode(&tmpCloudConfig); err != nil {
				return err
			}
		}

		if err := cloud.MergeFromExternal(arc.Options.External, &cloudConfig); err != nil {
			return err
		}
		if tmpCloudConfig == nil {
			tmpCloudConfig = make(map[string]interface{}, 3)
		}

		if _, ok := tmpCloudConfig["token"]; !ok && cloudConfig.Token.Valid {
			tmpCloudConfig["token"] = cloudConfig.Token
		}
		if _, ok := tmpCloudConfig["name"]; !ok && cloudConfig.Name.Valid {
			tmpCloudConfig["name"] = cloudConfig.Name
		}
		if _, ok := tmpCloudConfig["projectID"]; !ok && cloudConfig.ProjectID.Valid {
			tmpCloudConfig["projectID"] = cloudConfig.ProjectID
		}

		if arc.Options.External == nil {
			arc.Options.External = make(map[string]json.RawMessage)
		}
		arc.Options.External["loadimpact"], err = json.Marshal(tmpCloudConfig)
		if err != nil {
			return err
		}

		name := cloudConfig.Name.String
		if !cloudConfig.Name.Valid || cloudConfig.Name.String == "" {
			name = filepath.Base(filename)
		}

		// Start cloud test run
		printBar(progressBar, "Validating script options")
		client := cloud.NewClient(cloudConfig.Token.String, cloudConfig.Host.String, consts.Version)
		if err := client.ValidateOptions(arc.Options); err != nil {
			return err
		}

		printBar(progressBar, "Uploading archive")
		refID, err := client.StartCloudTestRun(name, cloudConfig.ProjectID.Int64, arc)
		if err != nil {
			return err
		}
		progressBar.Modify(pb.WithConstLeft("   Run"))
		printBar(progressBar, "Initializing the cloud test")

		testURL := cloud.URLForResults(refID, cloudConfig)
		fprintf(stdout, "\n\n")
		fprintf(stdout, "   executor: %s\n", ui.ValueColor.Sprint("cloud"))
		fprintf(stdout, "     script: %s\n", ui.ValueColor.Sprint(filename))
		fprintf(stdout, "     output: %s\n", ui.ValueColor.Sprint(testURL))
		//TODO: print executors information
		fprintf(stdout, "\n")
		printBar(progressBar, "Initializing the cloud test")

		// The quiet option hides the progress bar and disallow aborting the test
		if quiet {
			return nil
		}

		// Trap Interrupts, SIGINTs and SIGTERMs.
		sigC := make(chan os.Signal, 1)
		signal.Notify(sigC, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
		defer signal.Stop(sigC)

		var progressErr error
		testProgress := &cloud.TestProgressResponse{}
		percentageFmt := "[" + pb.GetFixedLengthFloatFormat(100, 2) + "%%] %s"
		progressBar.Modify(
			pb.WithProgress(func() (float64, string) {
				if testProgress.RunStatus < lib.RunStatusRunning {
					return 0, testProgress.RunStatusText
				}
				return testProgress.Progress, fmt.Sprintf(percentageFmt, testProgress.Progress*100, testProgress.RunStatusText)
			}),
		)

		ticker := time.NewTicker(time.Millisecond * 2000)
		shouldExitLoop := false

	runningLoop:
		for {
			select {
			case <-ticker.C:
				testProgress, progressErr = client.GetTestProgress(refID)
				if progressErr == nil {
					if (testProgress.RunStatus > lib.RunStatusRunning) || (exitOnRunning && testProgress.RunStatus == lib.RunStatusRunning) {
						shouldExitLoop = true
					}
					printBar(progressBar, "")
				} else {
					logrus.WithError(progressErr).Error("Test progress error")
				}
				if shouldExitLoop {
					break runningLoop
				}
			case sig := <-sigC:
				logrus.WithField("sig", sig).Print("Exiting in response to signal...")
				err := client.StopCloudTestRun(refID)
				if err != nil {
					logrus.WithError(err).Error("Stop cloud test error")
				}
				shouldExitLoop = true // Exit after the next GetTestProgress call
			}
		}

		if testProgress == nil {
			return ExitCode{errors.New("Test progress error"), 98}
		}

		fprintf(stdout, "     test status: %s\n", ui.ValueColor.Sprint(testProgress.RunStatusText))

		if testProgress.ResultStatus == cloud.ResultStatusFailed {
			return ExitCode{errors.New("The test has failed"), 99}
		}

		return nil
	},
}

func cloudCmdFlagSet() *pflag.FlagSet {
	flags := pflag.NewFlagSet("", pflag.ContinueOnError)
	flags.SortFlags = false
	flags.AddFlagSet(optionFlagSet())
	flags.AddFlagSet(runtimeOptionFlagSet(false))

	//TODO: Figure out a better way to handle the CLI flags:
	// - the default value is specified in this way so we don't overwrire whatever
	//   was specified via the environment variable
	// - global variables are not very testable... :/
	flags.BoolVar(&exitOnRunning, "exit-on-running", exitOnRunning, "exits when test reaches the running status")
	// We also need to explicitly set the default value for the usage message here, so setting
	// K6_EXIT_ON_RUNNING=true won't affect the usage message
	flags.Lookup("exit-on-running").DefValue = "false"

	return flags
}

func init() {
	RootCmd.AddCommand(cloudCmd)
	cloudCmd.Flags().SortFlags = false
	cloudCmd.Flags().AddFlagSet(cloudCmdFlagSet())
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"log"
	"time"

	"github.com/phil-mansfield/shellfish/cmd/catalog"
	"github.com/phil-mansfield/shellfish/cmd/env"
	"github.com/phil-mansfield/shellfish/cmd/halo"
	"github.com/phil-mansfield/shellfish/cmd/memo"
	"github.com/phil-mansfield/shellfish/io"
	"github.com/phil-mansfield/shellfish/logging"
	"github.com/phil-mansfield/shellfish/parse"
)

type CoordConfig struct {
	values []string
}

var _ Mode = &CoordConfig{}

func (config *CoordConfig) ExampleConfig() string {
	return`[config.coord]
# Values are the names of the values you want to write to an output catalog.
# The default order is the one which is needed by Shellfish. Any other order
# would correspond to a catalog which is for your personal use only.
Values = X, Y, Z, R200m
`
}

func (config *CoordConfig) ReadConfig(fname string, flags []string) error {
	vars := parse.NewConfigVars("coord.config")
	vars.Strings(&config.values, "Values", []string{"X", "Y", "Z", "R200m"})

	if fname == "" {
		if len(flags) == 0 {
			return nil
		}
		return parse.ReadFlags(flags, vars)
	}
	if err := parse.ReadConfig(fname, vars); err != nil {
		return err
	}
	return parse.ReadFlags(flags, vars)
}

func (config *CoordConfig) validate(vars *halo.VarColumns) error {
	for _, val := range config.values {
		if _, ok := vars.ColumnLookup[val]; !ok {
			return fmt.Errorf(
				"Value '%s' requested by coord mode, but isn't " +
				"in HaloVaueNames.", val,
			)
		}
	}
	return nil
}

func (config *CoordConfig) Run(
	gConfig *GlobalConfig, e *env.Environment, stdin []byte,
) ([]string, error) {

	if logging.Mode != logging.Nil {
		log.Println(`
#####################
## shellfish coord ##
#####################`,
		)
	}
	var t time.Time
	if logging.Mode == logging.Performance {
		t = time.Now()
	}

	intCols, _, err := catalog.Parse(stdin, []int{0, 1}, []int{})
	if err != nil {
		return nil, err
	}
	ids, snaps := intCols[0], intCols[1]

	if len(ids) == 0 {
		return nil, fmt.Errorf("In input IDs.")
	}

	vars := halo.NewVarColumns(
		gConfig.HaloValueNames, gConfig.HaloValueColumns,
		gConfig.HaloRadiusUnits,
	)
	if err := config.validate(vars); err != nil {
		return nil, err
	}

	buf, err := getVectorBuffer(
		e.ParticleCatalog(snaps[0], 0),
		gConfig.SnapshotType, gConfig.Endianness,
		gConfig.GadgetNpartNum,
	)
	if err != nil {
		return nil, err
	}

	cols, err := readHaloCoords(
		ids, snaps, config.values, vars, buf, e, gConfig,
	)
	if err != nil {
		return nil, err
	}

	icols := [][]int{ids, snaps}
	fcols := [][]float64{}
	icolOrder := []int{}
	fcolOrder := []int{}

	intNum := 0
	for _, valueName := range config.values {
		j := findString(valueName, gConfig.HaloValueNames)
		comment := gConfig.HaloValueComments[j]
		if isIntType(comment) { intNum++ }
	}

	colOrder := append([]int{0, 1}, make([]int, len(cols))...)
	for i, valueName := range config.values {
		j := findString(valueName, gConfig.HaloValueNames)
		comment := gConfig.HaloValueComments[j]

		if isIntType(comment) {
			fcol := cols[i]
			icol := make([]int, len(fcol))
			for i := range icol { icol[i] = int(fcol[i]) }
			icols = append(icols, icol)
			icolOrder = append(icolOrder, i + 2)
		} else {
			fcols = append(fcols, cols[i])
			fcolOrder = append(fcolOrder, i + 2)
		}
	}

	colOrder = append(icolOrder, fcolOrder...)
	lines := catalog.FormatCols(icols, fcols, colOrder)

	cString := makeCommentString(gConfig, config)

	if logging.Mode == logging.Performance {
		log.Printf("Time: %s", time.Since(t).String())
		log.Printf("Memory:\n%s", logging.MemString())
	}

	return append([]string{cString}, lines...), nil
}

func isIntType(comment string) bool {
	return comment == "int" || comment == "\"int\""
}

func makeCommentString(gConfig *GlobalConfig, config *CoordConfig) string {
	colNames := make([]string, len(config.values))
	for i := 0; i < len(config.values); i++ {
		switch config.values[i] {
		case "R200m", "R200c", "R500c", "Rs":
			colNames[i] = fmt.Sprintf(
				"%s [%s]", config.values[i], gConfig.HaloPositionUnits,
			)
			continue
		}

		j := findString(config.values[i], gConfig.HaloValueNames)
		if gConfig.HaloValueComments[j] == "" ||
			gConfig.HaloValueComments[j] == "\"\"" ||
			isIntType(gConfig.HaloValueComments[j]) {

			colNames[i] = config.values[i]
		} else {
			colNames[i] = fmt.Sprintf(
				"%s [%s]", config.values[i], gConfig.HaloValueComments[j],
			)
		}
	}

	colOrder := make([]int, 2 + len(config.values))
	colSizes := make([]int, 2 + len(config.values))
	for i := range colOrder {
		colOrder[i], colSizes[i] = i, 1
	}

	return catalog.CommentString(
		[]string{"ID", "Snapshot"}, colNames, colOrder, colSizes,
	)
}

func findString(x string, xs []string) int {
	for i := range xs {
		if xs[i] == x { return i }
	}
	panic("Impossible")
}

func readHaloCoords(
	ids, snaps []int, valNames []string, vars *halo.VarColumns,
	buf io.VectorBuffer, e *env.Environment, gConfig *GlobalConfig,
) (cols [][]float64, err error) {
	if len(snaps) == 0 { return nil, nil }

	snapBins, idxBins := binBySnap(snaps, ids)

	cols = make([][]float64, len(valNames))
	for i := range cols {
		cols[i] = make([]float64, len(ids))
	}

	for snap, _ := range snapBins {
		if snap == -1 {
			continue
		}

		hds, _, err := memo.ReadHeaders(snaps[0], buf, e)
		if err != nil { return nil, err }
		cosmo := &hds[0].Cosmo

		snapIDs := snapBins[snap]
		idxs := idxBins[snap]

		_, scols, err := memo.ReadRockstar(
			snap, valNames, snapIDs, vars, buf, e,
		)

		if err != nil {
			return nil, err
		}

		for i := range valNames {
			switch valNames[i] {
			case "X", "Y", "Z":
				ucf := halo.UnitConversionFactor(
					gConfig.HaloPositionUnits, cosmo,
				)
				for j := range scols[i] {
					scols[i][j] *= ucf
				}
			case "R200m", "R200c", "R500c", "Rs" :
				ucf := halo.UnitConversionFactor(
					gConfig.HaloRadiusUnits, cosmo,
				)
				for j := range scols[i] {
					scols[i][j] *= ucf
				}
			}
		}

		for i, idx := range idxs {
			for j := range cols {
				cols[j][idx] = scols[j][i]
			}
		}
	}

	return cols, nil
}
<commit_msg>Fixed coord ID writing.<commit_after>package cmd

import (
	"fmt"
	"log"
	"time"

	"github.com/phil-mansfield/shellfish/cmd/catalog"
	"github.com/phil-mansfield/shellfish/cmd/env"
	"github.com/phil-mansfield/shellfish/cmd/halo"
	"github.com/phil-mansfield/shellfish/cmd/memo"
	"github.com/phil-mansfield/shellfish/io"
	"github.com/phil-mansfield/shellfish/logging"
	"github.com/phil-mansfield/shellfish/parse"
)

type CoordConfig struct {
	values []string
}

var _ Mode = &CoordConfig{}

func (config *CoordConfig) ExampleConfig() string {
	return`[config.coord]
# Values are the names of the values you want to write to an output catalog.
# The default order is the one which is needed by Shellfish. Any other order
# would correspond to a catalog which is for your personal use only.
Values = X, Y, Z, R200m
`
}

func (config *CoordConfig) ReadConfig(fname string, flags []string) error {
	vars := parse.NewConfigVars("coord.config")
	vars.Strings(&config.values, "Values", []string{"X", "Y", "Z", "R200m"})

	if fname == "" {
		if len(flags) == 0 {
			return nil
		}
		return parse.ReadFlags(flags, vars)
	}
	if err := parse.ReadConfig(fname, vars); err != nil {
		return err
	}
	return parse.ReadFlags(flags, vars)
}

func (config *CoordConfig) validate(vars *halo.VarColumns) error {
	for _, val := range config.values {
		if _, ok := vars.ColumnLookup[val]; !ok {
			return fmt.Errorf(
				"Value '%s' requested by coord mode, but isn't " +
				"in HaloVaueNames.", val,
			)
		}
	}
	return nil
}

func (config *CoordConfig) Run(
	gConfig *GlobalConfig, e *env.Environment, stdin []byte,
) ([]string, error) {

	if logging.Mode != logging.Nil {
		log.Println(`
#####################
## shellfish coord ##
#####################`,
		)
	}
	var t time.Time
	if logging.Mode == logging.Performance {
		t = time.Now()
	}

	intCols, _, err := catalog.Parse(stdin, []int{0, 1}, []int{})
	if err != nil {
		return nil, err
	}
	ids, snaps := intCols[0], intCols[1]

	if len(ids) == 0 {
		return nil, fmt.Errorf("In input IDs.")
	}

	vars := halo.NewVarColumns(
		gConfig.HaloValueNames, gConfig.HaloValueColumns,
		gConfig.HaloRadiusUnits,
	)
	if err := config.validate(vars); err != nil {
		return nil, err
	}

	buf, err := getVectorBuffer(
		e.ParticleCatalog(snaps[0], 0),
		gConfig.SnapshotType, gConfig.Endianness,
		gConfig.GadgetNpartNum,
	)
	if err != nil {
		return nil, err
	}

	cols, err := readHaloCoords(
		ids, snaps, config.values, vars, buf, e, gConfig,
	)
	if err != nil {
		return nil, err
	}

	icols := [][]int{ids, snaps}
	fcols := [][]float64{}
	icolOrder := []int{0, 1}
	fcolOrder := []int{}

	intNum := 0
	for _, valueName := range config.values {
		j := findString(valueName, gConfig.HaloValueNames)
		comment := gConfig.HaloValueComments[j]
		if isIntType(comment) { intNum++ }
	}

	for i, valueName := range config.values {
		j := findString(valueName, gConfig.HaloValueNames)
		comment := gConfig.HaloValueComments[j]

		if isIntType(comment) {
			fcol := cols[i]
			icol := make([]int, len(fcol))
			for i := range icol { icol[i] = int(fcol[i]) }
			icols = append(icols, icol)
			icolOrder = append(icolOrder, i + 2)
		} else {
			fcols = append(fcols, cols[i])
			fcolOrder = append(fcolOrder, i + 2)
		}
	}

	colOrder := append(icolOrder, fcolOrder...)
	lines := catalog.FormatCols(icols, fcols, colOrder)

	cString := makeCommentString(gConfig, config)

	if logging.Mode == logging.Performance {
		log.Printf("Time: %s", time.Since(t).String())
		log.Printf("Memory:\n%s", logging.MemString())
	}

	return append([]string{cString}, lines...), nil
}

func isIntType(comment string) bool {
	return comment == "int" || comment == "\"int\""
}

func makeCommentString(gConfig *GlobalConfig, config *CoordConfig) string {
	colNames := make([]string, len(config.values))
	for i := 0; i < len(config.values); i++ {
		switch config.values[i] {
		case "R200m", "R200c", "R500c", "Rs":
			colNames[i] = fmt.Sprintf(
				"%s [%s]", config.values[i], gConfig.HaloPositionUnits,
			)
			continue
		}

		j := findString(config.values[i], gConfig.HaloValueNames)
		if gConfig.HaloValueComments[j] == "" ||
			gConfig.HaloValueComments[j] == "\"\"" ||
			isIntType(gConfig.HaloValueComments[j]) {

			colNames[i] = config.values[i]
		} else {
			colNames[i] = fmt.Sprintf(
				"%s [%s]", config.values[i], gConfig.HaloValueComments[j],
			)
		}
	}

	colOrder := make([]int, 2 + len(config.values))
	colSizes := make([]int, 2 + len(config.values))
	for i := range colOrder {
		colOrder[i], colSizes[i] = i, 1
	}

	return catalog.CommentString(
		[]string{"ID", "Snapshot"}, colNames, colOrder, colSizes,
	)
}

func findString(x string, xs []string) int {
	for i := range xs {
		if xs[i] == x { return i }
	}
	panic("Impossible")
}

func readHaloCoords(
	ids, snaps []int, valNames []string, vars *halo.VarColumns,
	buf io.VectorBuffer, e *env.Environment, gConfig *GlobalConfig,
) (cols [][]float64, err error) {
	if len(snaps) == 0 { return nil, nil }

	snapBins, idxBins := binBySnap(snaps, ids)

	cols = make([][]float64, len(valNames))
	for i := range cols {
		cols[i] = make([]float64, len(ids))
	}

	for snap, _ := range snapBins {
		if snap == -1 {
			continue
		}

		hds, _, err := memo.ReadHeaders(snaps[0], buf, e)
		if err != nil { return nil, err }
		cosmo := &hds[0].Cosmo

		snapIDs := snapBins[snap]
		idxs := idxBins[snap]

		_, scols, err := memo.ReadRockstar(
			snap, valNames, snapIDs, vars, buf, e,
		)

		if err != nil {
			return nil, err
		}

		for i := range valNames {
			switch valNames[i] {
			case "X", "Y", "Z":
				ucf := halo.UnitConversionFactor(
					gConfig.HaloPositionUnits, cosmo,
				)
				for j := range scols[i] {
					scols[i][j] *= ucf
				}
			case "R200m", "R200c", "R500c", "Rs" :
				ucf := halo.UnitConversionFactor(
					gConfig.HaloRadiusUnits, cosmo,
				)
				for j := range scols[i] {
					scols[i][j] *= ucf
				}
			}
		}

		for i, idx := range idxs {
			for j := range cols {
				cols[j][idx] = scols[j][i]
			}
		}
	}

	return cols, nil
}
<|endoftext|>


<commit_before>package cmd

import (
	"encoding/json"
	"errors"
	log "github.com/Sirupsen/logrus"
	"github.com/gocraft/health"
	"gopkg.in/dougEfresh/dbr.v2"
	"time"
)

type eventRecorder interface {
	recordEvent(event *Event) error
	resolveGeoEvent(event *Event) error
	get(id int64) *EventGeo
}

type eventLister interface {
	list() []EventGeo
}

type eventTransporter interface {
	eventLister
	eventRecorder
}

type eventClient struct {
	db        *dbr.Connection
	geoClient geoClientTransporter
}

var defaultEventClient *eventClient

func (c *eventClient) list() []EventGeo {
	sess := c.db.NewSession(nil)
	after := time.Now().UTC().AddDate(0, 0, -1)
	var geoEvents []EventGeo
	_, err := sess.Select("*").
		From("vw_event").
		Where("dt > ?", after).
		Limit(1000).
		LoadValues(&geoEvents)
	if err != nil {
		log.Errorf("Error getting events %s", err)
	}
	return geoEvents
}

func (c *eventClient) recordEvent(event *Event) error {
	log.Infof("Processing event %+v", event)
	job := stream.NewJob("record_event")
	if c.db == nil {
		return nil
	}
	sess := c.db.NewSession(nil)
	var ids []int64
	_, err := sess.InsertInto("event").
		Columns("dt", "username", "passwd", "remote_addr", "remote_port", "remote_name", "remote_version", "origin_addr", "application", "protocol").
		Record(event).
		Returning(&ids, "id")

	if err != nil {
		job.Complete(health.Error)
		return err
	}
	event.ID = ids[0]
	job.Complete(health.Success)
	return nil
}

func (c *eventClient) resolveGeoEvent(event *Event) error {
	job := stream.NewJob("resolve_geo_event")
	if event.ID == 0 {
		err := errors.New("Bad event recv")
		log.Errorf("Got bad event %s", event)
		job.EventErr("resolve_geo_event_invalid", err)
		job.Complete(health.ValidationError)
		return err
	}

	sess := c.db.NewSession(nil)
	geo, err := c.resolveAddr(event.RemoteAddr)
	if err != nil {
		log.Errorf("Error geting location for RemoteAddr %+v %s", event, err)
		job.Complete(health.ValidationError)
		return err
	}
	updateBuilder := sess.Update("event").Set("remote_geo_id", geo.ID).Where("id = ?", event.ID)
	if _, err = updateBuilder.Exec(); err != nil {
		log.Errorf("Error updating remote_addr_geo_id for id %d %s", event.ID, err)
		job.Complete(health.Error)
		return err
	}

	geo, err = c.resolveAddr(event.OriginAddr)
	if err != nil {
		log.Errorf("Errro getting location for origin %+v %s", event, err)
		job.Complete(health.Error)
		return err
	}
	updateBuilder = sess.Update("event").Set("origin_geo_id", geo.ID).Where("id = ?", event.ID)
	if _, err = updateBuilder.Exec(); err != nil {
		log.Errorf("Error updating origin for id %d %s", event.ID, err)
		job.Complete(health.Error)
		return err
	}
	job.Complete(health.Success)
	go c.broadcastEvent(event.ID)
	return nil
}

func (c *eventClient) broadcastEvent(id int64) {
	gEvent := c.get(id)
	if gEvent == nil {
		return
	}
	if b, err := json.Marshal(gEvent); err != nil {
		log.Errorf("Error decoding geo event %d %s", id, err)
	} else {
		hub.broadcast <- b
	}
}

func (c *eventClient) get(id int64) *EventGeo {
	job := stream.NewJob("get_event")
	sess := c.db.NewSession(nil)
	var event EventGeo
	if _, err := sess.Select("*").
		From("vw_event").
		Where("id = ?", id).
		Load(&event); err != nil {
		log.Errorf("Error getting event id %d %s", id, err)
		job.Complete(health.Error)
		return nil
	}
	job.Complete(health.Success)
	return &event
}
<commit_msg>moving get(id) to eventLister<commit_after>package cmd

import (
	"encoding/json"
	"errors"
	log "github.com/Sirupsen/logrus"
	"github.com/gocraft/health"
	"gopkg.in/dougEfresh/dbr.v2"
	"time"
)

type eventRecorder interface {
	recordEvent(event *Event) error
	resolveGeoEvent(event *Event) error
}

type eventLister interface {
	list() []EventGeo
	get(id int64) *EventGeo
}

type eventTransporter interface {
	eventLister
	eventRecorder
}

type eventClient struct {
	db        *dbr.Connection
	geoClient geoClientTransporter
}

var defaultEventClient *eventClient

func (c *eventClient) list() []EventGeo {
	sess := c.db.NewSession(nil)
	after := time.Now().UTC().AddDate(0, 0, -1)
	var geoEvents []EventGeo
	_, err := sess.Select("*").
		From("vw_event").
		Where("dt > ?", after).
		Limit(1000).
		LoadValues(&geoEvents)
	if err != nil {
		log.Errorf("Error getting events %s", err)
	}
	return geoEvents
}

func (c *eventClient) recordEvent(event *Event) error {
	log.Infof("Processing event %+v", event)
	job := stream.NewJob("record_event")
	if c.db == nil {
		return nil
	}
	sess := c.db.NewSession(nil)
	var ids []int64
	_, err := sess.InsertInto("event").
		Columns("dt", "username", "passwd", "remote_addr", "remote_port", "remote_name", "remote_version", "origin_addr", "application", "protocol").
		Record(event).
		Returning(&ids, "id")

	if err != nil {
		job.Complete(health.Error)
		return err
	}
	event.ID = ids[0]
	job.Complete(health.Success)
	return nil
}

func (c *eventClient) resolveGeoEvent(event *Event) error {
	job := stream.NewJob("resolve_geo_event")
	if event.ID == 0 {
		err := errors.New("Bad event recv")
		log.Errorf("Got bad event %s", event)
		job.EventErr("resolve_geo_event_invalid", err)
		job.Complete(health.ValidationError)
		return err
	}

	sess := c.db.NewSession(nil)
	geo, err := c.resolveAddr(event.RemoteAddr)
	if err != nil {
		log.Errorf("Error geting location for RemoteAddr %+v %s", event, err)
		job.Complete(health.ValidationError)
		return err
	}
	updateBuilder := sess.Update("event").Set("remote_geo_id", geo.ID).Where("id = ?", event.ID)
	if _, err = updateBuilder.Exec(); err != nil {
		log.Errorf("Error updating remote_addr_geo_id for id %d %s", event.ID, err)
		job.Complete(health.Error)
		return err
	}

	geo, err = c.resolveAddr(event.OriginAddr)
	if err != nil {
		log.Errorf("Errro getting location for origin %+v %s", event, err)
		job.Complete(health.Error)
		return err
	}
	updateBuilder = sess.Update("event").Set("origin_geo_id", geo.ID).Where("id = ?", event.ID)
	if _, err = updateBuilder.Exec(); err != nil {
		log.Errorf("Error updating origin for id %d %s", event.ID, err)
		job.Complete(health.Error)
		return err
	}
	job.Complete(health.Success)
	go c.broadcastEvent(event.ID)
	return nil
}

func (c *eventClient) broadcastEvent(id int64) {
	gEvent := c.get(id)
	if gEvent == nil {
		return
	}
	if b, err := json.Marshal(gEvent); err != nil {
		log.Errorf("Error decoding geo event %d %s", id, err)
	} else {
		hub.broadcast <- b
	}
}

func (c *eventClient) get(id int64) *EventGeo {
	job := stream.NewJob("get_event")
	sess := c.db.NewSession(nil)
	var event EventGeo
	if _, err := sess.Select("*").
		From("vw_event").
		Where("id = ?", id).
		Load(&event); err != nil {
		log.Errorf("Error getting event id %d %s", id, err)
		job.Complete(health.Error)
		return nil
	}
	job.Complete(health.Success)
	return &event
}
<|endoftext|>


<commit_before>// Copyright © 2017 Marcus Franke <marcus.franke@gmail.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"sync"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/julienschmidt/httprouter"
	"gopkg.in/russross/blackfriday.v2"
)

var (
	port   string
	debug  bool
	cmdOut []byte
	mutex  sync.Mutex
)

// serveCmd represents the serve command
var serveCmd = &cobra.Command{
	Use:   "serve",
	Short: "Starts the yummy webserver",
	Run: func(cmd *cobra.Command, args []string) {

		repoPath := viper.GetString("yum.repopath")

		router:= httprouter.New()
		router.Handler("GET", "/", http.FileServer(http.Dir(repoPath)))

		router.GET("/help", helpHandler)
		router.POST("/api/upload", apiUploadHandler)
		//router.PUT("/api/upload/:filename", apiUploadPut)
		//router.DELETE("/api/delete/:name", apiDeleteHandler)

		log.Fatal(http.ListenAndServe(":8080", router))
	},
}

func helpHandler(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	// get helpFile path from configuration
	helpFile := viper.GetString("yum.helpFile")

	// ingest the configured helpFile
	help, err := ioutil.ReadFile(helpFile)
	if err != nil {
		http.Error(w, "Could not load the help file", http.StatusInternalServerError)
		return
	}

	// render the Markdown file to HTML using the
	// blackfriday library
	output := blackfriday.Run(help)
	fmt.Fprintf(w, string(output))
}

func apiUploadHandler(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {

	repoPath := viper.GetString("yum.repopath")
	workers := viper.GetString("yum.workers")
	createrepoBinary := viper.GetString("yum.createrepoBinary")

	if debug {
		fmt.Println("Method:", r.Method)
		fmt.Println("Header:", r.Header)
		fmt.Println("repoPath:", repoPath)
	}

	// will handle file uploads
	if r.Method == "POST" {
		file, handler, err := r.FormFile("fileupload")
		if err != nil {
			http.Error(w, "FormFile does not match - use fileupload\n", http.StatusBadRequest)
			return
		}
		defer file.Close()

		if filepath.Ext(handler.Filename) != ".rpm" {
			http.Error(w, "File not RPM\n", http.StatusUnsupportedMediaType)
			return
		}

		// check if the uploaded file already exists
		// if the repository is configured in protected mode
		// the request will return status 403 (forbidden)
		if viper.GetBool("yum.protected") {
			if _, err := os.Stat(repoPath + "/" + handler.Filename); err == nil {
				http.Error(w, "File already exists, forbidden to overwrite!\n",
					http.StatusForbidden)
				log.Println(err)
				return
			}
			log.Println("File already exists, will overwrite: " + handler.Filename)
		}

		// create file handler to write uploaded file to
		f, err := os.OpenFile(repoPath+"/"+handler.Filename, os.O_WRONLY|os.O_CREATE, 0644)
		if err != nil {
			http.Error(w, "An error occurred", http.StatusInternalServerError)
			log.Fatal(err)
		}
		defer f.Close()

		// copy the file buffer into the file handle
		_, err = io.Copy(f, file)
		if err != nil {
			http.Error(w, "An error occurred applying the upload to the filesystem", http.StatusInternalServerError)
			log.Println(err)
			return
		}

		// process the uploaded file
		mutex.Lock()
		cmdOut, err = exec.Command(createrepoBinary, "--update", "--workers", workers, repoPath).CombinedOutput()
		if err != nil {
			fmt.Fprintln(w, string(cmdOut))
			http.Error(w, "Could not update repository", http.StatusInternalServerError)
			log.Println(cmdOut, err)
			mutex.Unlock()
			return
		}
		log.Println(string(cmdOut))
		mutex.Unlock()

	}

	// assume curl --upload-file style of upload type
	// this is currently not supported
	if r.Method == "PUT" {
		http.Error(w, "Method not allowed, POST binary to URI\n", http.StatusMethodNotAllowed)
	}
}

func init() {
	RootCmd.AddCommand(serveCmd)

	// Flags for the serve command.
	serveCmd.Flags().StringVarP(&port, "port", "p", "8080", "Port to listen on")
	serveCmd.Flags().BoolVarP(&debug, "debug", "d", false, "Enable debug output")
}
<commit_msg>[serve] refactored ApiUploadHandler<commit_after>// Copyright © 2017 Marcus Franke <marcus.franke@gmail.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"sync"

	"github.com/julienschmidt/httprouter"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"gopkg.in/russross/blackfriday.v2"
)

var (
	port   string
	debug  bool
	cmdOut []byte
	mutex  sync.Mutex
)

// serveCmd represents the serve command
var serveCmd = &cobra.Command{
	Use:   "serve",
	Short: "Starts the yummy webserver",
	Run: func(cmd *cobra.Command, args []string) {

		repoPath := viper.GetString("yum.repopath")

		router := httprouter.New()
		router.Handler("GET", "/", http.FileServer(http.Dir(repoPath)))

		router.GET("/help", helpHandler)
		router.POST("/api/upload", apiPostUploadHandler)
		//router.PUT("/api/upload/:filename", apiUploadPut)
		//router.DELETE("/api/delete/:name", apiDeleteHandler)

		log.Fatal(http.ListenAndServe(":8080", router))
	},
}

func helpHandler(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	// get helpFile path from configuration
	helpFile := viper.GetString("yum.helpFile")

	// ingest the configured helpFile
	help, err := ioutil.ReadFile(helpFile)
	if err != nil {
		http.Error(w, "Could not load the help file", http.StatusInternalServerError)
		return
	}

	// render the Markdown file to HTML using the
	// blackfriday library
	output := blackfriday.Run(help)
	fmt.Fprintf(w, string(output))
}

func apiPostUploadHandler(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {

	repoPath := viper.GetString("yum.repopath")
	workers := viper.GetString("yum.workers")
	createrepoBinary := viper.GetString("yum.createrepoBinary")

	if debug {
		fmt.Println("Method:", r.Method)
		fmt.Println("Header:", r.Header)
		fmt.Println("repoPath:", repoPath)
	}

	file, handler, err := r.FormFile("fileupload")
	if err != nil {
		http.Error(w, "FormFile does not match - use fileupload\n", http.StatusBadRequest)
		return
	}
	defer file.Close()

	if filepath.Ext(handler.Filename) != ".rpm" {
		http.Error(w, "File not RPM\n", http.StatusUnsupportedMediaType)
		return
	}

	// check if the uploaded file already exists
	// if the repository is configured in protected mode
	// the request will return status 403 (forbidden)
	if viper.GetBool("yum.protected") {
		if _, err := os.Stat(repoPath + "/" + handler.Filename); err == nil {
			http.Error(w, "File already exists, forbidden to overwrite!\n",
				http.StatusForbidden)
			log.Println(err)
			return
		}
		log.Println("File already exists, will overwrite: " + handler.Filename)
	}

	// create file handler to write uploaded file to
	f, err := os.OpenFile(repoPath+"/"+handler.Filename, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		http.Error(w, "An error occurred", http.StatusInternalServerError)
		log.Fatal(err)
	}
	defer f.Close()

	// copy the file buffer into the file handle
	_, err = io.Copy(f, file)
	if err != nil {
		http.Error(w, "An error occurred applying the upload to the filesystem", http.StatusInternalServerError)
		log.Println(err)
		return
	}

	// process the uploaded file
	mutex.Lock()
	cmdOut, err = exec.Command(createrepoBinary, "--update", "--workers", workers, repoPath).CombinedOutput()
	if err != nil {
		fmt.Fprintln(w, string(cmdOut))
		http.Error(w, "Could not update repository", http.StatusInternalServerError)
		log.Println(cmdOut, err)
		mutex.Unlock()
		return
	}
	log.Println(string(cmdOut))
	mutex.Unlock()
}

func init() {
	RootCmd.AddCommand(serveCmd)

	// Flags for the serve command.
	serveCmd.Flags().StringVarP(&port, "port", "p", "8080", "Port to listen on")
	serveCmd.Flags().BoolVarP(&debug, "debug", "d", false, "Enable debug output")
}
<|endoftext|>


<commit_before>package scummatlas

import (
	"fmt"
)

type Image bool
type Script bool
type Box struct {
    ulx int
    uly int
    urx int
    ury int
    lrx int
    lry int
    llx int
    lly int
    mask byte
    flags byte
    scale int
}
type BoxMatrix bool

type Room struct {
	data     []byte
	offset   int
	Width    int
	Height   int
	ObjCount int
	//ColorCycle ColorCycle
	//TranspColor TranspColor
	//Palette Palette
	Image         Image
	ObjectImage   Image
	ObjectScripts []Script
	ExitScript    Script
	EntryScript   Script
	LocalScript   Script
	Boxes       []Box
	BoxMatrix     BoxMatrix
}

func NewRoom(data []byte) *Room {
	room := new(Room)
	room.data = data
	room.offset = 0

	blockName := room.getBlockName()
	if blockName != "ROOM" {
		panic("Can't find ROOM")
	}

	room.offset = 8
	for room.offset < len(data) {
		blockName := room.getBlockName()
		fmt.Println("Parsing", blockName)

		switch blockName {
		case "RMHD":
			room.parseRMHD()
		case "BOXD":
			room.parseBOXD()
		}
		case "RMIM":
			room.parseRMIM()
		}

		room.nextBlock()
	}

	fmt.Println("New ROOM\n")
	room.Print()
	return room
}

func (r *Room) parseBOXD() {
    boxCount := LE16(r.data, r.offset + 8)
    var boxOffset int
    for i := 0 ; i < boxCount ; i ++ {
        boxOffset = 10 + i * 20
        box := NewBox(r.data[boxOffset:boxOffset + 20])
	    r.Boxes = append(r.Boxes, box)
    }
}

func (r *Room) parseRMHD() {
	fmt.Println("RMHD offset", r.offset)
	r.Width = LE16(r.data, r.offset+8)
	r.Height = LE16(r.data, r.offset+10)
	r.ObjCount = LE16(r.data, r.offset+12)
}

func (r Room) Print() {
	fmt.Println("Size: ", r.Width, r.Height)
	fmt.Println("Object count: ", r.ObjCount)
	fmt.Println("Boxes: ", len(r.Boxes))
}

func (r Room) getBlockName() string {
	return string(r.data[r.offset : r.offset+4])
}

func (r *Room) nextBlock() {
	blockSize := BE32(r.data, r.offset+4)
	r.offset += blockSize
}

func NewBox(data []byte) Box {
    box := new(Box)

    box.ulx = LE16(data, 0)
    box.uly = LE16(data, 2)
    box.urx = LE16(data, 4)
    box.ury = LE16(data, 6)
    box.lrx = LE16(data, 8)
    box.lry = LE16(data, 10)
    box.llx = LE16(data, 12)
    box.lly = LE16(data, 14)
    box.mask = data[16]
    box.flags = data[17]
    box.scale = LE16(data, 18)

    return *box
}
<commit_msg>Add initial script parsing<commit_after>package scummatlas

import (
	"fmt"
)

type Image bool
type Box struct {
	ulx   int
	uly   int
	urx   int
	ury   int
	lrx   int
	lry   int
	llx   int
	lly   int
	mask  byte
	flags byte
	scale int
}
type Script string
type BoxMatrix bool

type Room struct {
	data     []byte
	offset   int
	Width    int
	Height   int
	ObjCount int
	//ColorCycle ColorCycle
	//TranspColor TranspColor
	//Palette Palette
	Image         Image
	ObjectImage   Image
	ObjectScripts []Script
	ExitScript    Script
	EntryScript   Script
	Boxes         []Box
	LocalScripts  []Script
	BoxMatrix     BoxMatrix
}

func NewRoom(data []byte) *Room {
	room := new(Room)
	room.data = data
	room.offset = 0

	blockName := room.getBlockName()
	if blockName != "ROOM" {
		panic("Can't find ROOM")
	}

	room.offset = 8
	for room.offset < len(data) {
		blockName := room.getBlockName()
		fmt.Println("Parsing", blockName)

		switch blockName {
		case "RMHD":
			room.parseRMHD()
		case "BOXD":
			room.parseBOXD()
		case "EXCD":
			room.parseEXCD()
		case "ENCD":
			room.parseENCD()
		case "LSCR":
			room.parseLSCR()
		case "RMIM":
			room.parseRMIM()
		}

		room.nextBlock()
	}

	fmt.Println("New ROOM\n")
	room.Print()
	return room
}

func parseScriptBlock(data []byte) Script {
	fmt.Println("Script size", BE32(data, 4))
	return ""
}

func (r *Room) parseLSCR() {
	script := parseScriptBlock(
		r.data[r.offset : r.offset+r.getBlockSize()])
	r.LocalScripts = append(r.LocalScripts, script)
}

func (r *Room) parseENCD() {
	r.EntryScript = parseScriptBlock(r.data[r.offset : r.offset+r.getBlockSize()])
}

func (r *Room) parseEXCD() {
	r.EntryScript = parseScriptBlock(r.data[r.offset : r.offset+r.getBlockSize()])
}

func (r *Room) parseBOXD() {
	boxCount := LE16(r.data, r.offset+8)
	var boxOffset int
	for i := 0; i < boxCount; i++ {
		boxOffset = 10 + i*20
		box := NewBox(r.data[boxOffset : boxOffset+20])
		r.Boxes = append(r.Boxes, box)
	}
}

func (r *Room) parseRMHD() {
	fmt.Println("RMHD offset", r.offset)
	r.Width = LE16(r.data, r.offset+8)
	r.Height = LE16(r.data, r.offset+10)
	r.ObjCount = LE16(r.data, r.offset+12)
}

func (r Room) Print() {
	fmt.Println("Size: ", r.Width, r.Height)
	fmt.Println("Object count: ", r.ObjCount)
	fmt.Println("Boxes: ", len(r.Boxes))
}

func (r Room) getBlockName() string {
	return string(r.data[r.offset : r.offset+4])
}

func (r Room) getBlockSize() int {
	return BE32(r.data, r.offset+4)
}

func (r *Room) nextBlock() {
	r.offset += r.getBlockSize()
}

func NewBox(data []byte) Box {
	box := new(Box)

	box.ulx = LE16(data, 0)
	box.uly = LE16(data, 2)
	box.urx = LE16(data, 4)
	box.ury = LE16(data, 6)
	box.lrx = LE16(data, 8)
	box.lry = LE16(data, 10)
	box.llx = LE16(data, 12)
	box.lly = LE16(data, 14)
	box.mask = data[16]
	box.flags = data[17]
	box.scale = LE16(data, 18)

	return *box
}
<|endoftext|>