<commit_before>/**
 * Copyright 2017 IBM Corp.
 *
 * 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 database

import (
	"fmt"
	"github.com/IBM/ubiquity/utils/logs"
)

var migrations = make(map[string]interface{})

func RegisterMigration(obj interface{}) {
	defer logs.GetLogger().Trace(logs.DEBUG)()

	migrations[fmt.Sprintf("%v", obj)] = obj
}

func UnregisterAllMigrations() {
	defer logs.GetLogger().Trace(logs.DEBUG)()

	migrations = make(map[string]interface{})
}

func doMigrations(connection Connection) error {
	defer logs.GetLogger().Trace(logs.DEBUG)()

	logger := logs.GetLogger()
	for k, v := range migrations {
		logger.Info("migrating", logs.Args{{"migration", k}})
		if err := connection.GetDb().AutoMigrate(v).Error; err != nil {
			logger.ErrorRet(err, "failed")
		}
		delete(migrations, k)
	}
	return nil
}
<commit_msg>UB-1757 Fix for error: concurrent map writes #285<commit_after>/**
 * Copyright 2017 IBM Corp.
 *
 * 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 database

import (
	"fmt"
	"sync"
	"github.com/IBM/ubiquity/utils/logs"
)

var migrations = new(sync.Map)

func RegisterMigration(obj interface{}) {
	defer logs.GetLogger().Trace(logs.DEBUG)()
	migrations.Store(fmt.Sprintf("%v", obj), obj)
}

func UnregisterAllMigrations() {
	defer logs.GetLogger().Trace(logs.DEBUG)()
	migrations = new(sync.Map)
}

func doMigrations(connection Connection) error {
	defer logs.GetLogger().Trace(logs.DEBUG)()

	logger := logs.GetLogger()

	migrations.Range(func(k, v interface{}) bool {
		logger.Info("migrating", logs.Args{{"migration", k}})
		if err := connection.GetDb().AutoMigrate(v).Error; err != nil {
			logger.ErrorRet(err, "failed")
		}
		migrations.Delete(k)
		return true
	})

	return nil
}
<|endoftext|>


<commit_before>/****************************************************************************
 * This file is part of Builder.
 *
 * Copyright (C) 2015 Pier Luigi Fiorini
 *
 * Author(s):
 *    Pier Luigi Fiorini <pierluigi.fiorini@gmail.com>
 *
 * $BEGIN_LICENSE:AGPL3+$
 *
 * 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/>.
 *
 * $END_LICENSE$
 ***************************************************************************/

package database

import (
	"bytes"
	"encoding/json"
	"github.com/boltdb/bolt"
)

type Package struct {
	Name          string   `json:"name"`
	Architectures []string `json:"architectures"`
	Ci            bool     `json:"ci"`
	Vcs           VcsInfo  `json:"vcs"`
	UpstreamVcs   VcsInfo  `json:"upstream_vcs"`
}

// Return whether the package was stored into the db.
func (db *Database) HasPackage(name string) bool {
	var found bool = false
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("package"))
		if bucket == nil {
			return nil
		}

		c := bucket.Cursor()
		for k, _ := c.Seek([]byte(name)); bytes.Equal(k, []byte(name)); k, _ = c.Next() {
			found = true
			return nil
		}
		return nil
	})
	return found
}

// Return a list of package names.
func (db *Database) GetPackageNames() []string {
	var list = []string{}
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("package"))
		if bucket == nil {
			return nil
		}

		bucket.ForEach(func(k, v []byte) error {
			list = append(list, string(k))
			return nil
		})
		return nil
	})
	return list
}

// Return a list of all packages.
func (db *Database) ListAllPackages() []*Package {
	var list []*Package
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("package"))
		if bucket == nil {
			return nil
		}

		bucket.ForEach(func(k, v []byte) error {
			pkg := &Package{}
			json.Unmarshal(v, &pkg)
			list = append(list, pkg)
			return nil
		})
		return nil
	})
	return list
}

// Return a package from the database.
func (db *Database) GetPackage(name string) *Package {
	var pkg *Package = nil
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("package"))
		if bucket == nil {
			return nil
		}

		c := bucket.Cursor()
		for k, v := c.Seek([]byte(name)); bytes.Equal(k, []byte(name)); k, v = c.Next() {
			json.Unmarshal(v, &pkg)
			return nil
		}
		return nil
	})
	return pkg
}

// Add a package to the database.
func (db *Database) AddPackage(pkg *Package) error {
	encoded, err := json.Marshal(pkg)
	if err != nil {
		return nil
	}
	return db.db.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists([]byte("package"))
		if err != nil {
			return err
		}

		err = bucket.Put([]byte(pkg.Name), encoded)
		if err != nil {
			return err
		}

		return nil
	})
}

// Remove a package from the database.
func (db *Database) RemovePackage(name string) error {
	return db.db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("package"))
		if bucket == nil {
			return nil
		}

		err := bucket.Delete([]byte(name))
		if err != nil {
			return err
		}

		return nil
	})
}
<commit_msg>Return an error if the package cannot be serialized<commit_after>/****************************************************************************
 * This file is part of Builder.
 *
 * Copyright (C) 2015 Pier Luigi Fiorini
 *
 * Author(s):
 *    Pier Luigi Fiorini <pierluigi.fiorini@gmail.com>
 *
 * $BEGIN_LICENSE:AGPL3+$
 *
 * 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/>.
 *
 * $END_LICENSE$
 ***************************************************************************/

package database

import (
	"bytes"
	"encoding/json"
	"github.com/boltdb/bolt"
)

type Package struct {
	Name          string   `json:"name"`
	Architectures []string `json:"architectures"`
	Ci            bool     `json:"ci"`
	Vcs           VcsInfo  `json:"vcs"`
	UpstreamVcs   VcsInfo  `json:"upstream_vcs"`
}

// Return whether the package was stored into the db.
func (db *Database) HasPackage(name string) bool {
	var found bool = false
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("package"))
		if bucket == nil {
			return nil
		}

		c := bucket.Cursor()
		for k, _ := c.Seek([]byte(name)); bytes.Equal(k, []byte(name)); k, _ = c.Next() {
			found = true
			return nil
		}
		return nil
	})
	return found
}

// Return a list of package names.
func (db *Database) GetPackageNames() []string {
	var list = []string{}
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("package"))
		if bucket == nil {
			return nil
		}

		bucket.ForEach(func(k, v []byte) error {
			list = append(list, string(k))
			return nil
		})
		return nil
	})
	return list
}

// Return a list of all packages.
func (db *Database) ListAllPackages() []*Package {
	var list []*Package
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("package"))
		if bucket == nil {
			return nil
		}

		bucket.ForEach(func(k, v []byte) error {
			pkg := &Package{}
			json.Unmarshal(v, &pkg)
			list = append(list, pkg)
			return nil
		})
		return nil
	})
	return list
}

// Return a package from the database.
func (db *Database) GetPackage(name string) *Package {
	var pkg *Package = nil
	db.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("package"))
		if bucket == nil {
			return nil
		}

		c := bucket.Cursor()
		for k, v := c.Seek([]byte(name)); bytes.Equal(k, []byte(name)); k, v = c.Next() {
			json.Unmarshal(v, &pkg)
			return nil
		}
		return nil
	})
	return pkg
}

// Add a package to the database.
func (db *Database) AddPackage(pkg *Package) error {
	encoded, err := json.Marshal(pkg)
	if err != nil {
		return err
	}
	return db.db.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists([]byte("package"))
		if err != nil {
			return err
		}

		err = bucket.Put([]byte(pkg.Name), encoded)
		if err != nil {
			return err
		}

		return nil
	})
}

// Remove a package from the database.
func (db *Database) RemovePackage(name string) error {
	return db.db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte("package"))
		if bucket == nil {
			return nil
		}

		err := bucket.Delete([]byte(name))
		if err != nil {
			return err
		}

		return nil
	})
}
<|endoftext|>


<commit_before>package sarama

type MessageBlock struct {
	Offset int64
	Msg    *Message
}

// Messages convenience helper which returns either all the
// messages that are wrapped in this block
func (msb *MessageBlock) Messages() []*MessageBlock {
	if msb.Msg.Set != nil {
		return msb.Msg.Set.Messages
	}
	return []*MessageBlock{msb}
}

func (msb *MessageBlock) encode(pe packetEncoder) error {
	pe.putInt64(msb.Offset)
	pe.push(&lengthField{})
	err := msb.Msg.encode(pe)
	if err != nil {
		return err
	}
	return pe.pop()
}

func (msb *MessageBlock) decode(pd packetDecoder) (err error) {
	if msb.Offset, err = pd.getInt64(); err != nil {
		return err
	}

	if err = pd.push(&lengthField{}); err != nil {
		return err
	}

	msb.Msg = new(Message)
	if err = msb.Msg.decode(pd); err != nil {
		return err
	}

	if err = pd.pop(); err != nil {
		return err
	}

	return nil
}

type MessageSet struct {
	PartialTrailingMessage bool // whether the set on the wire contained an incomplete trailing MessageBlock
	Messages               []*MessageBlock
}

func (ms *MessageSet) encode(pe packetEncoder) error {
	for i := range ms.Messages {
		err := ms.Messages[i].encode(pe)
		if err != nil {
			return err
		}
	}
	return nil
}

func (ms *MessageSet) decode(pd packetDecoder) (err error) {
	ms.Messages = nil

	for pd.remaining() > 0 {
		magic, err := magicValue(pd)
		if err != nil {
			return err
		}

		if magic > 1 {
			return nil
		}

		msb := new(MessageBlock)
		err = msb.decode(pd)
		switch err {
		case nil:
			ms.Messages = append(ms.Messages, msb)
		case ErrInsufficientData:
			// As an optimization the server is allowed to return a partial message at the
			// end of the message set. Clients should handle this case. So we just ignore such things.
			ms.PartialTrailingMessage = true
			return nil
		default:
			return err
		}
	}

	return nil
}

func (ms *MessageSet) addMessage(msg *Message) {
	block := new(MessageBlock)
	block.Msg = msg
	ms.Messages = append(ms.Messages, block)
}
<commit_msg>Do not choke on MessageSet without version<commit_after>package sarama

type MessageBlock struct {
	Offset int64
	Msg    *Message
}

// Messages convenience helper which returns either all the
// messages that are wrapped in this block
func (msb *MessageBlock) Messages() []*MessageBlock {
	if msb.Msg.Set != nil {
		return msb.Msg.Set.Messages
	}
	return []*MessageBlock{msb}
}

func (msb *MessageBlock) encode(pe packetEncoder) error {
	pe.putInt64(msb.Offset)
	pe.push(&lengthField{})
	err := msb.Msg.encode(pe)
	if err != nil {
		return err
	}
	return pe.pop()
}

func (msb *MessageBlock) decode(pd packetDecoder) (err error) {
	if msb.Offset, err = pd.getInt64(); err != nil {
		return err
	}

	if err = pd.push(&lengthField{}); err != nil {
		return err
	}

	msb.Msg = new(Message)
	if err = msb.Msg.decode(pd); err != nil {
		return err
	}

	if err = pd.pop(); err != nil {
		return err
	}

	return nil
}

type MessageSet struct {
	PartialTrailingMessage bool // whether the set on the wire contained an incomplete trailing MessageBlock
	Messages               []*MessageBlock
}

func (ms *MessageSet) encode(pe packetEncoder) error {
	for i := range ms.Messages {
		err := ms.Messages[i].encode(pe)
		if err != nil {
			return err
		}
	}
	return nil
}

func (ms *MessageSet) decode(pd packetDecoder) (err error) {
	ms.Messages = nil

	for pd.remaining() > 0 {
		magic, err := magicValue(pd)
		if err != nil {
			if err == ErrInsufficientData {
				ms.PartialTrailingMessage = true
				return nil
			}
			return err
		}

		if magic > 1 {
			return nil
		}

		msb := new(MessageBlock)
		err = msb.decode(pd)
		switch err {
		case nil:
			ms.Messages = append(ms.Messages, msb)
		case ErrInsufficientData:
			// As an optimization the server is allowed to return a partial message at the
			// end of the message set. Clients should handle this case. So we just ignore such things.
			ms.PartialTrailingMessage = true
			return nil
		default:
			return err
		}
	}

	return nil
}

func (ms *MessageSet) addMessage(msg *Message) {
	block := new(MessageBlock)
	block.Msg = msg
	ms.Messages = append(ms.Messages, block)
}
<|endoftext|>


<commit_before>package db_test

import (
	"github.com/Lunchr/luncher-api/db"
	"github.com/Lunchr/luncher-api/db/model"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"testing"
)

var (
	dbClient                  *db.Client
	offersCollection          db.Offers
	offerGroupPostsCollection db.OfferGroupPosts
	tagsCollection            db.Tags
	regionsCollection         db.Regions
	restaurantsCollection     db.Restaurants
	usersCollection           db.Users
	mocks                     *Mocks
)

func TestDb(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Db Suite")
}

var _ = BeforeSuite(func(done Done) {
	defer close(done)
	mocks = createMocks()
	createClient()
	wipeDb()
	initCollections()
})

// RebuildDBAfterEach can be used on tests or test blocks that mess up the data
// in the DB. Most commonly inserts and updates
var RebuildDBAfterEach = func() {
	AfterEach(func(done Done) {
		defer close(done)
		wipeDb()
		initCollections()
	})
}
var _ = AfterSuite(func(done Done) {
	defer close(done)
	wipeDb()
	dbClient.Disconnect()
})

var _ = It("should work", func() {})

func createClient() {
	dbConfig := createTestDbConf()
	dbClient = db.NewClient(dbConfig)
	err := dbClient.Connect()
	Expect(err).NotTo(HaveOccurred())
}

func initCollections() {
	initOffersCollection()
	initOfferGroupPostsCollection()
	initTagsCollection()
	initRegionsCollection()
	initRestaurantsCollection()
	initUsersCollection()
}

func initOffersCollection() {
	var err error
	offersCollection, err = db.NewOffers(dbClient)
	Expect(err).NotTo(HaveOccurred())
	_, err = insertOffers()
	Expect(err).NotTo(HaveOccurred())
}

func initOfferGroupPostsCollection() {
	offerGroupPostsCollection = db.NewOfferGroupPosts(dbClient)
}

func initTagsCollection() {
	tagsCollection = db.NewTags(dbClient)
	err := insertTags()
	Expect(err).NotTo(HaveOccurred())
}

func initRegionsCollection() {
	regionsCollection = db.NewRegions(dbClient)
	err := insertRegions()
	Expect(err).NotTo(HaveOccurred())
}

func initRestaurantsCollection() {
	restaurantsCollection = db.NewRestaurants(dbClient)
	_, err := insertRestaurants()
	Expect(err).NotTo(HaveOccurred())
}

func initUsersCollection() {
	usersCollection = db.NewUsers(dbClient)
	err := insertUsers()
	Expect(err).NotTo(HaveOccurred())
}

func createTestDbConf() (dbConfig *db.Config) {
	dbConfig = &db.Config{
		DbURL:  "mongodb://localhost/test",
		DbName: "test",
	}
	return
}

func insertTags() (err error) {
	return tagsCollection.Insert(mocks.tags...)
}

func insertRegions() (err error) {
	return regionsCollection.Insert(mocks.regions...)
}

func insertRestaurants() ([]*model.Restaurant, error) {
	return restaurantsCollection.Insert(mocks.restaurants...)
}

func insertOffers() ([]*model.Offer, error) {
	return offersCollection.Insert(mocks.offers...)
}

func insertUsers() (err error) {
	return usersCollection.Insert(mocks.users...)
}

func wipeDb() {
	err := dbClient.WipeDb()
	Expect(err).NotTo(HaveOccurred())
}
<commit_msg>Make db tests' setup synchronous<commit_after>package db_test

import (
	"github.com/Lunchr/luncher-api/db"
	"github.com/Lunchr/luncher-api/db/model"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"testing"
)

var (
	dbClient                  *db.Client
	offersCollection          db.Offers
	offerGroupPostsCollection db.OfferGroupPosts
	tagsCollection            db.Tags
	regionsCollection         db.Regions
	restaurantsCollection     db.Restaurants
	usersCollection           db.Users
	mocks                     *Mocks
)

func TestDb(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Db Suite")
}

var _ = BeforeSuite(func() {
	mocks = createMocks()
	createClient()
	wipeDb()
	initCollections()
})

// RebuildDBAfterEach can be used on tests or test blocks that mess up the data
// in the DB. Most commonly inserts and updates
var RebuildDBAfterEach = func() {
	AfterEach(func() {
		wipeDb()
		initCollections()
	})
}
var _ = AfterSuite(func() {
	wipeDb()
	dbClient.Disconnect()
})

var _ = It("should work", func() {})

func createClient() {
	dbConfig := createTestDbConf()
	dbClient = db.NewClient(dbConfig)
	err := dbClient.Connect()
	Expect(err).NotTo(HaveOccurred())
}

func initCollections() {
	initOffersCollection()
	initOfferGroupPostsCollection()
	initTagsCollection()
	initRegionsCollection()
	initRestaurantsCollection()
	initUsersCollection()
}

func initOffersCollection() {
	var err error
	offersCollection, err = db.NewOffers(dbClient)
	Expect(err).NotTo(HaveOccurred())
	_, err = insertOffers()
	Expect(err).NotTo(HaveOccurred())
}

func initOfferGroupPostsCollection() {
	offerGroupPostsCollection = db.NewOfferGroupPosts(dbClient)
}

func initTagsCollection() {
	tagsCollection = db.NewTags(dbClient)
	err := insertTags()
	Expect(err).NotTo(HaveOccurred())
}

func initRegionsCollection() {
	regionsCollection = db.NewRegions(dbClient)
	err := insertRegions()
	Expect(err).NotTo(HaveOccurred())
}

func initRestaurantsCollection() {
	restaurantsCollection = db.NewRestaurants(dbClient)
	_, err := insertRestaurants()
	Expect(err).NotTo(HaveOccurred())
}

func initUsersCollection() {
	usersCollection = db.NewUsers(dbClient)
	err := insertUsers()
	Expect(err).NotTo(HaveOccurred())
}

func createTestDbConf() (dbConfig *db.Config) {
	dbConfig = &db.Config{
		DbURL:  "mongodb://localhost/test",
		DbName: "test",
	}
	return
}

func insertTags() (err error) {
	return tagsCollection.Insert(mocks.tags...)
}

func insertRegions() (err error) {
	return regionsCollection.Insert(mocks.regions...)
}

func insertRestaurants() ([]*model.Restaurant, error) {
	return restaurantsCollection.Insert(mocks.restaurants...)
}

func insertOffers() ([]*model.Offer, error) {
	return offersCollection.Insert(mocks.offers...)
}

func insertUsers() (err error) {
	return usersCollection.Insert(mocks.users...)
}

func wipeDb() {
	err := dbClient.WipeDb()
	Expect(err).NotTo(HaveOccurred())
}
<|endoftext|>


<commit_before>package resources

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/endpoints"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
)

func init() {
	register("S3Bucket", ListS3Buckets)
}

func ListS3Buckets(s *session.Session) ([]Resource, error) {
	svc := s3.New(s)

	buckets, err := DescribeS3Buckets(svc)
	if err != nil {
		return nil, err
	}

	resources := make([]Resource, 0)
	for _, name := range buckets {
		resources = append(resources, &S3Bucket{
			svc:  svc,
			name: name,
		})
	}

	return resources, nil
}

func DescribeS3Buckets(svc *s3.S3) ([]string, error) {
	resp, err := svc.ListBuckets(nil)
	if err != nil {
		return nil, err
	}

	buckets := make([]string, 0)
	for _, out := range resp.Buckets {
		bucketLocationResponse, err := svc.GetBucketLocation(&s3.GetBucketLocationInput{Bucket: out.Name})

		if err != nil {
			continue
		}

		location := UnPtrString(bucketLocationResponse.LocationConstraint, endpoints.UsEast1RegionID)
		region := UnPtrString(svc.Config.Region, endpoints.UsEast1RegionID)
		if location == region {
			buckets = append(buckets, *out.Name)
		}

	}

	return buckets, nil
}

type S3Bucket struct {
	svc  *s3.S3
	name string
}

func (e *S3Bucket) Remove() error {
	err := e.RemoveAllObjects()
	if err != nil {
		return err
	}

	params := &s3.DeleteBucketInput{
		Bucket: &e.name,
	}

	_, err = e.svc.DeleteBucket(params)
	if err != nil {
		return err
	}

	return nil
}

func (e *S3Bucket) RemoveAllObjects() error {
	params := &s3.ListObjectsInput{
		Bucket: &e.name,
	}

	iterator := s3manager.NewDeleteListIterator(e.svc, params)

	err := s3manager.NewBatchDeleteWithClient(e.svc).Delete(aws.BackgroundContext(), iterator)
	if err != nil {
		return err
	}

	return nil
}

func (e *S3Bucket) String() string {
	return fmt.Sprintf("s3://%s", e.name)
}
<commit_msg>remove policy of the bucket before removing all objects and the bucket itself<commit_after>package resources

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/endpoints"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
)

func init() {
	register("S3Bucket", ListS3Buckets)
}

func ListS3Buckets(s *session.Session) ([]Resource, error) {
	svc := s3.New(s)

	buckets, err := DescribeS3Buckets(svc)
	if err != nil {
		return nil, err
	}

	resources := make([]Resource, 0)
	for _, name := range buckets {
		resources = append(resources, &S3Bucket{
			svc:  svc,
			name: name,
		})
	}

	return resources, nil
}

func DescribeS3Buckets(svc *s3.S3) ([]string, error) {
	resp, err := svc.ListBuckets(nil)
	if err != nil {
		return nil, err
	}

	buckets := make([]string, 0)
	for _, out := range resp.Buckets {
		bucketLocationResponse, err := svc.GetBucketLocation(&s3.GetBucketLocationInput{Bucket: out.Name})

		if err != nil {
			continue
		}

		location := UnPtrString(bucketLocationResponse.LocationConstraint, endpoints.UsEast1RegionID)
		region := UnPtrString(svc.Config.Region, endpoints.UsEast1RegionID)
		if location == region {
			buckets = append(buckets, *out.Name)
		}

	}

	return buckets, nil
}

type S3Bucket struct {
	svc  *s3.S3
	name string
}

func (e *S3Bucket) Remove() error {
	_, err := e.svc.DeleteBucketPolicy(&s3.DeleteBucketPolicyInput{
		Bucket: &e.name,
	})

	err = e.RemoveAllObjects()
	if err != nil {
		return err
	}

	if err != nil {
		return err
	}

	params := &s3.DeleteBucketInput{
		Bucket: &e.name,
	}

	_, err = e.svc.DeleteBucket(params)
	if err != nil {
		return err
	}

	return nil
}

func (e *S3Bucket) RemoveAllObjects() error {
	params := &s3.ListObjectsInput{
		Bucket: &e.name,
	}

	iterator := s3manager.NewDeleteListIterator(e.svc, params)

	err := s3manager.NewBatchDeleteWithClient(e.svc).Delete(aws.BackgroundContext(), iterator)
	if err != nil {
		return err
	}

	return nil
}

func (e *S3Bucket) String() string {
	return fmt.Sprintf("s3://%s", e.name)
}
<|endoftext|>


<commit_before>package admin

import (
	"errors"
	"fmt"
	"io"
	"reflect"

	"github.com/spf13/cobra"

	kapi "k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/meta"
	"k8s.io/kubernetes/pkg/apimachinery/registered"
	"k8s.io/kubernetes/pkg/kubectl"
	kcmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/kubectl/resource"
	"k8s.io/kubernetes/pkg/runtime"

	authorizationapi "github.com/openshift/origin/pkg/authorization/api"
	policyregistry "github.com/openshift/origin/pkg/authorization/registry/policy"
	policyetcd "github.com/openshift/origin/pkg/authorization/registry/policy/etcd"
	policybindingregistry "github.com/openshift/origin/pkg/authorization/registry/policybinding"
	policybindingetcd "github.com/openshift/origin/pkg/authorization/registry/policybinding/etcd"
	rolestorage "github.com/openshift/origin/pkg/authorization/registry/role/policybased"
	rolebindingstorage "github.com/openshift/origin/pkg/authorization/registry/rolebinding/policybased"

	clusterpolicyregistry "github.com/openshift/origin/pkg/authorization/registry/clusterpolicy"
	clusterpolicyetcd "github.com/openshift/origin/pkg/authorization/registry/clusterpolicy/etcd"
	clusterpolicybindingregistry "github.com/openshift/origin/pkg/authorization/registry/clusterpolicybinding"
	clusterpolicybindingetcd "github.com/openshift/origin/pkg/authorization/registry/clusterpolicybinding/etcd"
	clusterrolestorage "github.com/openshift/origin/pkg/authorization/registry/clusterrole/proxy"
	clusterrolebindingstorage "github.com/openshift/origin/pkg/authorization/registry/clusterrolebinding/proxy"
	"github.com/openshift/origin/pkg/authorization/rulevalidation"

	"github.com/openshift/origin/pkg/client"
	"github.com/openshift/origin/pkg/cmd/cli/describe"
	configapilatest "github.com/openshift/origin/pkg/cmd/server/api/latest"
	cmdclientcmd "github.com/openshift/origin/pkg/cmd/util/clientcmd"
	templateapi "github.com/openshift/origin/pkg/template/api"
	"github.com/openshift/origin/pkg/util/restoptions"
)

const OverwriteBootstrapPolicyCommandName = "overwrite-policy"

type OverwriteBootstrapPolicyOptions struct {
	File             string
	MasterConfigFile string

	Force                        bool
	Out                          io.Writer
	CreateBootstrapPolicyCommand string
}

func NewCommandOverwriteBootstrapPolicy(commandName string, fullName string, createBootstrapPolicyCommand string, out io.Writer) *cobra.Command {
	options := &OverwriteBootstrapPolicyOptions{Out: out}
	options.CreateBootstrapPolicyCommand = createBootstrapPolicyCommand

	cmd := &cobra.Command{
		Use:   commandName,
		Short: "Reset the policy to the default values",
		Run: func(cmd *cobra.Command, args []string) {
			if err := options.Validate(args); err != nil {
				kcmdutil.CheckErr(kcmdutil.UsageError(cmd, err.Error()))
			}

			if err := options.OverwriteBootstrapPolicy(); err != nil {
				kcmdutil.CheckErr(err)
			}
		},
	}

	flags := cmd.Flags()

	flags.BoolVarP(&options.Force, "force", "f", false, "You must confirm you really want to reset your policy. This will delete any custom settings you may have.")
	flags.StringVar(&options.File, "filename", "", "The policy template file containing roles and bindings.  One can be created with '"+createBootstrapPolicyCommand+"'.")
	flags.StringVar(&options.MasterConfigFile, "master-config", "openshift.local.config/master/master-config.yaml", "Location of the master configuration file to run from in order to connect to etcd and directly modify the policy.")

	// autocompletion hints
	cmd.MarkFlagFilename("filename")
	cmd.MarkFlagFilename("master-config", "yaml", "yml")

	return cmd
}

func (o OverwriteBootstrapPolicyOptions) Validate(args []string) error {
	if len(args) != 0 {
		return errors.New("no arguments are supported")
	}
	if len(o.File) == 0 {
		return errors.New("filename must be provided")
	}
	if len(o.MasterConfigFile) == 0 {
		return errors.New("master-config must be provided")
	}

	return nil
}

func (o OverwriteBootstrapPolicyOptions) OverwriteBootstrapPolicy() error {
	masterConfig, err := configapilatest.ReadAndResolveMasterConfig(o.MasterConfigFile)
	if err != nil {
		return err
	}

	optsGetter := restoptions.NewConfigGetter(*masterConfig)

	return OverwriteBootstrapPolicy(optsGetter, o.File, o.CreateBootstrapPolicyCommand, o.Force, o.Out)
}

func OverwriteBootstrapPolicy(optsGetter restoptions.Getter, policyFile, createBootstrapPolicyCommand string, change bool, out io.Writer) error {
	if !change {
		fmt.Fprintf(out, "Performing a dry run of policy overwrite:\n\n")
	}

	mapper := cmdclientcmd.ShortcutExpander{RESTMapper: kubectl.ShortcutExpander{RESTMapper: registered.RESTMapper()}}
	typer := kapi.Scheme
	clientMapper := resource.ClientMapperFunc(func(mapping *meta.RESTMapping) (resource.RESTClient, error) {
		return nil, nil
	})

	r := resource.NewBuilder(mapper, typer, clientMapper, kapi.Codecs.UniversalDecoder()).
		FilenameParam(false, false, policyFile).
		Flatten().
		Do()

	if r.Err() != nil {
		return r.Err()
	}

	policyStorage, err := policyetcd.NewStorage(optsGetter)
	if err != nil {
		return err
	}
	policyRegistry := policyregistry.NewRegistry(policyStorage)

	policyBindingStorage, err := policybindingetcd.NewStorage(optsGetter)
	if err != nil {
		return err
	}
	policyBindingRegistry := policybindingregistry.NewRegistry(policyBindingStorage)

	clusterPolicyStorage, err := clusterpolicyetcd.NewStorage(optsGetter)
	if err != nil {
		return err
	}
	clusterPolicyRegistry := clusterpolicyregistry.NewRegistry(clusterPolicyStorage)

	clusterPolicyBindingStorage, err := clusterpolicybindingetcd.NewStorage(optsGetter)
	if err != nil {
		return err
	}
	clusterPolicyBindingRegistry := clusterpolicybindingregistry.NewRegistry(clusterPolicyBindingStorage)

	ruleResolver := rulevalidation.NewDefaultRuleResolver(
		policyListerNamespacer{registry: policyRegistry},
		policyBindingListerNamespacer{registry: policyBindingRegistry},
		clusterpolicyregistry.ReadOnlyClusterPolicy{Registry: clusterPolicyRegistry},
		clusterpolicybindingregistry.ReadOnlyClusterPolicyBinding{Registry: clusterPolicyBindingRegistry},
	)

	roleStorage := rolestorage.NewVirtualStorage(policyRegistry, ruleResolver)
	roleBindingStorage := rolebindingstorage.NewVirtualStorage(policyBindingRegistry, ruleResolver)
	clusterRoleStorage := clusterrolestorage.NewClusterRoleStorage(clusterPolicyRegistry, clusterPolicyBindingRegistry)
	clusterRoleBindingStorage := clusterrolebindingstorage.NewClusterRoleBindingStorage(clusterPolicyRegistry, clusterPolicyBindingRegistry)

	return r.Visit(func(info *resource.Info, err error) error {
		if err != nil {
			return err
		}
		template, ok := info.Object.(*templateapi.Template)
		if !ok {
			return errors.New("policy must be contained in a template.  One can be created with '" + createBootstrapPolicyCommand + "'.")
		}
		runtime.DecodeList(template.Objects, kapi.Codecs.UniversalDecoder())

		for _, item := range template.Objects {
			switch t := item.(type) {
			case *authorizationapi.Role:
				ctx := kapi.WithNamespace(kapi.NewContext(), t.Namespace)
				if change {
					roleStorage.Delete(ctx, t.Name, nil)
					if _, err := roleStorage.CreateRoleWithEscalation(ctx, t); err != nil {
						return err
					}
				} else {
					fmt.Fprintf(out, "Overwrite role %s/%s\n", t.Namespace, t.Name)
					if s, err := describe.DescribeRole(t); err == nil {
						fmt.Fprintf(out, "%s\n", s)
					}
				}
			case *authorizationapi.RoleBinding:
				ctx := kapi.WithNamespace(kapi.NewContext(), t.Namespace)
				if change {
					roleBindingStorage.Delete(ctx, t.Name, nil)
					if _, err := roleBindingStorage.CreateRoleBindingWithEscalation(ctx, t); err != nil {
						return err
					}
				} else {
					fmt.Fprintf(out, "Overwrite role binding %s/%s\n", t.Namespace, t.Name)
					if s, err := describe.DescribeRoleBinding(t, nil, nil); err == nil {
						fmt.Fprintf(out, "%s\n", s)
					}
				}

			case *authorizationapi.ClusterRole:
				ctx := kapi.WithNamespace(kapi.NewContext(), t.Namespace)
				if change {
					clusterRoleStorage.Delete(ctx, t.Name, nil)
					if _, err := clusterRoleStorage.CreateClusterRoleWithEscalation(ctx, t); err != nil {
						return err
					}
				} else {
					fmt.Fprintf(out, "Overwrite role %s/%s\n", t.Namespace, t.Name)
					if s, err := describe.DescribeRole(authorizationapi.ToRole(t)); err == nil {
						fmt.Fprintf(out, "%s\n", s)
					}
				}
			case *authorizationapi.ClusterRoleBinding:
				ctx := kapi.WithNamespace(kapi.NewContext(), t.Namespace)
				if change {
					clusterRoleBindingStorage.Delete(ctx, t.Name, nil)
					if _, err := clusterRoleBindingStorage.CreateClusterRoleBindingWithEscalation(ctx, t); err != nil {
						return err
					}
				} else {
					fmt.Fprintf(out, "Overwrite role binding %s/%s\n", t.Namespace, t.Name)
					if s, err := describe.DescribeRoleBinding(authorizationapi.ToRoleBinding(t), nil, nil); err == nil {
						fmt.Fprintf(out, "%s\n", s)
					}
				}

			default:
				return fmt.Errorf("only roles and rolebindings may be created in this mode, not: %v", reflect.TypeOf(t))
			}
		}
		if !change {
			fmt.Fprintf(out, "To make the changes described above, pass --force\n")
		}
		return nil
	})
}

type policyListerNamespacer struct {
	registry policyregistry.Registry
}

func (s policyListerNamespacer) Policies(namespace string) client.PolicyLister {
	return policyLister{registry: s.registry, namespace: namespace}
}

type policyLister struct {
	registry  policyregistry.Registry
	namespace string
}

func (s policyLister) List(options kapi.ListOptions) (*authorizationapi.PolicyList, error) {
	return s.registry.ListPolicies(kapi.WithNamespace(kapi.NewContext(), s.namespace), &options)
}

func (s policyLister) Get(name string) (*authorizationapi.Policy, error) {
	return s.registry.GetPolicy(kapi.WithNamespace(kapi.NewContext(), s.namespace), name)
}

type policyBindingListerNamespacer struct {
	registry policybindingregistry.Registry
}

func (s policyBindingListerNamespacer) PolicyBindings(namespace string) client.PolicyBindingLister {
	return policyBindingLister{registry: s.registry, namespace: namespace}
}

type policyBindingLister struct {
	registry  policybindingregistry.Registry
	namespace string
}

func (s policyBindingLister) List(options kapi.ListOptions) (*authorizationapi.PolicyBindingList, error) {
	return s.registry.ListPolicyBindings(kapi.WithNamespace(kapi.NewContext(), s.namespace), &options)
}

func (s policyBindingLister) Get(name string) (*authorizationapi.PolicyBinding, error) {
	return s.registry.GetPolicyBinding(kapi.WithNamespace(kapi.NewContext(), s.namespace), name)
}
<commit_msg>Improve bootstrap policy overwrite error handling<commit_after>package admin

import (
	"errors"
	"fmt"
	"io"
	"reflect"

	"github.com/spf13/cobra"

	kapi "k8s.io/kubernetes/pkg/api"
	kapierrors "k8s.io/kubernetes/pkg/api/errors"
	"k8s.io/kubernetes/pkg/api/meta"
	"k8s.io/kubernetes/pkg/apimachinery/registered"
	"k8s.io/kubernetes/pkg/kubectl"
	kcmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/kubectl/resource"
	"k8s.io/kubernetes/pkg/runtime"
	kerrors "k8s.io/kubernetes/pkg/util/errors"

	authorizationapi "github.com/openshift/origin/pkg/authorization/api"
	policyregistry "github.com/openshift/origin/pkg/authorization/registry/policy"
	policyetcd "github.com/openshift/origin/pkg/authorization/registry/policy/etcd"
	policybindingregistry "github.com/openshift/origin/pkg/authorization/registry/policybinding"
	policybindingetcd "github.com/openshift/origin/pkg/authorization/registry/policybinding/etcd"
	rolestorage "github.com/openshift/origin/pkg/authorization/registry/role/policybased"
	rolebindingstorage "github.com/openshift/origin/pkg/authorization/registry/rolebinding/policybased"

	clusterpolicyregistry "github.com/openshift/origin/pkg/authorization/registry/clusterpolicy"
	clusterpolicyetcd "github.com/openshift/origin/pkg/authorization/registry/clusterpolicy/etcd"
	clusterpolicybindingregistry "github.com/openshift/origin/pkg/authorization/registry/clusterpolicybinding"
	clusterpolicybindingetcd "github.com/openshift/origin/pkg/authorization/registry/clusterpolicybinding/etcd"
	clusterrolestorage "github.com/openshift/origin/pkg/authorization/registry/clusterrole/proxy"
	clusterrolebindingstorage "github.com/openshift/origin/pkg/authorization/registry/clusterrolebinding/proxy"
	"github.com/openshift/origin/pkg/authorization/rulevalidation"

	"github.com/openshift/origin/pkg/client"
	"github.com/openshift/origin/pkg/cmd/cli/describe"
	configapilatest "github.com/openshift/origin/pkg/cmd/server/api/latest"
	cmdclientcmd "github.com/openshift/origin/pkg/cmd/util/clientcmd"
	templateapi "github.com/openshift/origin/pkg/template/api"
	"github.com/openshift/origin/pkg/util/restoptions"
)

const OverwriteBootstrapPolicyCommandName = "overwrite-policy"

type OverwriteBootstrapPolicyOptions struct {
	File             string
	MasterConfigFile string

	Force                        bool
	Out                          io.Writer
	CreateBootstrapPolicyCommand string
}

func NewCommandOverwriteBootstrapPolicy(commandName string, fullName string, createBootstrapPolicyCommand string, out io.Writer) *cobra.Command {
	options := &OverwriteBootstrapPolicyOptions{Out: out}
	options.CreateBootstrapPolicyCommand = createBootstrapPolicyCommand

	cmd := &cobra.Command{
		Use:   commandName,
		Short: "Reset the policy to the default values",
		Run: func(cmd *cobra.Command, args []string) {
			if err := options.Validate(args); err != nil {
				kcmdutil.CheckErr(kcmdutil.UsageError(cmd, err.Error()))
			}

			if err := options.OverwriteBootstrapPolicy(); err != nil {
				kcmdutil.CheckErr(err)
			}
		},
	}

	flags := cmd.Flags()

	flags.BoolVarP(&options.Force, "force", "f", false, "You must confirm you really want to reset your policy. This will delete any custom settings you may have.")
	flags.StringVar(&options.File, "filename", "", "The policy template file containing roles and bindings.  One can be created with '"+createBootstrapPolicyCommand+"'.")
	flags.StringVar(&options.MasterConfigFile, "master-config", "openshift.local.config/master/master-config.yaml", "Location of the master configuration file to run from in order to connect to etcd and directly modify the policy.")

	// autocompletion hints
	cmd.MarkFlagFilename("filename")
	cmd.MarkFlagFilename("master-config", "yaml", "yml")

	return cmd
}

func (o OverwriteBootstrapPolicyOptions) Validate(args []string) error {
	if len(args) != 0 {
		return errors.New("no arguments are supported")
	}
	if len(o.File) == 0 {
		return errors.New("filename must be provided")
	}
	if len(o.MasterConfigFile) == 0 {
		return errors.New("master-config must be provided")
	}

	return nil
}

func (o OverwriteBootstrapPolicyOptions) OverwriteBootstrapPolicy() error {
	masterConfig, err := configapilatest.ReadAndResolveMasterConfig(o.MasterConfigFile)
	if err != nil {
		return err
	}

	optsGetter := restoptions.NewConfigGetter(*masterConfig)

	return OverwriteBootstrapPolicy(optsGetter, o.File, o.CreateBootstrapPolicyCommand, o.Force, o.Out)
}

func OverwriteBootstrapPolicy(optsGetter restoptions.Getter, policyFile, createBootstrapPolicyCommand string, change bool, out io.Writer) error {
	if !change {
		fmt.Fprintf(out, "Performing a dry run of policy overwrite:\n\n")
	}

	mapper := cmdclientcmd.ShortcutExpander{RESTMapper: kubectl.ShortcutExpander{RESTMapper: registered.RESTMapper()}}
	typer := kapi.Scheme
	clientMapper := resource.ClientMapperFunc(func(mapping *meta.RESTMapping) (resource.RESTClient, error) {
		return nil, nil
	})

	r := resource.NewBuilder(mapper, typer, clientMapper, kapi.Codecs.UniversalDecoder()).
		FilenameParam(false, false, policyFile).
		Flatten().
		Do()

	if r.Err() != nil {
		return r.Err()
	}

	policyStorage, err := policyetcd.NewStorage(optsGetter)
	if err != nil {
		return err
	}
	policyRegistry := policyregistry.NewRegistry(policyStorage)

	policyBindingStorage, err := policybindingetcd.NewStorage(optsGetter)
	if err != nil {
		return err
	}
	policyBindingRegistry := policybindingregistry.NewRegistry(policyBindingStorage)

	clusterPolicyStorage, err := clusterpolicyetcd.NewStorage(optsGetter)
	if err != nil {
		return err
	}
	clusterPolicyRegistry := clusterpolicyregistry.NewRegistry(clusterPolicyStorage)

	clusterPolicyBindingStorage, err := clusterpolicybindingetcd.NewStorage(optsGetter)
	if err != nil {
		return err
	}
	clusterPolicyBindingRegistry := clusterpolicybindingregistry.NewRegistry(clusterPolicyBindingStorage)

	ruleResolver := rulevalidation.NewDefaultRuleResolver(
		policyListerNamespacer{registry: policyRegistry},
		policyBindingListerNamespacer{registry: policyBindingRegistry},
		clusterpolicyregistry.ReadOnlyClusterPolicy{Registry: clusterPolicyRegistry},
		clusterpolicybindingregistry.ReadOnlyClusterPolicyBinding{Registry: clusterPolicyBindingRegistry},
	)

	roleStorage := rolestorage.NewVirtualStorage(policyRegistry, ruleResolver)
	roleBindingStorage := rolebindingstorage.NewVirtualStorage(policyBindingRegistry, ruleResolver)
	clusterRoleStorage := clusterrolestorage.NewClusterRoleStorage(clusterPolicyRegistry, clusterPolicyBindingRegistry)
	clusterRoleBindingStorage := clusterrolebindingstorage.NewClusterRoleBindingStorage(clusterPolicyRegistry, clusterPolicyBindingRegistry)

	return r.Visit(func(info *resource.Info, err error) error {
		if err != nil {
			return err
		}
		template, ok := info.Object.(*templateapi.Template)
		if !ok {
			return errors.New("policy must be contained in a template.  One can be created with '" + createBootstrapPolicyCommand + "'.")
		}
		runtime.DecodeList(template.Objects, kapi.Codecs.UniversalDecoder())

		// For each object, we attempt the following to maximize our ability to persist the desired objects, while minimizing etcd write thrashing:
		// 1. Create the object (no-ops if the object already exists)
		// 2. If the object already exists, attempt to update the object (no-ops if an identical object is already persisted)
		// 3. If we encounter any error updating, delete and recreate
		errs := []error{}
		for _, item := range template.Objects {
			switch t := item.(type) {
			case *authorizationapi.Role:
				ctx := kapi.WithNamespace(kapi.NewContext(), t.Namespace)
				if change {
					// Attempt to create
					_, err := roleStorage.CreateRoleWithEscalation(ctx, t)
					// Unconditional replace if it already exists
					if kapierrors.IsAlreadyExists(err) {
						_, _, err = roleStorage.UpdateRoleWithEscalation(ctx, t)
					}
					// Delete and recreate as a last resort
					if err != nil {
						roleStorage.Delete(ctx, t.Name, nil)
						_, err = roleStorage.CreateRoleWithEscalation(ctx, t)
					}
					// Gather any error
					if err != nil {
						errs = append(errs, err)
					}
				} else {
					fmt.Fprintf(out, "Overwrite role %s/%s\n", t.Namespace, t.Name)
					if s, err := describe.DescribeRole(t); err == nil {
						fmt.Fprintf(out, "%s\n", s)
					}
				}
			case *authorizationapi.RoleBinding:
				ctx := kapi.WithNamespace(kapi.NewContext(), t.Namespace)
				if change {
					// Attempt to create
					_, err := roleBindingStorage.CreateRoleBindingWithEscalation(ctx, t)
					// Unconditional replace if it already exists
					if kapierrors.IsAlreadyExists(err) {
						_, _, err = roleBindingStorage.UpdateRoleBindingWithEscalation(ctx, t)
					}
					// Delete and recreate as a last resort
					if err != nil {
						roleBindingStorage.Delete(ctx, t.Name, nil)
						_, err = roleBindingStorage.CreateRoleBindingWithEscalation(ctx, t)
					}
					// Gather any error
					if err != nil {
						errs = append(errs, err)
					}
				} else {
					fmt.Fprintf(out, "Overwrite role binding %s/%s\n", t.Namespace, t.Name)
					if s, err := describe.DescribeRoleBinding(t, nil, nil); err == nil {
						fmt.Fprintf(out, "%s\n", s)
					}
				}

			case *authorizationapi.ClusterRole:
				ctx := kapi.WithNamespace(kapi.NewContext(), t.Namespace)
				if change {
					// Attempt to create
					_, err := clusterRoleStorage.CreateClusterRoleWithEscalation(ctx, t)
					// Unconditional replace if it already exists
					if kapierrors.IsAlreadyExists(err) {
						_, _, err = clusterRoleStorage.UpdateClusterRoleWithEscalation(ctx, t)
					}
					// Delete and recreate as a last resort
					if err != nil {
						clusterRoleStorage.Delete(ctx, t.Name, nil)
						_, err = clusterRoleStorage.CreateClusterRoleWithEscalation(ctx, t)
					}
					// Gather any error
					if err != nil {
						errs = append(errs, err)
					}
				} else {
					fmt.Fprintf(out, "Overwrite role %s/%s\n", t.Namespace, t.Name)
					if s, err := describe.DescribeRole(authorizationapi.ToRole(t)); err == nil {
						fmt.Fprintf(out, "%s\n", s)
					}
				}
			case *authorizationapi.ClusterRoleBinding:
				ctx := kapi.WithNamespace(kapi.NewContext(), t.Namespace)
				if change {
					// Attempt to create
					_, err := clusterRoleBindingStorage.CreateClusterRoleBindingWithEscalation(ctx, t)
					// Unconditional replace if it already exists
					if kapierrors.IsAlreadyExists(err) {
						_, _, err = clusterRoleBindingStorage.UpdateClusterRoleBindingWithEscalation(ctx, t)
					}
					// Delete and recreate as a last resort
					if err != nil {
						clusterRoleBindingStorage.Delete(ctx, t.Name, nil)
						_, err = clusterRoleBindingStorage.CreateClusterRoleBindingWithEscalation(ctx, t)
					}
					// Gather any error
					if err != nil {
						errs = append(errs, err)
					}
				} else {
					fmt.Fprintf(out, "Overwrite role binding %s/%s\n", t.Namespace, t.Name)
					if s, err := describe.DescribeRoleBinding(authorizationapi.ToRoleBinding(t), nil, nil); err == nil {
						fmt.Fprintf(out, "%s\n", s)
					}
				}

			default:
				errs = append(errs, fmt.Errorf("only roles and rolebindings may be created in this mode, not: %v", reflect.TypeOf(t)))
			}
		}
		if !change {
			fmt.Fprintf(out, "To make the changes described above, pass --force\n")
		}
		return kerrors.NewAggregate(errs)
	})
}

type policyListerNamespacer struct {
	registry policyregistry.Registry
}

func (s policyListerNamespacer) Policies(namespace string) client.PolicyLister {
	return policyLister{registry: s.registry, namespace: namespace}
}

type policyLister struct {
	registry  policyregistry.Registry
	namespace string
}

func (s policyLister) List(options kapi.ListOptions) (*authorizationapi.PolicyList, error) {
	return s.registry.ListPolicies(kapi.WithNamespace(kapi.NewContext(), s.namespace), &options)
}

func (s policyLister) Get(name string) (*authorizationapi.Policy, error) {
	return s.registry.GetPolicy(kapi.WithNamespace(kapi.NewContext(), s.namespace), name)
}

type policyBindingListerNamespacer struct {
	registry policybindingregistry.Registry
}

func (s policyBindingListerNamespacer) PolicyBindings(namespace string) client.PolicyBindingLister {
	return policyBindingLister{registry: s.registry, namespace: namespace}
}

type policyBindingLister struct {
	registry  policybindingregistry.Registry
	namespace string
}

func (s policyBindingLister) List(options kapi.ListOptions) (*authorizationapi.PolicyBindingList, error) {
	return s.registry.ListPolicyBindings(kapi.WithNamespace(kapi.NewContext(), s.namespace), &options)
}

func (s policyBindingLister) Get(name string) (*authorizationapi.PolicyBinding, error) {
	return s.registry.GetPolicyBinding(kapi.WithNamespace(kapi.NewContext(), s.namespace), name)
}
<|endoftext|>


<commit_before>package middlewares

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"strconv"

	"strings"

	"github.com/auth0/go-jwt-middleware"
	"github.com/dgrijalva/jwt-go"
	"github.com/prest/config"
	"github.com/urfave/negroni"
)

// HandlerSet add content type header
func HandlerSet() negroni.Handler {
	return negroni.HandlerFunc(func(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
		format := r.URL.Query().Get("_renderer")
		recorder := httptest.NewRecorder()
		negroniResp := negroni.NewResponseWriter(recorder)
		next(negroniResp, r)
		renderFormat(w, recorder, format)
	})
}

// AccessControl is a middleware to handle permissions on tables in pREST
func AccessControl() negroni.Handler {
	return negroni.HandlerFunc(func(rw http.ResponseWriter, rq *http.Request, next http.HandlerFunc) {
		mapPath := getVars(rq.URL.Path)
		if mapPath == nil {
			next(rw, rq)
			return
		}

		permission := permissionByMethod(rq.Method)
		if permission == "" {
			next(rw, rq)
			return
		}

		if config.PrestConf.Adapter.TablePermissions(mapPath["table"], permission) {
			next(rw, rq)
			return
		}

		err := fmt.Errorf("required authorization to table %s", mapPath["table"])
		http.Error(rw, err.Error(), http.StatusUnauthorized)
	})
}

// JwtMiddleware check if actual request have JWT
func JwtMiddleware(key string) negroni.Handler {
	jwtMiddleware := jwtmiddleware.New(jwtmiddleware.Options{
		ValidationKeyGetter: func(token *jwt.Token) (interface{}, error) {
			return []byte(key), nil
		},
		SigningMethod: jwt.SigningMethodHS256,
	})
	return negroni.HandlerFunc(jwtMiddleware.HandlerWithNext)
}

// Cors middleware
func Cors(origin []string, headers []string) negroni.Handler {
	return negroni.HandlerFunc(func(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
		w.Header().Set(headerAllowOrigin, strings.Join(origin, ","))
		if r.Method == "OPTIONS" && r.Header.Get("Access-Control-Request-Method") != "" {
			w.Header().Set(headerAllowMethods, strings.Join(defaultAllowMethods, ","))
			w.Header().Set(headerAllowHeaders, strings.Join(headers, ","))
			w.Header().Set(headerAllowCredentials, strconv.FormatBool(true))
			if allowed := checkCors(r, origin); !allowed {
				w.WriteHeader(http.StatusForbidden)
				return
			}
			w.WriteHeader(http.StatusOK)
			return
		}
		next(w, r)
	})
}
<commit_msg>add allowCredentials header to every request<commit_after>package middlewares

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"strconv"

	"strings"

	"github.com/auth0/go-jwt-middleware"
	"github.com/dgrijalva/jwt-go"
	"github.com/prest/config"
	"github.com/urfave/negroni"
)

// HandlerSet add content type header
func HandlerSet() negroni.Handler {
	return negroni.HandlerFunc(func(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
		format := r.URL.Query().Get("_renderer")
		recorder := httptest.NewRecorder()
		negroniResp := negroni.NewResponseWriter(recorder)
		next(negroniResp, r)
		renderFormat(w, recorder, format)
	})
}

// AccessControl is a middleware to handle permissions on tables in pREST
func AccessControl() negroni.Handler {
	return negroni.HandlerFunc(func(rw http.ResponseWriter, rq *http.Request, next http.HandlerFunc) {
		mapPath := getVars(rq.URL.Path)
		if mapPath == nil {
			next(rw, rq)
			return
		}

		permission := permissionByMethod(rq.Method)
		if permission == "" {
			next(rw, rq)
			return
		}

		if config.PrestConf.Adapter.TablePermissions(mapPath["table"], permission) {
			next(rw, rq)
			return
		}

		err := fmt.Errorf("required authorization to table %s", mapPath["table"])
		http.Error(rw, err.Error(), http.StatusUnauthorized)
	})
}

// JwtMiddleware check if actual request have JWT
func JwtMiddleware(key string) negroni.Handler {
	jwtMiddleware := jwtmiddleware.New(jwtmiddleware.Options{
		ValidationKeyGetter: func(token *jwt.Token) (interface{}, error) {
			return []byte(key), nil
		},
		SigningMethod: jwt.SigningMethodHS256,
	})
	return negroni.HandlerFunc(jwtMiddleware.HandlerWithNext)
}

// Cors middleware
func Cors(origin []string, headers []string) negroni.Handler {
	return negroni.HandlerFunc(func(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
		w.Header().Set(headerAllowOrigin, strings.Join(origin, ","))
		w.Header().Set(headerAllowCredentials, strconv.FormatBool(true))
		if r.Method == "OPTIONS" && r.Header.Get("Access-Control-Request-Method") != "" {
			w.Header().Set(headerAllowMethods, strings.Join(defaultAllowMethods, ","))
			w.Header().Set(headerAllowHeaders, strings.Join(headers, ","))
			if allowed := checkCors(r, origin); !allowed {
				w.WriteHeader(http.StatusForbidden)
				return
			}
			w.WriteHeader(http.StatusOK)
			return
		}
		next(w, r)
	})
}
<|endoftext|>


<commit_before>package auth

import (
	"github.com/fluxxu/session"
	"github.com/fluxxu/util"
	"github.com/zenazn/goji/web"
	"net/http"
)

func sessionMiddleware(c *web.C, h http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		//fmt.Println("session in")
		var sess *session.Session = nil
		var err error
		auth := r.Header.Get("Authorization")
		methodLen := len("Session ")
		if len(auth) > methodLen {
			sid := auth[methodLen:]
			if sid != "" {
				sess, err = sessionStore.Load(sid)
				if err != nil {
					util.Response(w).Error("load session error: " + err.Error())
					return
				}
				if sess != nil {
					c.Env["session"] = sess
				}
			}
		}
		h.ServeHTTP(w, r)
		if sess != nil {
			sessionStore.Save(sess)
		}
		//fmt.Println("session out")
	}
	return http.HandlerFunc(fn)
}

var skipAuthPath = make(map[string]bool)

func Skip(path string) {
	skipAuthPath[path] = true
}

func authMiddleware(c *web.C, h http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		// allow POST /token
		if r.Method == "POST" && r.URL.Path == muxBase+"/token" {
			h.ServeHTTP(w, r)
			return
		}

		// skip
		if _, ok := skipAuthPath[r.URL.Path]; ok {
			h.ServeHTTP(w, r)
			return
		}

		// allow login user only
		if _, ok := c.Env["session"]; !ok {
			util.Response(w).Error("Unauthorized", 401)
			return
		}

		h.ServeHTTP(w, r)
	}
	return http.HandlerFunc(fn)
}
<commit_msg>Allow skip method<commit_after>package auth

import (
	"github.com/fluxxu/session"
	"github.com/fluxxu/util"
	"github.com/zenazn/goji/web"
	"net/http"
)

func sessionMiddleware(c *web.C, h http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		//fmt.Println("session in")
		var sess *session.Session = nil
		var err error
		auth := r.Header.Get("Authorization")
		methodLen := len("Session ")
		if len(auth) > methodLen {
			sid := auth[methodLen:]
			if sid != "" {
				sess, err = sessionStore.Load(sid)
				if err != nil {
					util.Response(w).Error("load session error: " + err.Error())
					return
				}
				if sess != nil {
					c.Env["session"] = sess
				}
			}
		}
		h.ServeHTTP(w, r)
		if sess != nil {
			sessionStore.Save(sess)
		}
		//fmt.Println("session out")
	}
	return http.HandlerFunc(fn)
}

var skipAuthPath = make(map[string][]string)

func Skip(path string, args ...string) {
	skipAuthPath[path] = args
}

func authMiddleware(c *web.C, h http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		// allow POST /token
		if r.Method == "POST" && r.URL.Path == muxBase+"/token" {
			h.ServeHTTP(w, r)
			return
		}

		// skip
		if methods, ok := skipAuthPath[r.URL.Path]; ok {
			if len(methods) == 0 {
				h.ServeHTTP(w, r)
				return
			} else {
				if util.IndexOfString(methods, r.Method) != -1 {
					h.ServeHTTP(w, r)
					return
				}
			}
		}

		// allow login user only
		if _, ok := c.Env["session"]; !ok {
			util.Response(w).Error("Unauthorized", 401)
			return
		}

		h.ServeHTTP(w, r)
	}
	return http.HandlerFunc(fn)
}
<|endoftext|>


<commit_before>package minion

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"math/rand"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/dnaeon/backoff"
	"github.com/dnaeon/gru/classifier"
	"github.com/dnaeon/gru/task"
	"github.com/dnaeon/gru/utils"
	"github.com/libgit2/git2go"

	etcdclient "github.com/coreos/etcd/client"
	"github.com/pborman/uuid"
	"golang.org/x/net/context"
)

// EtcdMinionSpace is the keyspace in etcd used by minions
const EtcdMinionSpace = "/gru/minion"

// Etcd Minion
type etcdMinion struct {
	// Name of the minion
	name string

	// Minion root directory in etcd
	rootDir string

	// Minion queue directory in etcd
	queueDir string

	// Log directory of previously executed tasks
	logDir string

	// Classifier directory in etcd
	classifierDir string

	// Minion unique identifier
	id uuid.UUID

	// KeysAPI client to etcd
	kapi etcdclient.KeysAPI

	// Channel over which tasks are sent for processing
	taskQueue chan *task.Task

	// Git repo from which to sync modules and data
	gitRepo string

	// Site directory contains the modules and data synced from Git
	siteDir string

	// Channel used to signal for shutdown time
	done chan struct{}
}

// EtcdMinionConfig contains configuration settings for
// minions which use the etcd backend implementation
type EtcdMinionConfig struct {
	// Name of the minion
	Name string

	// GitRepo provides the url from which to sync modules and data
	GitRepo string

	// EtcdConfig provides etcd-related configuration settings
	EtcdConfig etcdclient.Config
}

// NewEtcdMinion creates a new minion with etcd backend
func NewEtcdMinion(config *EtcdMinionConfig) (Minion, error) {
	c, err := etcdclient.New(config.EtcdConfig)
	if err != nil {
		return nil, err
	}

	cwd, err := os.Getwd()
	if err != nil {
		return nil, err
	}

	kapi := etcdclient.NewKeysAPI(c)
	id := utils.GenerateUUID(config.Name)
	rootDir := filepath.Join(EtcdMinionSpace, id.String())
	queueDir := filepath.Join(rootDir, "queue")
	classifierDir := filepath.Join(rootDir, "classifier")
	logDir := filepath.Join(rootDir, "log")
	siteDir := filepath.Join(cwd, "site")
	taskQueue := make(chan *task.Task)
	done := make(chan struct{})

	m := &etcdMinion{
		name:          config.Name,
		rootDir:       rootDir,
		queueDir:      queueDir,
		classifierDir: classifierDir,
		logDir:        logDir,
		id:            id,
		kapi:          kapi,
		taskQueue:     taskQueue,
		gitRepo:       config.GitRepo,
		siteDir:       siteDir,
		done:          done,
	}

	return m, nil
}

// Checks for any pending tasks and sends them
// for processing if there are any
func (m *etcdMinion) checkQueue() error {
	opts := &etcdclient.GetOptions{
		Recursive: true,
		Sort:      true,
	}

	// Get backlog tasks if any
	// If the directory key in etcd is missing that is okay, since
	// it means there are no pending tasks for processing
	resp, err := m.kapi.Get(context.Background(), m.queueDir, opts)
	if err != nil {
		if eerr, ok := err.(etcdclient.Error); !ok || eerr.Code == etcdclient.ErrorCodeKeyNotFound {
			return err
		}
	}

	backlog := resp.Node.Nodes
	if len(backlog) == 0 {
		// No backlog tasks found
		return nil
	}

	log.Printf("Found %d pending tasks in queue", len(backlog))
	for _, node := range backlog {
		t, err := EtcdUnmarshalTask(node)
		m.kapi.Delete(context.Background(), node.Key, nil)

		if err != nil {
			continue
		}

		m.taskQueue <- t
	}

	return nil
}

// Runs periodic jobs such as refreshing classifiers and
// updating the lastseen time every five minutes
func (m *etcdMinion) periodicRunner() {
	schedule := time.Minute * 5
	ticker := time.NewTicker(schedule)
	log.Printf("Periodic scheduler set to run every %s\n", schedule)

	for {
		select {
		case <-m.done:
			break
		case now := <-ticker.C:
			// Run periodic jobs
			m.classify()
			m.checkQueue()
			m.SetLastseen(now.Unix())
		}
	}
}

// Processes new tasks
func (m *etcdMinion) processTask(t *task.Task) error {
	// Update state of task to indicate that we are now processing it
	t.State = task.TaskStateProcessing
	if err := m.SaveTaskResult(t); err != nil {
		return err
	}

	var buf bytes.Buffer
	fmt.Fprintf(&buf, "Loaded %d resources from catalog\n", t.Catalog.Len())
	catalogErr := t.Catalog.Run(&buf)
	t.TimeProcessed = time.Now().Unix()
	t.Result = buf.String()

	if catalogErr != nil {
		t.State = task.TaskStateFailed
	} else {
		t.State = task.TaskStateSuccess
	}

	if err := m.SaveTaskResult(t); err != nil {
		return err
	}

	return catalogErr
}

// Classifies the minion
func (m *etcdMinion) classify() error {
	for key := range classifier.Registry {
		klassifier, err := classifier.Get(key)

		if err != nil {
			log.Printf("Failed to get classifier %s: %s\n", key, err)
			continue
		}

		m.SetClassifier(klassifier)
	}

	return nil
}

// EtcdUnmarshalTask unmarshals tasks from an etcd node
func EtcdUnmarshalTask(node *etcdclient.Node) (*task.Task, error) {
	task := new(task.Task)
	err := json.Unmarshal([]byte(node.Value), &task)

	return task, err
}

// ID returns the minion unique identifier
func (m *etcdMinion) ID() uuid.UUID {
	return m.id
}

// SetName sets the human-readable name of the minion in etcd
func (m *etcdMinion) SetName(name string) error {
	nameKey := filepath.Join(m.rootDir, "name")
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), nameKey, name, opts)
	if err != nil {
		log.Printf("Failed to set name of minion: %s\n", err)
	}

	return err
}

// SetLastseen sets the time the minion was last seen in
// seconds since the Epoch
func (m *etcdMinion) SetLastseen(s int64) error {
	lastseenKey := filepath.Join(m.rootDir, "lastseen")
	lastseenValue := strconv.FormatInt(s, 10)
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), lastseenKey, lastseenValue, opts)
	if err != nil {
		log.Printf("Failed to set lastseen time: %s\n", err)
	}

	return err
}

// SetClassifier sets a classifier for the minion in etcd
func (m *etcdMinion) SetClassifier(c *classifier.Classifier) error {
	// Classifiers in etcd expire after an hour
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
		TTL:       time.Hour,
	}

	// Serialize classifier to JSON and save it in etcd
	data, err := json.Marshal(c)
	if err != nil {
		log.Printf("Failed to serialize classifier %s: %s\n", c.Key, err)
		return err
	}

	// Classifier key in etcd
	klassifierKey := filepath.Join(m.classifierDir, c.Key)
	_, err = m.kapi.Set(context.Background(), klassifierKey, string(data), opts)

	if err != nil {
		log.Printf("Failed to set classifier %s: %s\n", c.Key, err)
		return err
	}

	return nil
}

// TaskListener monitors etcd for new tasks
func (m *etcdMinion) TaskListener(c chan<- *task.Task) error {
	log.Printf("Task listener is watching %s\n", m.queueDir)

	rand.Seed(time.Now().UTC().UnixNano())
	b := backoff.Backoff{
		Min:    1 * time.Second,
		Max:    10 * time.Minute,
		Factor: 2.0,
		Jitter: true,
	}

	watcherOpts := &etcdclient.WatcherOptions{
		Recursive: true,
	}
	watcher := m.kapi.Watcher(m.queueDir, watcherOpts)

	for {
		resp, err := watcher.Next(context.Background())
		if err != nil {
			// Use a backoff and retry later again
			duration := b.Duration()
			log.Printf("%s, retrying in %s\n", err, duration)
			time.Sleep(duration)
			continue
		}

		// Reset the backoff counter on successful receive
		b.Reset()

		// Ignore "delete" events when removing a task from the queue
		action := strings.ToLower(resp.Action)
		if strings.EqualFold(action, "delete") {
			continue
		}

		// Unmarshal and remove task from the queue
		t, err := EtcdUnmarshalTask(resp.Node)
		m.kapi.Delete(context.Background(), resp.Node.Key, nil)

		if err != nil {
			log.Printf("Received invalid task %s: %s\n", resp.Node.Key, err)
			continue
		}

		// Send the task for processing
		log.Printf("Received task %s\n", t.TaskID)
		t.State = task.TaskStateQueued
		t.TimeReceived = time.Now().Unix()
		if err := m.SaveTaskResult(t); err != nil {
			log.Printf("Unable to save task state: %s\n", err)
			continue
		}

		c <- t
	}

	return nil
}

// TaskRunner processes new tasks
func (m *etcdMinion) TaskRunner(c <-chan *task.Task) error {
	log.Println("Starting task runner")

	for {
		select {
		case <-m.done:
			break
		case t := <-c:
			// Sync the module and data files first, then process the task
			// TODO: In case of failures to sync site repo, we may consider
			//       continuing the task processing from the last known state
			m.Sync()

			// Switch to the correct environment
			if err := m.setEnvironment(t.Environment); err != nil {
				log.Printf("Unable to set environment: %s\n", err)
				continue
			}

			log.Printf("Processing task %s\n", t.TaskID)
			err := m.processTask(t)
			if err != nil {
				log.Printf("Failed to process task %s: %s\n", t.TaskID, err)
			} else {
				log.Printf("Finished processing task %s\n", t.TaskID)
			}
		}
	}

	return nil
}

// SaveTaskResult stores the result of a task in etcd
func (m *etcdMinion) SaveTaskResult(t *task.Task) error {
	taskKey := filepath.Join(m.logDir, t.TaskID.String())

	data, err := json.Marshal(t)
	if err != nil {
		return err
	}

	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err = m.kapi.Set(context.Background(), taskKey, string(data), opts)

	return err
}

// setEnvironment checks out the branch with the
// provided name in the site repo dir and sets HEAD to detached
func (m *etcdMinion) setEnvironment(name string) error {
	repo, err := git.OpenRepository(m.siteDir)
	if err != nil {
		return err
	}

	remoteBranchName := filepath.Join("origin", name)
	remoteBranch, err := repo.LookupBranch(remoteBranchName, git.BranchRemote)
	if err != nil {
		return err
	}

	err = repo.SetHeadDetached(remoteBranch.Target())
	if err != nil {
		return err
	}

	checkoutOpts := &git.CheckoutOpts{
		Strategy: git.CheckoutForce,
	}

	err = repo.CheckoutHead(checkoutOpts)
	if err != nil {
		return err
	}

	return nil
}

// Sync syncs module and data files from the upstream Git repository
func (m *etcdMinion) Sync() error {
	fi, err := os.Stat(m.siteDir)
	if err != nil {
		return err
	}

	// Site directory does not exist, clone the Git repository from upstream
	if os.IsNotExist(err) {
		log.Printf("site repo is missing, performing initial sync from %s\n", m.gitRepo)
		opts := &git.CloneOptions{}
		if _, err := git.Clone(m.gitRepo, m.siteDir, opts); err != nil {
			return err
		}
	}

	// File exists, ensure that it is a valid Git repository
	if !fi.IsDir() {
		return fmt.Errorf("%s exists, but is not a directory", m.siteDir)
	}

	// Open the repository and fetch from the default remote
	log.Println("starting site repo sync")
	repo, err := git.OpenRepository(m.siteDir)
	if err != nil {
		return err
	}

	origin, err := repo.Remotes.Lookup("origin")
	if err != nil {
		return err
	}

	err = origin.Fetch([]string{}, &git.FetchOptions{}, "")
	if err != nil {
		log.Printf("site repo sync failed: %s\n", err)
	} else {
		log.Println("site repo sync completed")
	}

	return err
}

// Serve starts the minion
func (m *etcdMinion) Serve() error {
	err := m.SetName(m.name)
	if err != nil {
		return err
	}

	now := time.Now().Unix()
	err = m.SetLastseen(now)
	if err != nil {
		return err
	}

	// Start minion services
	go m.classify()
	go m.checkQueue()
	go m.periodicRunner()
	go m.TaskRunner(m.taskQueue)
	go m.TaskListener(m.taskQueue)

	log.Printf("Minion %s is ready to serve", m.ID())

	return nil
}

// Stop shutdowns the minions and its services
func (m *etcdMinion) Stop() error {
	log.Println("Minion is shutting down")

	close(m.taskQueue)
	close(m.done)

	return nil
}
<commit_msg>Check if the minion has a configured Git repo<commit_after>package minion

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"math/rand"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/dnaeon/backoff"
	"github.com/dnaeon/gru/classifier"
	"github.com/dnaeon/gru/task"
	"github.com/dnaeon/gru/utils"
	"github.com/libgit2/git2go"

	etcdclient "github.com/coreos/etcd/client"
	"github.com/pborman/uuid"
	"golang.org/x/net/context"
)

// EtcdMinionSpace is the keyspace in etcd used by minions
const EtcdMinionSpace = "/gru/minion"

// Etcd Minion
type etcdMinion struct {
	// Name of the minion
	name string

	// Minion root directory in etcd
	rootDir string

	// Minion queue directory in etcd
	queueDir string

	// Log directory of previously executed tasks
	logDir string

	// Classifier directory in etcd
	classifierDir string

	// Minion unique identifier
	id uuid.UUID

	// KeysAPI client to etcd
	kapi etcdclient.KeysAPI

	// Channel over which tasks are sent for processing
	taskQueue chan *task.Task

	// Git repo from which to sync modules and data
	gitRepo string

	// Site directory contains the modules and data synced from Git
	siteDir string

	// Channel used to signal for shutdown time
	done chan struct{}
}

// EtcdMinionConfig contains configuration settings for
// minions which use the etcd backend implementation
type EtcdMinionConfig struct {
	// Name of the minion
	Name string

	// GitRepo provides the url from which to sync modules and data
	GitRepo string

	// EtcdConfig provides etcd-related configuration settings
	EtcdConfig etcdclient.Config
}

// NewEtcdMinion creates a new minion with etcd backend
func NewEtcdMinion(config *EtcdMinionConfig) (Minion, error) {
	c, err := etcdclient.New(config.EtcdConfig)
	if err != nil {
		return nil, err
	}

	cwd, err := os.Getwd()
	if err != nil {
		return nil, err
	}

	kapi := etcdclient.NewKeysAPI(c)
	id := utils.GenerateUUID(config.Name)
	rootDir := filepath.Join(EtcdMinionSpace, id.String())
	queueDir := filepath.Join(rootDir, "queue")
	classifierDir := filepath.Join(rootDir, "classifier")
	logDir := filepath.Join(rootDir, "log")
	siteDir := filepath.Join(cwd, "site")
	taskQueue := make(chan *task.Task)
	done := make(chan struct{})

	m := &etcdMinion{
		name:          config.Name,
		rootDir:       rootDir,
		queueDir:      queueDir,
		classifierDir: classifierDir,
		logDir:        logDir,
		id:            id,
		kapi:          kapi,
		taskQueue:     taskQueue,
		gitRepo:       config.GitRepo,
		siteDir:       siteDir,
		done:          done,
	}

	return m, nil
}

// Checks for any pending tasks and sends them
// for processing if there are any
func (m *etcdMinion) checkQueue() error {
	opts := &etcdclient.GetOptions{
		Recursive: true,
		Sort:      true,
	}

	// Get backlog tasks if any
	// If the directory key in etcd is missing that is okay, since
	// it means there are no pending tasks for processing
	resp, err := m.kapi.Get(context.Background(), m.queueDir, opts)
	if err != nil {
		if eerr, ok := err.(etcdclient.Error); !ok || eerr.Code == etcdclient.ErrorCodeKeyNotFound {
			return err
		}
	}

	backlog := resp.Node.Nodes
	if len(backlog) == 0 {
		// No backlog tasks found
		return nil
	}

	log.Printf("Found %d pending tasks in queue", len(backlog))
	for _, node := range backlog {
		t, err := EtcdUnmarshalTask(node)
		m.kapi.Delete(context.Background(), node.Key, nil)

		if err != nil {
			continue
		}

		m.taskQueue <- t
	}

	return nil
}

// Runs periodic jobs such as refreshing classifiers and
// updating the lastseen time every five minutes
func (m *etcdMinion) periodicRunner() {
	schedule := time.Minute * 5
	ticker := time.NewTicker(schedule)
	log.Printf("Periodic scheduler set to run every %s\n", schedule)

	for {
		select {
		case <-m.done:
			break
		case now := <-ticker.C:
			// Run periodic jobs
			m.classify()
			m.checkQueue()
			m.SetLastseen(now.Unix())
		}
	}
}

// Processes new tasks
func (m *etcdMinion) processTask(t *task.Task) error {
	// Update state of task to indicate that we are now processing it
	t.State = task.TaskStateProcessing
	if err := m.SaveTaskResult(t); err != nil {
		return err
	}

	var buf bytes.Buffer
	fmt.Fprintf(&buf, "Loaded %d resources from catalog\n", t.Catalog.Len())
	catalogErr := t.Catalog.Run(&buf)
	t.TimeProcessed = time.Now().Unix()
	t.Result = buf.String()

	if catalogErr != nil {
		t.State = task.TaskStateFailed
	} else {
		t.State = task.TaskStateSuccess
	}

	if err := m.SaveTaskResult(t); err != nil {
		return err
	}

	return catalogErr
}

// Classifies the minion
func (m *etcdMinion) classify() error {
	for key := range classifier.Registry {
		klassifier, err := classifier.Get(key)

		if err != nil {
			log.Printf("Failed to get classifier %s: %s\n", key, err)
			continue
		}

		m.SetClassifier(klassifier)
	}

	return nil
}

// EtcdUnmarshalTask unmarshals tasks from an etcd node
func EtcdUnmarshalTask(node *etcdclient.Node) (*task.Task, error) {
	task := new(task.Task)
	err := json.Unmarshal([]byte(node.Value), &task)

	return task, err
}

// ID returns the minion unique identifier
func (m *etcdMinion) ID() uuid.UUID {
	return m.id
}

// SetName sets the human-readable name of the minion in etcd
func (m *etcdMinion) SetName(name string) error {
	nameKey := filepath.Join(m.rootDir, "name")
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), nameKey, name, opts)
	if err != nil {
		log.Printf("Failed to set name of minion: %s\n", err)
	}

	return err
}

// SetLastseen sets the time the minion was last seen in
// seconds since the Epoch
func (m *etcdMinion) SetLastseen(s int64) error {
	lastseenKey := filepath.Join(m.rootDir, "lastseen")
	lastseenValue := strconv.FormatInt(s, 10)
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), lastseenKey, lastseenValue, opts)
	if err != nil {
		log.Printf("Failed to set lastseen time: %s\n", err)
	}

	return err
}

// SetClassifier sets a classifier for the minion in etcd
func (m *etcdMinion) SetClassifier(c *classifier.Classifier) error {
	// Classifiers in etcd expire after an hour
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
		TTL:       time.Hour,
	}

	// Serialize classifier to JSON and save it in etcd
	data, err := json.Marshal(c)
	if err != nil {
		log.Printf("Failed to serialize classifier %s: %s\n", c.Key, err)
		return err
	}

	// Classifier key in etcd
	klassifierKey := filepath.Join(m.classifierDir, c.Key)
	_, err = m.kapi.Set(context.Background(), klassifierKey, string(data), opts)

	if err != nil {
		log.Printf("Failed to set classifier %s: %s\n", c.Key, err)
		return err
	}

	return nil
}

// TaskListener monitors etcd for new tasks
func (m *etcdMinion) TaskListener(c chan<- *task.Task) error {
	log.Printf("Task listener is watching %s\n", m.queueDir)

	rand.Seed(time.Now().UTC().UnixNano())
	b := backoff.Backoff{
		Min:    1 * time.Second,
		Max:    10 * time.Minute,
		Factor: 2.0,
		Jitter: true,
	}

	watcherOpts := &etcdclient.WatcherOptions{
		Recursive: true,
	}
	watcher := m.kapi.Watcher(m.queueDir, watcherOpts)

	for {
		resp, err := watcher.Next(context.Background())
		if err != nil {
			// Use a backoff and retry later again
			duration := b.Duration()
			log.Printf("%s, retrying in %s\n", err, duration)
			time.Sleep(duration)
			continue
		}

		// Reset the backoff counter on successful receive
		b.Reset()

		// Ignore "delete" events when removing a task from the queue
		action := strings.ToLower(resp.Action)
		if strings.EqualFold(action, "delete") {
			continue
		}

		// Unmarshal and remove task from the queue
		t, err := EtcdUnmarshalTask(resp.Node)
		m.kapi.Delete(context.Background(), resp.Node.Key, nil)

		if err != nil {
			log.Printf("Received invalid task %s: %s\n", resp.Node.Key, err)
			continue
		}

		// Send the task for processing
		log.Printf("Received task %s\n", t.TaskID)
		t.State = task.TaskStateQueued
		t.TimeReceived = time.Now().Unix()
		if err := m.SaveTaskResult(t); err != nil {
			log.Printf("Unable to save task state: %s\n", err)
			continue
		}

		c <- t
	}

	return nil
}

// TaskRunner processes new tasks
func (m *etcdMinion) TaskRunner(c <-chan *task.Task) error {
	log.Println("Starting task runner")

	for {
		select {
		case <-m.done:
			break
		case t := <-c:
			// Sync the module and data files first, then process the task
			// TODO: In case of failures to sync site repo, we may consider
			//       continuing the task processing from the last known state
			m.Sync()

			// Switch to the correct environment
			if err := m.setEnvironment(t.Environment); err != nil {
				log.Printf("Unable to set environment: %s\n", err)
				continue
			}

			log.Printf("Processing task %s\n", t.TaskID)
			err := m.processTask(t)
			if err != nil {
				log.Printf("Failed to process task %s: %s\n", t.TaskID, err)
			} else {
				log.Printf("Finished processing task %s\n", t.TaskID)
			}
		}
	}

	return nil
}

// SaveTaskResult stores the result of a task in etcd
func (m *etcdMinion) SaveTaskResult(t *task.Task) error {
	taskKey := filepath.Join(m.logDir, t.TaskID.String())

	data, err := json.Marshal(t)
	if err != nil {
		return err
	}

	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err = m.kapi.Set(context.Background(), taskKey, string(data), opts)

	return err
}

// setEnvironment checks out the branch with the
// provided name in the site repo dir and sets HEAD to detached
func (m *etcdMinion) setEnvironment(name string) error {
	repo, err := git.OpenRepository(m.siteDir)
	if err != nil {
		return err
	}

	remoteBranchName := filepath.Join("origin", name)
	remoteBranch, err := repo.LookupBranch(remoteBranchName, git.BranchRemote)
	if err != nil {
		return err
	}

	err = repo.SetHeadDetached(remoteBranch.Target())
	if err != nil {
		return err
	}

	checkoutOpts := &git.CheckoutOpts{
		Strategy: git.CheckoutForce,
	}

	err = repo.CheckoutHead(checkoutOpts)
	if err != nil {
		return err
	}

	return nil
}

// Sync syncs module and data files from the upstream Git repository
func (m *etcdMinion) Sync() error {
	if m.gitRepo == "" {
		log.Printf("site repo url is not set, skipping sync")
		return nil
	}

	fi, err := os.Stat(m.siteDir)
	if err != nil {
		return err
	}

	// Site directory does not exist, clone the Git repository from upstream
	if os.IsNotExist(err) {
		log.Printf("site repo is missing, performing initial sync from %s\n", m.gitRepo)
		opts := &git.CloneOptions{}
		if _, err := git.Clone(m.gitRepo, m.siteDir, opts); err != nil {
			return err
		}
	}

	// File exists, ensure that it is a valid Git repository
	if !fi.IsDir() {
		return fmt.Errorf("%s exists, but is not a directory", m.siteDir)
	}

	// Open the repository and fetch from the default remote
	log.Println("starting site repo sync")
	repo, err := git.OpenRepository(m.siteDir)
	if err != nil {
		return err
	}

	origin, err := repo.Remotes.Lookup("origin")
	if err != nil {
		return err
	}

	err = origin.Fetch([]string{}, &git.FetchOptions{}, "")
	if err != nil {
		log.Printf("site repo sync failed: %s\n", err)
	} else {
		log.Println("site repo sync completed")
	}

	return err
}

// Serve starts the minion
func (m *etcdMinion) Serve() error {
	err := m.SetName(m.name)
	if err != nil {
		return err
	}

	now := time.Now().Unix()
	err = m.SetLastseen(now)
	if err != nil {
		return err
	}

	// Start minion services
	go m.classify()
	go m.checkQueue()
	go m.periodicRunner()
	go m.TaskRunner(m.taskQueue)
	go m.TaskListener(m.taskQueue)

	log.Printf("Minion %s is ready to serve", m.ID())

	return nil
}

// Stop shutdowns the minions and its services
func (m *etcdMinion) Stop() error {
	log.Println("Minion is shutting down")

	close(m.taskQueue)
	close(m.done)

	return nil
}
<|endoftext|>


<commit_before>package types

import ()

type FieldRef struct {
	Collection string
	OnDelete   int
}

const (
	RefOnDeletePrevent = iota
	RefOnDeleteCascade
	RefOnDeleteNull
)

func (f *FieldRef) GetMysqlDef() string { return "INT(11) UNSIGNED NULL" }

func (f *FieldRef) IsSearchable() bool { return false }

func (f *FieldRef) Init(raw map[string]interface{}) error {
	collection, ok := raw["collection"]
	if !ok {
		return ModelDefinitionError{"Ref Field has no key 'collection'", ""}
	}
	collectionString := collection.(string)
	f.Collection = collectionString

	onDelete, ok := raw["onDelete"]
	if ok {
		if onDelete == "CASCADE" {
			f.OnDelete = RefOnDeleteCascade
		} else if onDelete == "NULL" {
			f.OnDelete = RefOnDeleteNull
		}
	} else {
		f.OnDelete = RefOnDeletePrevent
	}
	return nil
}

func (f *FieldRef) FromDb(stored interface{}) (interface{}, error) {
	// uInt64 -> Iunt64
	storedInt, ok := stored.(*uint64)
	if !ok {
		return nil, MakeFromDbErrorFromString("Incorrect Type in DB (expected uint64)")
	}
	if storedInt == nil {
		return nil, nil
	}
	return *storedInt, nil
}

func (f *FieldRef) ToDb(input interface{}) (string, error) {
	return unsignedIntToDb(input)
}

func (f *FieldRef) GetScanReciever() interface{} {
	var v uint64
	var vp *uint64 = &v
	return &vp
}
<commit_msg>Fix ref type<commit_after>package types

import ()

type FieldRef struct {
	Collection string
	OnDelete   int
	Limit      map[string]interface{}
}

const (
	RefOnDeletePrevent = iota
	RefOnDeleteCascade
	RefOnDeleteNull
)

func (f *FieldRef) GetMysqlDef() string { return "INT(11) UNSIGNED NULL" }

func (f *FieldRef) IsSearchable() bool { return false }

func (f *FieldRef) Init(raw map[string]interface{}) error {
	collection, ok := raw["collection"]
	if !ok {
		return ModelDefinitionError{"Ref Field has no key 'collection'", ""}
	}
	collectionString := collection.(string)
	f.Collection = collectionString

	onDelete, ok := raw["onDelete"]
	if ok {
		if onDelete == "CASCADE" {
			f.OnDelete = RefOnDeleteCascade
		} else if onDelete == "NULL" {
			f.OnDelete = RefOnDeleteNull
		}
	} else {
		f.OnDelete = RefOnDeletePrevent
	}

	f.Limit = map[string]interface{}{}

	rawLimit, ok := raw["limit"]
	if ok {
		l, ok := rawLimit.(map[string]interface{})
		if ok {
			f.Limit = l
		}
	}
	return nil
}

func (f *FieldRef) FromDb(stored interface{}) (interface{}, error) {
	// uInt64 -> Iunt64
	storedInt, ok := stored.(*uint64)
	if !ok {
		return nil, MakeFromDbErrorFromString("Incorrect Type in DB (expected uint64)")
	}
	if storedInt == nil {
		return nil, nil
	}
	return *storedInt, nil
}

func (f *FieldRef) ToDb(input interface{}) (string, error) {
	return unsignedIntToDb(input)
}

func (f *FieldRef) GetScanReciever() interface{} {
	var v uint64
	var vp *uint64 = &v
	return &vp
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/zhangpeihao/goflv"
	rtmp "github.com/zhangpeihao/gortmp"
	"github.com/zhangpeihao/log"
	"os"
	"time"
)

const (
	programName = "RtmpPublisher"
	version     = "0.0.1"
)

var (
	url         *string = flag.String("URL", "rtmp://192.168.20.111/vid3", "The rtmp url to connect.")
	streamName  *string = flag.String("Stream", "camstream", "Stream name to play.")
	flvFileName *string = flag.String("FLV", "", "Dump FLV into file.")
)

type TestOutboundConnHandler struct {
}

var obConn rtmp.OutboundConn
var createStreamChan chan rtmp.OutboundStream
var videoDataSize int64
var audioDataSize int64
var flvFile *flv.File

var status uint

func (handler *TestOutboundConnHandler) OnStatus(conn rtmp.OutboundConn) {
	var err error
	status, err = obConn.Status()
	fmt.Printf("@@@@@@@@@@@@@status: %d, err: %v\n", status, err)
}

func (handler *TestOutboundConnHandler) OnClosed(conn rtmp.Conn) {
	fmt.Printf("@@@@@@@@@@@@@Closed\n")
}

func (handler *TestOutboundConnHandler) OnReceived(conn rtmp.Conn, message *rtmp.Message) {
}

func (handler *TestOutboundConnHandler) OnReceivedRtmpCommand(conn rtmp.Conn, command *rtmp.Command) {
	fmt.Printf("ReceviedRtmpCommand: %+v\n", command)
}

func (handler *TestOutboundConnHandler) OnStreamCreated(conn rtmp.OutboundConn, stream rtmp.OutboundStream) {
	fmt.Printf("Stream created: %d\n", stream.ID())
	createStreamChan <- stream
}
func (handler *TestOutboundConnHandler) OnPlayStart(stream rtmp.OutboundStream) {

}
func (handler *TestOutboundConnHandler) OnPublishStart(stream rtmp.OutboundStream) {
	// Set chunk buffer size
	go publish(stream)
}

func publish(stream rtmp.OutboundStream) {

	var err error
	flvFile, err = flv.OpenFile(*flvFileName)
	if err != nil {
		fmt.Println("Open FLV dump file error:", err)
		return
	}
	defer flvFile.Close()
	startTs := uint32(0)
	startAt := time.Now().UnixNano()
	preTs := uint32(0)
	for status == rtmp.OUTBOUND_CONN_STATUS_CREATE_STREAM_OK {
		if flvFile.IsFinished() {
			fmt.Println("@@@@@@@@@@@@@@File finished")
			flvFile.LoopBack()
			startAt = time.Now().UnixNano()
			startTs = uint32(0)
			preTs = uint32(0)
		}
		header, data, err := flvFile.ReadTag()
		if err != nil {
			fmt.Println("flvFile.ReadTag() error:", err)
			break
		}
		switch header.TagType {
		case flv.VIDEO_TAG:
			videoDataSize += int64(len(data))
		case flv.AUDIO_TAG:
			audioDataSize += int64(len(data))
		}

		if startTs == uint32(0) {
			startTs = header.Timestamp
		}
		diff1 := uint32(0)
		//		deltaTs := uint32(0)
		if header.Timestamp > startTs {
			diff1 = header.Timestamp - startTs
		} else {
			fmt.Println("@@@@@@@@@@@@@@diff1")
		}
		if diff1 > preTs {
			//			deltaTs = diff1 - preTs
			preTs = diff1
		}
		if err = stream.PublishData(header.TagType, data, diff1); err != nil {
			fmt.Println("PublishData() error:", err)
			break
		}
		diff2 := uint32((time.Now().UnixNano() - startAt) / 1000000)
		//		fmt.Printf("diff1: %d, diff2: %d\n", diff1, diff2)
		if diff1 > diff2+100 {
			//			fmt.Printf("header.Timestamp: %d, now: %d\n", header.Timestamp, time.Now().UnixNano())
			time.Sleep(time.Millisecond * time.Duration(diff1-diff2))
		}
	}
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "%s version[%s]\r\nUsage: %s [OPTIONS]\r\n", programName, version, os.Args[0])
		flag.PrintDefaults()
	}
	flag.Parse()

	l := log.NewLogger(".", "publisher", nil, 60, 3600*24, true)
	rtmp.InitLogger(l)
	defer l.Close()
	createStreamChan = make(chan rtmp.OutboundStream)
	testHandler := &TestOutboundConnHandler{}
	fmt.Println("to dial")
	var err error
	obConn, err = rtmp.Dial(*url, testHandler, 100)
	if err != nil {
		fmt.Println("Dial error", err)
		os.Exit(-1)
	}
	defer obConn.Close()
	fmt.Println("to connect")
	err = obConn.Connect()
	if err != nil {
		fmt.Printf("Connect error: %s", err.Error())
		os.Exit(-1)
	}
	for {
		select {
		case stream := <-createStreamChan:
			// Publish
			stream.Attach(testHandler)
			err = stream.Publish(*streamName, "live")
			if err != nil {
				fmt.Printf("Publish error: %s", err.Error())
				os.Exit(-1)
			}

		case <-time.After(1 * time.Second):
			fmt.Printf("Audio size: %d bytes; Vedio size: %d bytes\n", audioDataSize, videoDataSize)
		}
	}
}
<commit_msg>Fix option text mistake<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"time"

	rtmp "github.com/zhangpeihao/gortmp"
	"github.com/zhangpeihao/log"
)

const (
	programName = "RtmpPublisher"
	version     = "0.0.1"
)

var (
	url         *string = flag.String("URL", "rtmp://192.168.20.111/vid3", "The rtmp url to connect.")
	streamName  *string = flag.String("Stream", "camstream", "Stream name to play.")
	flvFileName *string = flag.String("FLV", "", "FLV file to publishs.")
)

type TestOutboundConnHandler struct {
}

var obConn rtmp.OutboundConn
var createStreamChan chan rtmp.OutboundStream
var videoDataSize int64
var audioDataSize int64
var flvFile *flv.File

var status uint

func (handler *TestOutboundConnHandler) OnStatus(conn rtmp.OutboundConn) {
	var err error
	status, err = obConn.Status()
	fmt.Printf("@@@@@@@@@@@@@status: %d, err: %v\n", status, err)
}

func (handler *TestOutboundConnHandler) OnClosed(conn rtmp.Conn) {
	fmt.Printf("@@@@@@@@@@@@@Closed\n")
}

func (handler *TestOutboundConnHandler) OnReceived(conn rtmp.Conn, message *rtmp.Message) {
}

func (handler *TestOutboundConnHandler) OnReceivedRtmpCommand(conn rtmp.Conn, command *rtmp.Command) {
	fmt.Printf("ReceviedRtmpCommand: %+v\n", command)
}

func (handler *TestOutboundConnHandler) OnStreamCreated(conn rtmp.OutboundConn, stream rtmp.OutboundStream) {
	fmt.Printf("Stream created: %d\n", stream.ID())
	createStreamChan <- stream
}
func (handler *TestOutboundConnHandler) OnPlayStart(stream rtmp.OutboundStream) {

}
func (handler *TestOutboundConnHandler) OnPublishStart(stream rtmp.OutboundStream) {
	// Set chunk buffer size
	go publish(stream)
}

func publish(stream rtmp.OutboundStream) {

	var err error
	flvFile, err = flv.OpenFile(*flvFileName)
	if err != nil {
		fmt.Println("Open FLV dump file error:", err)
		return
	}
	defer flvFile.Close()
	startTs := uint32(0)
	startAt := time.Now().UnixNano()
	preTs := uint32(0)
	for status == rtmp.OUTBOUND_CONN_STATUS_CREATE_STREAM_OK {
		if flvFile.IsFinished() {
			fmt.Println("@@@@@@@@@@@@@@File finished")
			flvFile.LoopBack()
			startAt = time.Now().UnixNano()
			startTs = uint32(0)
			preTs = uint32(0)
		}
		header, data, err := flvFile.ReadTag()
		if err != nil {
			fmt.Println("flvFile.ReadTag() error:", err)
			break
		}
		switch header.TagType {
		case flv.VIDEO_TAG:
			videoDataSize += int64(len(data))
		case flv.AUDIO_TAG:
			audioDataSize += int64(len(data))
		}

		if startTs == uint32(0) {
			startTs = header.Timestamp
		}
		diff1 := uint32(0)
		//		deltaTs := uint32(0)
		if header.Timestamp > startTs {
			diff1 = header.Timestamp - startTs
		} else {
			fmt.Println("@@@@@@@@@@@@@@diff1")
		}
		if diff1 > preTs {
			//			deltaTs = diff1 - preTs
			preTs = diff1
		}
		if err = stream.PublishData(header.TagType, data, diff1); err != nil {
			fmt.Println("PublishData() error:", err)
			break
		}
		diff2 := uint32((time.Now().UnixNano() - startAt) / 1000000)
		//		fmt.Printf("diff1: %d, diff2: %d\n", diff1, diff2)
		if diff1 > diff2+100 {
			//			fmt.Printf("header.Timestamp: %d, now: %d\n", header.Timestamp, time.Now().UnixNano())
			time.Sleep(time.Millisecond * time.Duration(diff1-diff2))
		}
	}
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "%s version[%s]\r\nUsage: %s [OPTIONS]\r\n", programName, version, os.Args[0])
		flag.PrintDefaults()
	}
	flag.Parse()

	l := log.NewLogger(".", "publisher", nil, 60, 3600*24, true)
	rtmp.InitLogger(l)
	defer l.Close()
	createStreamChan = make(chan rtmp.OutboundStream)
	testHandler := &TestOutboundConnHandler{}
	fmt.Println("to dial")
	var err error
	obConn, err = rtmp.Dial(*url, testHandler, 100)
	if err != nil {
		fmt.Println("Dial error", err)
		os.Exit(-1)
	}
	defer obConn.Close()
	fmt.Println("to connect")
	err = obConn.Connect()
	if err != nil {
		fmt.Printf("Connect error: %s", err.Error())
		os.Exit(-1)
	}
	for {
		select {
		case stream := <-createStreamChan:
			// Publish
			stream.Attach(testHandler)
			err = stream.Publish(*streamName, "live")
			if err != nil {
				fmt.Printf("Publish error: %s", err.Error())
				os.Exit(-1)
			}

		case <-time.After(1 * time.Second):
			fmt.Printf("Audio size: %d bytes; Vedio size: %d bytes\n", audioDataSize, videoDataSize)
		}
	}
}
<|endoftext|>


<commit_before>package bridge

import (
	"fmt"
	"net"
	"os"

	"github.com/cloudfoundry-incubator/switchboard/domain"
	"github.com/pivotal-golang/lager"
)

type Runner struct {
	logger             lager.Logger
	port               uint
	trafficEnabledChan <-chan bool
	activeBackendChan  <-chan domain.Backend
}

func NewRunner(activeBackendChan <-chan domain.Backend, trafficEnabledChan <-chan bool, port uint, logger lager.Logger) Runner {
	return Runner{
		logger:             logger,
		activeBackendChan:  activeBackendChan,
		trafficEnabledChan: trafficEnabledChan,
		port:               port,
	}
}

func (r Runner) Run(signals <-chan os.Signal, ready chan<- struct{}) error {
	r.logger.Info(fmt.Sprintf("Proxy listening on port %d", r.port))

	listener, err := net.Listen("tcp", fmt.Sprintf("0.0.0.0:%d", r.port))
	if err != nil {
		return err
	}

	shutdown := make(chan interface{})
	go func(shutdown <-chan interface{}, listener net.Listener) {
		trafficEnabled := true
		var activeBackend domain.Backend
		for {

			e := make(chan error)
			c := make(chan net.Conn)

			go blockingAccept(listener, c, e)
			select {
			case <-shutdown:
				return
			case t := <-r.trafficEnabledChan:
				// ENABLED -> DISABLED
				if trafficEnabled && !t {
					if activeBackend != nil {
						activeBackend.SeverConnections()
					}
				}

				trafficEnabled = t
				close(e)
				close(c)
				continue
			case clientConn := <-c:
				if !trafficEnabled {
					clientConn.Close()
					continue
				}

				go func(clientConn net.Conn, activeBackend domain.Backend) {
					if activeBackend == nil {
						clientConn.Close()
						r.logger.Error("No active backend", err)
						return
					}

					err := activeBackend.Bridge(clientConn)
					if err != nil {
						clientConn.Close()
						r.logger.Error("Error routing to backend", err)
					}
				}(clientConn, activeBackend)
			case err := <-e:
				if err != nil {
					r.logger.Error("Error accepting client connection", err)
					continue
				}

			}
		}
	}(shutdown, listener)

	close(ready)

	signal := <-signals
	r.logger.Info("Received signal", lager.Data{"signal": signal})
	close(shutdown)
	listener.Close()

	r.logger.Info("Proxy runner has exited")
	return nil
}

func blockingAccept(l net.Listener, c chan<- net.Conn, e chan<- error) {
	clientConn, err := l.Accept()

	if err != nil {
		e <- err
		return
	}
	c <- clientConn
}
<commit_msg>Listen for changes to active backend, sever existing connections<commit_after>package bridge

import (
	"fmt"
	"net"
	"os"

	"github.com/cloudfoundry-incubator/switchboard/domain"
	"github.com/pivotal-golang/lager"
)

type Runner struct {
	logger             lager.Logger
	port               uint
	trafficEnabledChan <-chan bool
	activeBackendChan  <-chan domain.Backend
}

func NewRunner(activeBackendChan <-chan domain.Backend, trafficEnabledChan <-chan bool, port uint, logger lager.Logger) Runner {
	return Runner{
		logger:             logger,
		activeBackendChan:  activeBackendChan,
		trafficEnabledChan: trafficEnabledChan,
		port:               port,
	}
}

func (r Runner) Run(signals <-chan os.Signal, ready chan<- struct{}) error {
	r.logger.Info(fmt.Sprintf("Proxy listening on port %d", r.port))

	listener, err := net.Listen("tcp", fmt.Sprintf("0.0.0.0:%d", r.port))
	if err != nil {
		return err
	}

	shutdown := make(chan interface{})
	go func(shutdown <-chan interface{}, listener net.Listener) {
		trafficEnabled := true
		var activeBackend domain.Backend
		for {

			e := make(chan error)
			c := make(chan net.Conn)

			go blockingAccept(listener, c, e)
			select {
			case <-shutdown:
				return
			case t := <-r.trafficEnabledChan:
				// ENABLED -> DISABLED
				if trafficEnabled && !t {
					if activeBackend != nil {
						activeBackend.SeverConnections()
					}
				}

				trafficEnabled = t
				close(e)
				close(c)
				continue
			case a := <- r.activeBackendChan:
				// NEW ACTIVE BACKEND
				if activeBackend != nil {
					activeBackend.SeverConnections()
				}

				activeBackend = a
				close(e)
				close(c)
				continue
			case clientConn := <-c:
				if !trafficEnabled {
					clientConn.Close()
					continue
				}

				go func(clientConn net.Conn, activeBackend domain.Backend) {
					if activeBackend == nil {
						clientConn.Close()
						r.logger.Error("No active backend", err)
						return
					}

					err := activeBackend.Bridge(clientConn)
					if err != nil {
						clientConn.Close()
						r.logger.Error("Error routing to backend", err)
					}
				}(clientConn, activeBackend)
			case err := <-e:
				if err != nil {
					r.logger.Error("Error accepting client connection", err)
					continue
				}
			}
		}
	}(shutdown, listener)

	close(ready)

	signal := <-signals
	r.logger.Info("Received signal", lager.Data{"signal": signal})
	close(shutdown)
	listener.Close()

	r.logger.Info("Proxy runner has exited")
	return nil
}

func blockingAccept(l net.Listener, c chan<- net.Conn, e chan<- error) {
	clientConn, err := l.Accept()

	if err != nil {
		e <- err
		return
	}
	c <- clientConn
}
<|endoftext|>