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

package sql

import (
	"errors"
	"time"

	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/util"
	"github.com/gogo/protobuf/proto"
)

// Future home of the aysynchronous schema changer that picks up
// queued schema changes and processes them.
//
// applyMutations applies the queued mutations for a table.
func (p *planner) applyMutations(tableDesc *TableDescriptor) error {
	if len(tableDesc.Mutations) == 0 {
		return nil
	}
	newTableDesc := proto.Clone(tableDesc).(*TableDescriptor)
	p.applyUpVersion(newTableDesc)
	// Make all mutations active.
	for _, mutation := range newTableDesc.Mutations {
		newTableDesc.makeMutationComplete(mutation)
	}
	newTableDesc.Mutations = nil
	if err := newTableDesc.Validate(); err != nil {
		return err
	}

	b := client.Batch{}
	if err := p.backfillBatch(&b, tableDesc, newTableDesc); err != nil {
		return err
	}

	b.Put(MakeDescMetadataKey(newTableDesc.GetID()), wrapDescriptor(newTableDesc))

	if err := p.txn.Run(&b); err != nil {
		return convertBatchError(newTableDesc, b, err)
	}
	p.notifyCompletedSchemaChange(newTableDesc.ID)
	return nil
}

// applyUpVersion only increments the version of the table descriptor.
func (p *planner) applyUpVersion(tableDesc *TableDescriptor) {
	if tableDesc.UpVersion {
		tableDesc.UpVersion = false
		tableDesc.Version++
	}
}

// SchemaChanger is used to change the schema on a table.
type SchemaChanger struct {
	tableID ID
	nodeID  roachpb.NodeID
	db      client.DB
}

// NewSchemaChangerForTesting only for tests.
func NewSchemaChangerForTesting(tableID ID, nodeID roachpb.NodeID, db client.DB) SchemaChanger {
	return SchemaChanger{tableID: tableID, nodeID: nodeID, db: db}
}

var errExistingSchemaChangeLease = errors.New("an outstanding schema change lease exists")

func (sc *SchemaChanger) createSchemaChangeLease() TableDescriptor_SchemaChangeLease {
	return TableDescriptor_SchemaChangeLease{NodeID: sc.nodeID, ExpirationTime: time.Now().Add(jitteredLeaseDuration()).UnixNano()}
}

// AcquireLease acquires a schema change lease on the table if
// an unexpired lease doesn't exist. It returns the lease.
func (sc *SchemaChanger) AcquireLease() (TableDescriptor_SchemaChangeLease, error) {
	var lease TableDescriptor_SchemaChangeLease
	err := sc.db.Txn(func(txn *client.Txn) error {
		txn.SetSystemDBTrigger()
		tableDesc, err := getTableDescFromID(txn, sc.tableID)
		if err != nil {
			return err
		}

		// A second to deal with the time uncertainty across nodes.
		// It is perfectly valid for two or more goroutines to hold a valid
		// lease and execute a schema change in parallel, because schema
		// changes are executed using transactions that run sequentially.
		// This just reduces the probability of a write collision.
		expirationTimeUncertainty := time.Second

		if tableDesc.Lease != nil && time.Unix(0, tableDesc.Lease.ExpirationTime).Add(expirationTimeUncertainty).After(time.Now()) {
			return errExistingSchemaChangeLease
		}
		lease = sc.createSchemaChangeLease()
		tableDesc.Lease = &lease
		return txn.Put(MakeDescMetadataKey(tableDesc.ID), wrapDescriptor(tableDesc))
	})
	return lease, err
}

func (sc *SchemaChanger) findTableWithLease(txn *client.Txn, lease TableDescriptor_SchemaChangeLease) (*TableDescriptor, error) {
	tableDesc, err := getTableDescFromID(txn, sc.tableID)
	if err != nil {
		return nil, err
	}
	if tableDesc.Lease == nil {
		return nil, util.Errorf("no lease present for tableID: %d", sc.tableID)
	}
	if *tableDesc.Lease != lease {
		return nil, util.Errorf("table: %d has lease: %v, expected: %v", sc.tableID, tableDesc.Lease, lease)
	}
	return tableDesc, nil
}

// ReleaseLease the table lease if it is the one registered with
// the table descriptor.
func (sc *SchemaChanger) ReleaseLease(lease TableDescriptor_SchemaChangeLease) error {
	return sc.db.Txn(func(txn *client.Txn) error {
		tableDesc, err := sc.findTableWithLease(txn, lease)
		if err != nil {
			return err
		}
		tableDesc.Lease = nil
		txn.SetSystemDBTrigger()
		return txn.Put(MakeDescMetadataKey(tableDesc.ID), wrapDescriptor(tableDesc))
	})
}

// ExtendLease for the current leaser.
func (sc *SchemaChanger) ExtendLease(lease TableDescriptor_SchemaChangeLease) (TableDescriptor_SchemaChangeLease, error) {
	err := sc.db.Txn(func(txn *client.Txn) error {
		tableDesc, err := sc.findTableWithLease(txn, lease)
		if err != nil {
			return err
		}
		lease = sc.createSchemaChangeLease()
		tableDesc.Lease = &lease
		txn.SetSystemDBTrigger()
		return txn.Put(MakeDescMetadataKey(tableDesc.ID), wrapDescriptor(tableDesc))
	})
	return lease, err
}
<commit_msg>sql: typo fix<commit_after>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License. See the AUTHORS file
// for names of contributors.
//
// Author: Vivek Menezes (vivek@cockroachlabs.com)

package sql

import (
	"errors"
	"time"

	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/util"
	"github.com/gogo/protobuf/proto"
)

// Future home of the asynchronous schema changer that picks up
// queued schema changes and processes them.
//
// applyMutations applies the queued mutations for a table.
func (p *planner) applyMutations(tableDesc *TableDescriptor) error {
	if len(tableDesc.Mutations) == 0 {
		return nil
	}
	newTableDesc := proto.Clone(tableDesc).(*TableDescriptor)
	p.applyUpVersion(newTableDesc)
	// Make all mutations active.
	for _, mutation := range newTableDesc.Mutations {
		newTableDesc.makeMutationComplete(mutation)
	}
	newTableDesc.Mutations = nil
	if err := newTableDesc.Validate(); err != nil {
		return err
	}

	b := client.Batch{}
	if err := p.backfillBatch(&b, tableDesc, newTableDesc); err != nil {
		return err
	}

	b.Put(MakeDescMetadataKey(newTableDesc.GetID()), wrapDescriptor(newTableDesc))

	if err := p.txn.Run(&b); err != nil {
		return convertBatchError(newTableDesc, b, err)
	}
	p.notifyCompletedSchemaChange(newTableDesc.ID)
	return nil
}

// applyUpVersion only increments the version of the table descriptor.
func (p *planner) applyUpVersion(tableDesc *TableDescriptor) {
	if tableDesc.UpVersion {
		tableDesc.UpVersion = false
		tableDesc.Version++
	}
}

// SchemaChanger is used to change the schema on a table.
type SchemaChanger struct {
	tableID ID
	nodeID  roachpb.NodeID
	db      client.DB
}

// NewSchemaChangerForTesting only for tests.
func NewSchemaChangerForTesting(tableID ID, nodeID roachpb.NodeID, db client.DB) SchemaChanger {
	return SchemaChanger{tableID: tableID, nodeID: nodeID, db: db}
}

var errExistingSchemaChangeLease = errors.New("an outstanding schema change lease exists")

func (sc *SchemaChanger) createSchemaChangeLease() TableDescriptor_SchemaChangeLease {
	return TableDescriptor_SchemaChangeLease{NodeID: sc.nodeID, ExpirationTime: time.Now().Add(jitteredLeaseDuration()).UnixNano()}
}

// AcquireLease acquires a schema change lease on the table if
// an unexpired lease doesn't exist. It returns the lease.
func (sc *SchemaChanger) AcquireLease() (TableDescriptor_SchemaChangeLease, error) {
	var lease TableDescriptor_SchemaChangeLease
	err := sc.db.Txn(func(txn *client.Txn) error {
		txn.SetSystemDBTrigger()
		tableDesc, err := getTableDescFromID(txn, sc.tableID)
		if err != nil {
			return err
		}

		// A second to deal with the time uncertainty across nodes.
		// It is perfectly valid for two or more goroutines to hold a valid
		// lease and execute a schema change in parallel, because schema
		// changes are executed using transactions that run sequentially.
		// This just reduces the probability of a write collision.
		expirationTimeUncertainty := time.Second

		if tableDesc.Lease != nil && time.Unix(0, tableDesc.Lease.ExpirationTime).Add(expirationTimeUncertainty).After(time.Now()) {
			return errExistingSchemaChangeLease
		}
		lease = sc.createSchemaChangeLease()
		tableDesc.Lease = &lease
		return txn.Put(MakeDescMetadataKey(tableDesc.ID), wrapDescriptor(tableDesc))
	})
	return lease, err
}

func (sc *SchemaChanger) findTableWithLease(txn *client.Txn, lease TableDescriptor_SchemaChangeLease) (*TableDescriptor, error) {
	tableDesc, err := getTableDescFromID(txn, sc.tableID)
	if err != nil {
		return nil, err
	}
	if tableDesc.Lease == nil {
		return nil, util.Errorf("no lease present for tableID: %d", sc.tableID)
	}
	if *tableDesc.Lease != lease {
		return nil, util.Errorf("table: %d has lease: %v, expected: %v", sc.tableID, tableDesc.Lease, lease)
	}
	return tableDesc, nil
}

// ReleaseLease the table lease if it is the one registered with
// the table descriptor.
func (sc *SchemaChanger) ReleaseLease(lease TableDescriptor_SchemaChangeLease) error {
	return sc.db.Txn(func(txn *client.Txn) error {
		tableDesc, err := sc.findTableWithLease(txn, lease)
		if err != nil {
			return err
		}
		tableDesc.Lease = nil
		txn.SetSystemDBTrigger()
		return txn.Put(MakeDescMetadataKey(tableDesc.ID), wrapDescriptor(tableDesc))
	})
}

// ExtendLease for the current leaser.
func (sc *SchemaChanger) ExtendLease(lease TableDescriptor_SchemaChangeLease) (TableDescriptor_SchemaChangeLease, error) {
	err := sc.db.Txn(func(txn *client.Txn) error {
		tableDesc, err := sc.findTableWithLease(txn, lease)
		if err != nil {
			return err
		}
		lease = sc.createSchemaChangeLease()
		tableDesc.Lease = &lease
		txn.SetSystemDBTrigger()
		return txn.Put(MakeDescMetadataKey(tableDesc.ID), wrapDescriptor(tableDesc))
	})
	return lease, err
}
<|endoftext|>


<commit_before>package kv

import (
	"encoding/json"
	"fmt"
	"strings"

	"github.com/movio/kasper/util"
	"golang.org/x/net/context"
	elastic "gopkg.in/olivere/elastic.v5"
)

// ElasticsearchKeyValueStore is a key-value storage that uses ElasticSearch.
// In this key-value store, all keys must have the format "<index>/<type>/<_id>".
// See: https://www.elastic.co/products/elasticsearch
type ElasticsearchKeyValueStore struct {
	witness            *util.StructPtrWitness
	client             *elastic.Client
	context            context.Context
	existingIndexNames []string
}

// NewElasticsearchKeyValueStore creates new ElasticsearchKeyValueStore instance.
// Host must of the format hostname:port.
// StructPtr should be a pointer to struct type that is used
// for serialization and deserialization of store values.
func NewElasticsearchKeyValueStore(host string, structPtr interface{}) *ElasticsearchKeyValueStore {
	url := fmt.Sprintf("http://%s", host)
	client, err := elastic.NewClient(
		elastic.SetURL(url),
	)
	if err != nil {
		panic(fmt.Sprintf("Cannot create ElasticSearch Client to '%s': %s", url, err))
	}
	return &ElasticsearchKeyValueStore{
		witness: util.NewStructPtrWitness(structPtr),
		client:  client,
		context: context.Background(),
	}
}

func (s *ElasticsearchKeyValueStore) checkOrCreateIndex(indexName string) {
	for _, existingIndexName := range s.existingIndexNames {
		if existingIndexName == indexName {
			return
		}
	}
	exists, err := s.client.IndexExists(indexName).Do(s.context)
	if err != nil {
		panic(fmt.Sprintf("Failed to check if index exists: %s", err))
	}
	if !exists {
		_, err := s.client.CreateIndex(indexName).Do(s.context)
		if err != nil {
			panic(fmt.Sprintf("Failed to create index: %s", err))
		}
	}
	s.existingIndexNames = append(s.existingIndexNames, indexName)
}

// Get gets value by key from store
func (s *ElasticsearchKeyValueStore) Get(key string) (interface{}, error) {
	keyParts := strings.Split(key, "/")
	if len(keyParts) != 3 {
		return nil, fmt.Errorf("invalid key: '%s'", key)
	}
	indexName := keyParts[0]
	indexType := keyParts[1]
	valueID := keyParts[2]

	s.checkOrCreateIndex(indexName)

	rawValue, err := s.client.Get().
		Index(indexName).
		Type(indexType).
		Id(valueID).
		Do(s.context)

	if fmt.Sprintf("%s", err) == "elastic: Error 404 (Not Found)" {
		return s.witness.Nil(), nil
	}

	if err != nil {
		return s.witness.Nil(), err
	}

	if !rawValue.Found {
		return s.witness.Nil(), nil
	}

	structPtr := s.witness.Allocate()
	err = json.Unmarshal(*rawValue.Source, structPtr)
	if err != nil {
		return s.witness.Nil(), err
	}
	return structPtr, nil
}

// Put updates key in store with serialized value
func (s *ElasticsearchKeyValueStore) Put(key string, structPtr interface{}) error {
	s.witness.Assert(structPtr)
	keyParts := strings.Split(key, "/")
	if len(keyParts) != 3 {
		return fmt.Errorf("invalid key: '%s'", key)
	}
	indexName := keyParts[0]
	indexType := keyParts[1]
	valueID := keyParts[2]

	s.checkOrCreateIndex(indexName)

	_, err := s.client.Index().
		Index(indexName).
		Type(indexType).
		Id(valueID).
		BodyJson(structPtr).
		Do(s.context)

	return err
}

// Delete removes key from store
func (s *ElasticsearchKeyValueStore) Delete(key string) error {
	keyParts := strings.Split(key, "/")
	if len(keyParts) != 3 {
		return fmt.Errorf("invalid key: '%s'", key)
	}
	indexName := keyParts[0]
	indexType := keyParts[1]
	valueID := keyParts[2]

	s.checkOrCreateIndex(indexName)

	_, err := s.client.Delete().
		Index(indexName).
		Type(indexType).
		Id(valueID).
		Do(s.context)

	return err
}
<commit_msg>Disable Sniff in ES KV (workaround for ES in docker)<commit_after>package kv

import (
	"encoding/json"
	"fmt"
	"strings"

	"github.com/movio/kasper/util"
	"golang.org/x/net/context"
	elastic "gopkg.in/olivere/elastic.v5"
)

// ElasticsearchKeyValueStore is a key-value storage that uses ElasticSearch.
// In this key-value store, all keys must have the format "<index>/<type>/<_id>".
// See: https://www.elastic.co/products/elasticsearch
type ElasticsearchKeyValueStore struct {
	witness            *util.StructPtrWitness
	client             *elastic.Client
	context            context.Context
	existingIndexNames []string
}

// NewElasticsearchKeyValueStore creates new ElasticsearchKeyValueStore instance.
// Host must of the format hostname:port.
// StructPtr should be a pointer to struct type that is used
// for serialization and deserialization of store values.
func NewElasticsearchKeyValueStore(host string, structPtr interface{}) *ElasticsearchKeyValueStore {
	url := fmt.Sprintf("http://%s", host)
	client, err := elastic.NewClient(
		elastic.SetURL(url),
		elastic.SetSniff(false), // FIXME: workaround for issues with ES in docker
	)
	if err != nil {
		panic(fmt.Sprintf("Cannot create ElasticSearch Client to '%s': %s", url, err))
	}
	return &ElasticsearchKeyValueStore{
		witness: util.NewStructPtrWitness(structPtr),
		client:  client,
		context: context.Background(),
	}
}

func (s *ElasticsearchKeyValueStore) checkOrCreateIndex(indexName string) {
	for _, existingIndexName := range s.existingIndexNames {
		if existingIndexName == indexName {
			return
		}
	}
	exists, err := s.client.IndexExists(indexName).Do(s.context)
	if err != nil {
		panic(fmt.Sprintf("Failed to check if index exists: %s", err))
	}
	if !exists {
		_, err := s.client.CreateIndex(indexName).Do(s.context)
		if err != nil {
			panic(fmt.Sprintf("Failed to create index: %s", err))
		}
	}
	s.existingIndexNames = append(s.existingIndexNames, indexName)
}

// Get gets value by key from store
func (s *ElasticsearchKeyValueStore) Get(key string) (interface{}, error) {
	keyParts := strings.Split(key, "/")
	if len(keyParts) != 3 {
		return nil, fmt.Errorf("invalid key: '%s'", key)
	}
	indexName := keyParts[0]
	indexType := keyParts[1]
	valueID := keyParts[2]

	s.checkOrCreateIndex(indexName)

	rawValue, err := s.client.Get().
		Index(indexName).
		Type(indexType).
		Id(valueID).
		Do(s.context)

	if fmt.Sprintf("%s", err) == "elastic: Error 404 (Not Found)" {
		return s.witness.Nil(), nil
	}

	if err != nil {
		return s.witness.Nil(), err
	}

	if !rawValue.Found {
		return s.witness.Nil(), nil
	}

	structPtr := s.witness.Allocate()
	err = json.Unmarshal(*rawValue.Source, structPtr)
	if err != nil {
		return s.witness.Nil(), err
	}
	return structPtr, nil
}

// Put updates key in store with serialized value
func (s *ElasticsearchKeyValueStore) Put(key string, structPtr interface{}) error {
	s.witness.Assert(structPtr)
	keyParts := strings.Split(key, "/")
	if len(keyParts) != 3 {
		return fmt.Errorf("invalid key: '%s'", key)
	}
	indexName := keyParts[0]
	indexType := keyParts[1]
	valueID := keyParts[2]

	s.checkOrCreateIndex(indexName)

	_, err := s.client.Index().
		Index(indexName).
		Type(indexType).
		Id(valueID).
		BodyJson(structPtr).
		Do(s.context)

	return err
}

// Delete removes key from store
func (s *ElasticsearchKeyValueStore) Delete(key string) error {
	keyParts := strings.Split(key, "/")
	if len(keyParts) != 3 {
		return fmt.Errorf("invalid key: '%s'", key)
	}
	indexName := keyParts[0]
	indexType := keyParts[1]
	valueID := keyParts[2]

	s.checkOrCreateIndex(indexName)

	_, err := s.client.Delete().
		Index(indexName).
		Type(indexType).
		Id(valueID).
		Do(s.context)

	return err
}
<|endoftext|>


<commit_before>/* linked_list.go : Implement linked list with too many memory allocations
Author: James Fairbanks
Date: 2012-09-02
Liscence: BSD
*/

package list

import (
//    "fmt"
//    "github.com/jpfairbanks/timing"
//    "math/rand"
)

//Dtype: the data type of the arrays
type Dtype int64

var NULL *Node

//NULL := 0

//Node: Struct containing a datum and a pointer to the next Node.
type Node struct {
	Datum Dtype
	Next  *Node
}

//List: Struct for containing the head of the list.
type List struct {
	Head *Node
}

//New: Make a new list of Nodes.
func New() List {
	var ell List
	head := &Node{Datum: 0, Next: NULL}
	ell.Head = head
	return ell
}

//Insert: Insert a Node into the list.
func (ell List) Insert(pos int, element Dtype) {
	node := ell.Head
	for i := 0; i < pos; i++ {
		node = node.Next
	}
	next := node.Next
	node.Next = &Node{Datum: element, Next: next}
}

//Remove: Remove a Node into the list.
func (ell List) Remove(pos int) Dtype {
	node := ell.Head
	for i := 0; i < pos; i++ {
		node = node.Next
	}
	next := node.Next
	node.Next = node.Next.Next
	return next.Datum
}

//Strider: a slice of type dtype with a fixed stride
type Strider interface {
	Walk() Dtype
}
<commit_msg>methods must take pointers to modify the List.Head pointer; add String for printing lists<commit_after>/* linked_list.go : Implement linked list with too many memory allocations
Author: James Fairbanks
Date: 2012-09-02
Liscence: BSD
*/

package list

import (
	"fmt"
	"strings"

//    "github.com/jpfairbanks/timing"
//    "math/rand"
)

//Dtype: the data type of the arrays
type Dtype int64

var NULL *Node

//NULL := 0

//Node: Struct containing a datum and a pointer to the next Node.
type Node struct {
	Datum Dtype
	Next  *Node
}

//List: Struct for containing the head of the list.
type List struct {
	Head *Node
}

//New: Make a new list of Nodes.
func New() List {
	var ell List
	head := &Node{Datum: 0, Next: NULL}
	ell.Head = head
	return ell
}

//Insert: Insert a Node into the list.
func (ell *List) Insert(pos int, element Dtype) {
	var prev *Node
	node := ell.Head
	for i := 0; i < pos; i++ {
		//fmt.Println("getting the next")
		if i == pos-1 {
			prev = node
		}
		node = node.Next
	}
	//fmt.Printf("found pos node: %v\n", node)
	//fmt.Printf("found prev: %v\n", prev)
	newNode := &Node{Datum: element, Next: node}
	//fmt.Printf("made newNode: %v\n", newNode)
	if pos > 0 {
		prev.Next = newNode
	} else {
		ell.Head = newNode
	}
	//fmt.Println("did previous")
	//fmt.Println(node.Next)
}

//Remove: Remove a Node into the list.
func (ell *List) Remove(pos int) Dtype {
	node := ell.Head
	for i := 0; i < pos; i++ {
		node = node.Next
	}
	next := node.Next
	node.Next = node.Next.Next
	return next.Datum
}

//String: print  list by printing each element on a line
func (ell *List) String() string {
	var curr *Node
	curr = ell.Head
	var repr map[int]string
	repr = make(map[int]string)
	var sterm string
	var i int
	for curr != nil {
		sterm = fmt.Sprintf("%d %v %v",
			i, curr.Datum, curr.Next)
		//fmt.Println(sterm)
		repr[i] = sterm
		i++
		curr = curr.Next
	}
	ordered := make([]string, len(repr))
	for j := 0; j < i; j++ {
		ordered[j] = repr[j]
	}
	var joined string
	joined = strings.Join(ordered, "\n")
	//fmt.Println(joined)
	return joined
}

//Strider: a slice of type dtype with a fixed stride
type Strider interface {
	Walk() Dtype
}
<|endoftext|>


<commit_before>package cc_messages

import (
	"encoding/json"

	"github.com/cloudfoundry-incubator/bbs/models"
)

type HealthCheckType string

const UnspecifiedHealthCheckType HealthCheckType = "" // backwards-compatibility
const PortHealthCheckType HealthCheckType = "port"
const NoneHealthCheckType HealthCheckType = "none"

const CC_HTTP_ROUTES = "http_routes"

const CC_TCP_ROUTES = "tcp_routes"

const (
	TaskStatePending   = "PENDING"
	TaskStateRunning   = "RUNNING"
	TaskStateCanceling = "CANCELING"
	TaskStateSucceeded = "SUCCEEDED"
)

type DesireAppRequestFromCC struct {
	ProcessGuid                 string                        `json:"process_guid"`
	DropletUri                  string                        `json:"droplet_uri"`
	DockerImageUrl              string                        `json:"docker_image"`
	DockerLoginServer           string                        `json:"docker_login_server,omitempty"`
	DockerUser                  string                        `json:"docker_user,omitempty"`
	DockerPassword              string                        `json:"docker_password,omitempty"`
	DockerEmail                 string                        `json:"docker_email,omitempty"`
	Stack                       string                        `json:"stack"`
	StartCommand                string                        `json:"start_command"`
	ExecutionMetadata           string                        `json:"execution_metadata"`
	Environment                 []*models.EnvironmentVariable `json:"environment"`
	MemoryMB                    int                           `json:"memory_mb"`
	DiskMB                      int                           `json:"disk_mb"`
	FileDescriptors             uint64                        `json:"file_descriptors"`
	NumInstances                int                           `json:"num_instances"`
	RoutingInfo                 CCRouteInfo                   `json:"routing_info"`
	AllowSSH                    bool                          `json:"allow_ssh"`
	LogGuid                     string                        `json:"log_guid"`
	HealthCheckType             HealthCheckType               `json:"health_check_type"`
	HealthCheckTimeoutInSeconds uint                          `json:"health_check_timeout_in_seconds"`
	EgressRules                 []*models.SecurityGroupRule   `json:"egress_rules,omitempty"`
	ETag                        string                        `json:"etag"`
	Ports                       []uint32                      `json:"ports,omitempty"`
	LogSource                   string                        `json:"log_source,omitempty"`
}

type CCRouteInfo map[string]*json.RawMessage

type CCHTTPRoutes []CCHTTPRoute

func (r CCHTTPRoutes) CCRouteInfo() (CCRouteInfo, error) {
	routesJson, err := json.Marshal(r)
	if err != nil {
		return nil, err
	}

	routesPayload := json.RawMessage(routesJson)
	routingInfo := make(map[string]*json.RawMessage)
	routingInfo[CC_HTTP_ROUTES] = &routesPayload
	return routingInfo, nil
}

type CCHTTPRoute struct {
	Hostname        string `json:"hostname"`
	RouteServiceUrl string `json:"route_service_url,omitempty"`
	Port            uint32 `json:"port,omitempty"`
}

type CCTCPRoutes []CCTCPRoute

type CCTCPRoute struct {
	RouterGroupGuid string `json:"router_group_guid"`
	ExternalPort    uint32 `json:"external_port,omitempty"`
	ContainerPort   uint32 `json:"container_port,omitempty"`
}

func (r CCTCPRoutes) CCRouteInfo() (CCRouteInfo, error) {
	routesJson, err := json.Marshal(r)
	if err != nil {
		return nil, err
	}

	routesPayload := json.RawMessage(routesJson)
	routingInfo := make(map[string]*json.RawMessage)
	routingInfo[CC_TCP_ROUTES] = &routesPayload
	return routingInfo, nil
}

type CCDesiredStateServerResponse struct {
	Apps        []DesireAppRequestFromCC `json:"apps"`
	CCBulkToken *json.RawMessage         `json:"token"`
}

type CCDesiredAppFingerprint struct {
	ProcessGuid string `json:"process_guid"`
	ETag        string `json:"etag"`
}

type CCTaskState struct {
	TaskGuid              string `json:"task_guid"`
	State                 string `json:"state"`
	CompletionCallbackUrl string `json:"completion_callback"`
}

type CCDesiredStateFingerprintResponse struct {
	Fingerprints []CCDesiredAppFingerprint `json:"fingerprints"`
	CCBulkToken  *json.RawMessage          `json:"token"`
}

type CCTaskStatesResponse struct {
	TaskStates  []CCTaskState    `json:"task_states"`
	CCBulkToken *json.RawMessage `json:"token"`
}

type CCBulkToken struct {
	Id int `json:"id"`
}

type TaskErrorID string

type TaskRequestFromCC struct {
	TaskGuid              string                        `json:"task_guid"`
	LogGuid               string                        `json:"log_guid"`
	MemoryMb              int                           `json:"memory_mb"`
	DiskMb                int                           `json:"disk_mb"`
	Lifecycle             string                        `json:"lifecycle"`
	EnvironmentVariables  []*models.EnvironmentVariable `json:"environment"`
	EgressRules           []*models.SecurityGroupRule   `json:"egress_rules,omitempty"`
	DropletUri            string                        `json:"droplet_uri"`
	DockerPath            string                        `json:"docker_path"`
	RootFs                string                        `json:"rootfs"`
	CompletionCallbackUrl string                        `json:"completion_callback"`
	Command               string                        `json:"command"`
	LogSource             string                        `json:"log_source,omit_empty"`
}

type TaskFailResponseForCC struct {
	TaskGuid      string `json:"task_guid"`
	Failed        bool   `json:"failed"`
	FailureReason string `json:"failure_reason"`
}

type TaskError struct {
	Id      TaskErrorID `json:"id"`
	Message string      `json:"message"`
}
<commit_msg>Add network field to the desired app message<commit_after>package cc_messages

import (
	"encoding/json"

	"github.com/cloudfoundry-incubator/bbs/models"
)

type HealthCheckType string

const UnspecifiedHealthCheckType HealthCheckType = "" // backwards-compatibility
const PortHealthCheckType HealthCheckType = "port"
const NoneHealthCheckType HealthCheckType = "none"

const CC_HTTP_ROUTES = "http_routes"

const CC_TCP_ROUTES = "tcp_routes"

const (
	TaskStatePending   = "PENDING"
	TaskStateRunning   = "RUNNING"
	TaskStateCanceling = "CANCELING"
	TaskStateSucceeded = "SUCCEEDED"
)

type DesireAppRequestFromCC struct {
	ProcessGuid                 string                        `json:"process_guid"`
	DropletUri                  string                        `json:"droplet_uri"`
	DockerImageUrl              string                        `json:"docker_image"`
	DockerLoginServer           string                        `json:"docker_login_server,omitempty"`
	DockerUser                  string                        `json:"docker_user,omitempty"`
	DockerPassword              string                        `json:"docker_password,omitempty"`
	DockerEmail                 string                        `json:"docker_email,omitempty"`
	Stack                       string                        `json:"stack"`
	StartCommand                string                        `json:"start_command"`
	ExecutionMetadata           string                        `json:"execution_metadata"`
	Environment                 []*models.EnvironmentVariable `json:"environment"`
	MemoryMB                    int                           `json:"memory_mb"`
	DiskMB                      int                           `json:"disk_mb"`
	FileDescriptors             uint64                        `json:"file_descriptors"`
	NumInstances                int                           `json:"num_instances"`
	RoutingInfo                 CCRouteInfo                   `json:"routing_info"`
	AllowSSH                    bool                          `json:"allow_ssh"`
	LogGuid                     string                        `json:"log_guid"`
	HealthCheckType             HealthCheckType               `json:"health_check_type"`
	HealthCheckTimeoutInSeconds uint                          `json:"health_check_timeout_in_seconds"`
	EgressRules                 []*models.SecurityGroupRule   `json:"egress_rules,omitempty"`
	ETag                        string                        `json:"etag"`
	Ports                       []uint32                      `json:"ports,omitempty"`
	LogSource                   string                        `json:"log_source,omitempty"`
	Network                     *models.Network               `json:"network,omitempty"`
}

type CCRouteInfo map[string]*json.RawMessage

type CCHTTPRoutes []CCHTTPRoute

func (r CCHTTPRoutes) CCRouteInfo() (CCRouteInfo, error) {
	routesJson, err := json.Marshal(r)
	if err != nil {
		return nil, err
	}

	routesPayload := json.RawMessage(routesJson)
	routingInfo := make(map[string]*json.RawMessage)
	routingInfo[CC_HTTP_ROUTES] = &routesPayload
	return routingInfo, nil
}

type CCHTTPRoute struct {
	Hostname        string `json:"hostname"`
	RouteServiceUrl string `json:"route_service_url,omitempty"`
	Port            uint32 `json:"port,omitempty"`
}

type CCTCPRoutes []CCTCPRoute

type CCTCPRoute struct {
	RouterGroupGuid string `json:"router_group_guid"`
	ExternalPort    uint32 `json:"external_port,omitempty"`
	ContainerPort   uint32 `json:"container_port,omitempty"`
}

func (r CCTCPRoutes) CCRouteInfo() (CCRouteInfo, error) {
	routesJson, err := json.Marshal(r)
	if err != nil {
		return nil, err
	}

	routesPayload := json.RawMessage(routesJson)
	routingInfo := make(map[string]*json.RawMessage)
	routingInfo[CC_TCP_ROUTES] = &routesPayload
	return routingInfo, nil
}

type CCDesiredStateServerResponse struct {
	Apps        []DesireAppRequestFromCC `json:"apps"`
	CCBulkToken *json.RawMessage         `json:"token"`
}

type CCDesiredAppFingerprint struct {
	ProcessGuid string `json:"process_guid"`
	ETag        string `json:"etag"`
}

type CCTaskState struct {
	TaskGuid              string `json:"task_guid"`
	State                 string `json:"state"`
	CompletionCallbackUrl string `json:"completion_callback"`
}

type CCDesiredStateFingerprintResponse struct {
	Fingerprints []CCDesiredAppFingerprint `json:"fingerprints"`
	CCBulkToken  *json.RawMessage          `json:"token"`
}

type CCTaskStatesResponse struct {
	TaskStates  []CCTaskState    `json:"task_states"`
	CCBulkToken *json.RawMessage `json:"token"`
}

type CCBulkToken struct {
	Id int `json:"id"`
}

type TaskErrorID string

type TaskRequestFromCC struct {
	TaskGuid              string                        `json:"task_guid"`
	LogGuid               string                        `json:"log_guid"`
	MemoryMb              int                           `json:"memory_mb"`
	DiskMb                int                           `json:"disk_mb"`
	Lifecycle             string                        `json:"lifecycle"`
	EnvironmentVariables  []*models.EnvironmentVariable `json:"environment"`
	EgressRules           []*models.SecurityGroupRule   `json:"egress_rules,omitempty"`
	DropletUri            string                        `json:"droplet_uri"`
	DockerPath            string                        `json:"docker_path"`
	RootFs                string                        `json:"rootfs"`
	CompletionCallbackUrl string                        `json:"completion_callback"`
	Command               string                        `json:"command"`
	LogSource             string                        `json:"log_source,omit_empty"`
}

type TaskFailResponseForCC struct {
	TaskGuid      string `json:"task_guid"`
	Failed        bool   `json:"failed"`
	FailureReason string `json:"failure_reason"`
}

type TaskError struct {
	Id      TaskErrorID `json:"id"`
	Message string      `json:"message"`
}
<|endoftext|>


<commit_before>package test

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"sort"
	"strconv"
	"testing"

	"github.com/livegrep/livegrep/client"
	"gopkg.in/check.v1"
)

func Test(t *testing.T) { check.TestingT(t) }

var repo = flag.String("test.repo", "", "Git repository to run integration tests against.")

type IntegrationSuite struct {
	config *os.File
	index  *os.File
	client client.Client
}

var _ = check.Suite(&IntegrationSuite{})

type J map[string]interface{}

func (i *IntegrationSuite) SetUpSuite(c *check.C) {
	if *repo == "" {
		c.Skip("No test.root specified.")
		return
	}

	var err error

	i.config, err = ioutil.TempFile("", "livegrep")
	if err != nil {
		c.Fatalf("TempFile: %s", err)
	}

	i.index, err = ioutil.TempFile("", "livegrep")
	if err != nil {
		c.Fatalf("TempFile: %s", err)
	}
	enc := json.NewEncoder(i.config)
	config := J{
		"name": "livegrep",
		"repositories": []J{
			J{
				"name":      "test",
				"path":      *repo,
				"revisions": []string{"HEAD"},
			},
		},
	}
	if err = enc.Encode(config); err != nil {
		c.Fatalf("Encode: %s", err)
	}

	i.client, err = NewClient(
		"--dump_index", i.index.Name(),
		"--max_matches", "100000",
		"--timeout", "100000",
		i.config.Name())
	if err != nil {
		c.Fatalf("NewClient: %s", err)
	}
}

func (i *IntegrationSuite) TearDownSuite(*check.C) {
	if i.client != nil {
		i.client.Close()
	}
	if i.index != nil {
		os.Remove(i.index.Name())
		i.index.Close()
	}
	if i.config != nil {
		os.Remove(i.config.Name())
		i.config.Close()
	}
}

func (i *IntegrationSuite) TestBasic(c *check.C) {
	search, err := i.client.Query(&client.Query{Line: "."})
	if err != nil {
		c.Fatalf("Query: %s", err)
	}
	for _ = range search.Results() {
	}
	stats, err := search.Close()
	c.Check(stats, check.Not(check.IsNil))
	c.Check(err, check.IsNil)
}

type Match struct {
	Path string
	Line int
}

type SortMatches []Match

func (s SortMatches) Less(i, j int) bool {
	l, r := s[i], s[j]
	if l.Path != r.Path {
		return l.Path < r.Path
	}
	return l.Line < r.Line
}
func (s SortMatches) Len() int {
	return len(s)
}

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

var grepRE = regexp.MustCompile(`^HEAD:([^:]+):(\d+):`)

func gitGrep(path, regex string) ([]Match, error) {
	cmd := exec.Command("git", "grep", "-n", "-I", "-E", "-e", regex, "HEAD")
	cmd.Dir = path
	out, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	var matches []Match

	cmd.Start()
	scan := bufio.NewScanner(out)
	for scan.Scan() {
		line := scan.Bytes()
		m := grepRE.FindSubmatch(line)
		if m == nil {
			cmd.Process.Kill()
			return nil, fmt.Errorf("unparsable: `%s'", line)
		}
		lno, _ := strconv.Atoi(string(m[2]))
		matches = append(matches, Match{string(m[1]), lno})
	}
	if err := scan.Err(); err != nil {
		cmd.Process.Kill()
		return nil, err
	}
	cmd.Wait()

	return matches, nil
}

func (i *IntegrationSuite) crosscheck(c *check.C, regex string) {
	gitMatches, err := gitGrep(*repo, regex)
	if err != nil {
		c.Fatalf("git grep: %s", err)
	}

	var livegrepMatches []Match
	search, err := i.client.Query(&client.Query{Line: regex})
	if err != nil {
		c.Fatalf("Query: %s", err)
	}
	for m := range search.Results() {
		livegrepMatches = append(livegrepMatches, Match{m.Path, m.LineNumber})
	}
	search.Close()

	sort.Sort(SortMatches(gitMatches))
	sort.Sort(SortMatches(livegrepMatches))

	c.Check(livegrepMatches, check.DeepEquals, gitMatches)
}

func (i *IntegrationSuite) TestCrosscheck(c *check.C) {
	i.crosscheck(c, `hello`)
	i.crosscheck(c, "\t{10}")
	i.crosscheck(c, `^(\s.*\S)?printk\s*\(`)
}
<commit_msg>integration tests: handle long lines<commit_after>package test

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"sort"
	"strconv"
	"testing"

	"github.com/livegrep/livegrep/client"
	"gopkg.in/check.v1"
)

const LineLimit = 1024

func Test(t *testing.T) { check.TestingT(t) }

var repo = flag.String("test.repo", "", "Git repository to run integration tests against.")

type IntegrationSuite struct {
	config *os.File
	index  *os.File
	client client.Client
}

var _ = check.Suite(&IntegrationSuite{})

type J map[string]interface{}

func (i *IntegrationSuite) SetUpSuite(c *check.C) {
	if *repo == "" {
		c.Skip("No test.root specified.")
		return
	}

	var err error

	i.config, err = ioutil.TempFile("", "livegrep")
	if err != nil {
		c.Fatalf("TempFile: %s", err)
	}

	i.index, err = ioutil.TempFile("", "livegrep")
	if err != nil {
		c.Fatalf("TempFile: %s", err)
	}
	enc := json.NewEncoder(i.config)
	config := J{
		"name": "livegrep",
		"repositories": []J{
			J{
				"name":      "test",
				"path":      *repo,
				"revisions": []string{"HEAD"},
			},
		},
	}
	if err = enc.Encode(config); err != nil {
		c.Fatalf("Encode: %s", err)
	}

	i.client, err = NewClient(
		"--dump_index", i.index.Name(),
		"--max_matches", "100000",
		"--timeout", "100000",
		"--line_limit", strconv.Itoa(LineLimit),
		i.config.Name())
	if err != nil {
		c.Fatalf("NewClient: %s", err)
	}
}

func (i *IntegrationSuite) TearDownSuite(*check.C) {
	if i.client != nil {
		i.client.Close()
	}
	if i.index != nil {
		os.Remove(i.index.Name())
		i.index.Close()
	}
	if i.config != nil {
		os.Remove(i.config.Name())
		i.config.Close()
	}
}

func (i *IntegrationSuite) TestBasic(c *check.C) {
	search, err := i.client.Query(&client.Query{Line: "."})
	if err != nil {
		c.Fatalf("Query: %s", err)
	}
	for _ = range search.Results() {
	}
	stats, err := search.Close()
	c.Check(stats, check.Not(check.IsNil))
	c.Check(err, check.IsNil)
}

type Match struct {
	Path string
	Line int
}

type SortMatches []Match

func (s SortMatches) Less(i, j int) bool {
	l, r := s[i], s[j]
	if l.Path != r.Path {
		return l.Path < r.Path
	}
	return l.Line < r.Line
}
func (s SortMatches) Len() int {
	return len(s)
}

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

var grepRE = regexp.MustCompile(`^HEAD:([^:]+):(\d+):`)

func readLineDropTooLong(r *bufio.Reader) ([]byte, error) {
	for {
		line, isPrefix, err := r.ReadLine()
		if err != nil {
			return nil, err
		}
		if !isPrefix {
			line = bytes.TrimRight(line, "\n")
			return line, nil
		}
		for isPrefix {
			line, isPrefix, err = r.ReadLine()
			if err != nil {
				return nil, err
			}
		}
	}
}

func gitGrep(path, regex string) ([]Match, error) {
	cmd := exec.Command("git", "grep", "-n", "-I", "-E", "-e", regex, "HEAD")
	cmd.Dir = path
	out, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	var matches []Match

	cmd.Start()
	buf := bufio.NewReader(out)
	for {
		line, err := readLineDropTooLong(buf)
		if err != nil {
			if err == io.EOF {
				break
			}
			cmd.Process.Kill()
			return nil, err
		}
		m := grepRE.FindSubmatch(line)
		if m == nil {
			cmd.Process.Kill()
			return nil, fmt.Errorf("unparsable: `%s'", line)
		}
		if len(line)-len(m[0]) > LineLimit {
			continue
		}
		lno, _ := strconv.Atoi(string(m[2]))
		matches = append(matches, Match{string(m[1]), lno})
	}
	cmd.Wait()

	return matches, nil
}

func (i *IntegrationSuite) crosscheck(c *check.C, regex string) {
	gitMatches, err := gitGrep(*repo, regex)
	if err != nil {
		c.Fatalf("git grep: %s", err)
	}

	var livegrepMatches []Match
	search, err := i.client.Query(&client.Query{Line: regex})
	if err != nil {
		c.Fatalf("Query: %s", err)
	}
	for m := range search.Results() {
		livegrepMatches = append(livegrepMatches, Match{m.Path, m.LineNumber})
	}
	search.Close()

	sort.Sort(SortMatches(gitMatches))
	sort.Sort(SortMatches(livegrepMatches))

	c.Check(livegrepMatches, check.DeepEquals, gitMatches)
}

func (i *IntegrationSuite) TestCrosscheck(c *check.C) {
	i.crosscheck(c, `hello`)
	i.crosscheck(c, "\t{10}")
	i.crosscheck(c, `^(\s.*\S)?printk\s*\(`)
}
<|endoftext|>


<commit_before>package jpegsegs

import (
	"errors"
	"fmt"
	"io"
)

const (
	TEM  = 0x01
	SOF0 = 0xC0 // SOFn = SOF0+n, n = 0-15 excluding 4, 8 and 12
	DHT  = 0xC4
	JPG  = 0xC8
	DAC  = 0xCC
	RST0 = 0xD0 // RSTn = RST0+n, n = 0-7
	SOI  = 0xD8
	EOI  = 0xD9
	SOS  = 0xDA
	DQT  = 0xDB
	DNL  = 0xDC
	DRI  = 0xDD
	DHP  = 0xDE
	EXP  = 0xDF
	APP0 = 0xE0 // APPn = APP0+n, n = 0-15
	JPG0 = 0xF0 // JPGn = JPG0+n  n = 0-13
	COM  = 0xFE
)

// Marker represents a JPEG marker, which usually indicates the start of a
// segment.
type Marker uint8

var markerNames [256]string

// Initialize markerNames
func init() {
	markerNames[0] = "NUL"
	markerNames[TEM] = "TEM"
	markerNames[DHT] = "DHT"
	markerNames[JPG] = "JPG"
	markerNames[DAC] = "DAC"
	markerNames[SOI] = "SOI"
	markerNames[EOI] = "EOI"
	markerNames[SOS] = "SOS"
	markerNames[DQT] = "DQT"
	markerNames[DNL] = "DNL"
	markerNames[DRI] = "DRI"
	markerNames[DHP] = "DHP"
	markerNames[EXP] = "EXP"
	markerNames[COM] = "COM"
	markerNames[0xFF] = "FILL"

	var i Marker
	for i = 0x02; i <= 0xBF; i++ {
		markerNames[i] = fmt.Sprintf("RES%.2X", i) // Reserved
	}
	for i = SOF0; i <= SOF0+0xF; i++ {
		if i == SOF0+4 || i == SOF0+8 || i == SOF0+12 {
			continue
		}
		markerNames[i] = fmt.Sprintf("SOF%d", i-SOF0)
	}
	for i = RST0; i <= RST0+7; i++ {
		markerNames[i] = fmt.Sprintf("RST%d", i-RST0)
	}
	for i = APP0; i <= APP0+0xF; i++ {
		markerNames[i] = fmt.Sprintf("APP%d", i-APP0)
	}
	for i = JPG0; i <= JPG0+0xD; i++ {
		markerNames[i] = fmt.Sprintf("JPG%d", i-JPG0)
	}
}

// Name returns the name of a marker value.
func (m Marker) Name() string {
	return markerNames[m]
}

// Size of a JPEG file header.
const HeaderSize = 2

// Indicate if buffer contains a JPEG header.
func IsJPEGHeader(buf []byte) bool {
	return buf[0] == 0xFF && buf[1] == SOI
}

// The JPEG header is a SOI marker. Filler bytes aren't allowed.
func ReadHeader(reader io.Reader, buf []byte) error {
	buf = buf[0:2]
	if _, err := io.ReadFull(reader, buf); err != nil {
		return err
	}
	if !IsJPEGHeader(buf) {
		return errors.New("SOI marker not found")
	}
	return nil
}

func ReadMarker(reader io.Reader, buf []byte) (Marker, error) {
	buf = buf[0:2]
	if _, err := io.ReadFull(reader, buf); err != nil {
		return 0, err
	}
	if buf[0] != 0xFF {
		return 0, errors.New("0xFF expected in marker")
	}
	buf = buf[1:2] // Look at the 2nd byte only.
	for {
		// Skip 0xFF fill bytes. Fill bytes don't seem to have
		// any purpose, so can be discarded.
		if buf[0] != 0xFF {
			break
		}
		if _, err := reader.Read(buf); err != nil {
			return 0, err
		}
	}
	if buf[0] == 0 {
		return 0, errors.New("Invalid marker 0")
	}
	return Marker(buf[0]), nil
}

func WriteMarker(writer io.Writer, marker Marker, buf []byte) error {
	buf = buf[0:2]
	buf[0] = 0xFF
	buf[1] = byte(marker)
	_, err := writer.Write(buf)
	return err
}

func ReadData(reader io.Reader, buf []byte) ([]byte, error) {
	buf = buf[0:2]
	if _, err := io.ReadFull(reader, buf); err != nil {
		return nil, err
	}
	length := int(buf[0])<<8 + int(buf[1]) - 2
	buf = buf[0:length]
	_, err := io.ReadFull(reader, buf)
	return buf, err
}

func WriteData(writer io.Writer, buf []byte, lenbuf []byte) error {
	len := len(buf) + 2
	if len >= 2<<15 {
		return errors.New(fmt.Sprintf("writeData: data is too long (%d), max 2^16 - 3 (%d)", len-2, 2<<15-3))
	}
	lenbuf[0] = byte(len / 256)
	lenbuf[1] = byte(len % 256)
	if _, err := writer.Write(lenbuf); err != nil {
		return err
	}
	_, err := writer.Write(buf)
	return err
}

// Read image scan data up to the next marker. 'buf' is either a
// buffer to read into, which will be reallocated if required, or nil
// to allocate a new buffer. Returns a buffer with the image data and
// the following marker.
func ReadImageData(reader io.ByteReader, buf []byte) ([]byte, Marker, error) {
	pos := 0
	ff := false
	if buf == nil {
		buf = make([]byte, 10000)
	} else {
		buf = buf[:cap(buf)]
	}
	for {
		var err error
		if buf[pos], err = reader.ReadByte(); err != nil {
			return buf[:pos], 0, err
		}
		if ff {
			if buf[pos] == 0 {
				// Escaped 0xFF in data stream, delete
				// the 0 by not incrementing pos.
				ff = false
				continue
			}
			// Marker
			return buf[:pos-1], Marker(buf[pos]), nil
		} else if buf[pos] == 0xFF {
			ff = true
		}
		pos++
		if cap(buf) < pos+1 {
			newbuf := make([]byte, 2*cap(buf))
			copy(newbuf, buf)
			buf = newbuf
		}
	}

}

// Write a block of image data.
func WriteImageData(writer io.ByteWriter, buf []byte) error {
	for pos := range buf {
		if err := writer.WriteByte(buf[pos]); err != nil {
			return err
		}
		if buf[pos] == 0xFF {
			if err := writer.WriteByte(0); err != nil {
				return err
			}
		}
		pos++
	}
	return nil
}

// Scanner represents a reader for JPEG markers and segments up to the
// SOS marker.
type Scanner struct {
	reader io.Reader
	buf    []byte // buffer of size 2^16 - 3
}

// NewScanner creates a new Scanner and checks the JPEG header.
func NewScanner(reader io.Reader) (*Scanner, error) {
	scanner := new(Scanner)
	scanner.reader = reader
	scanner.buf = make([]byte, 2<<15-3)
	if err := ReadHeader(reader, scanner.buf); err != nil {
		return nil, err
	}
	return scanner, nil
}

// Scan reads the next JPEG marker and its data segment if it has one.
// All markers are expected to have data except for SOS, which indicates
// the start of scan data. Scan doesn't work past that point. The data
// buffer is only valid until Scan is called again.
func (scanner *Scanner) Scan() (Marker, []byte, error) {
	marker, err := ReadMarker(scanner.reader, scanner.buf)
	if err != nil {
		return 0, nil, err
	}
	if marker == SOS {
		return marker, nil, nil
	}
	segment, err := ReadData(scanner.reader, scanner.buf)
	return marker, segment, err

}

// Dumper represents a writer for JPEG markers and segments up to the SOS
// marker.
type Dumper struct {
	writer io.Writer
	buf    []byte // buffer of size 2
}

// NewDumper creates a new Dumper and writes the JPEG header.
func NewDumper(writer io.Writer) (*Dumper, error) {
	dumper := new(Dumper)
	dumper.writer = writer
	dumper.buf = make([]byte, 2)
	if err := WriteMarker(writer, SOI, dumper.buf); err != nil {
		return nil, err
	}
	return dumper, nil
}

// Dump writes a marker and its data segment from buf. buf should be nil if
// it's the SOS marker (start of scan).
func (dumper *Dumper) Dump(marker Marker, buf []byte) error {
	if err := WriteMarker(dumper.writer, marker, dumper.buf); err != nil {
		return err
	}
	if buf == nil {
		return nil
	}
	return WriteData(dumper.writer, buf, dumper.buf)
}

// Copy reads all remaining data from a Scanner and lets the Dumper write it.
func (dumper *Dumper) Copy(scanner *Scanner) error {
	_, err := io.Copy(dumper.writer, scanner.reader)
	return err
}

// Segment represents a marker and its segment data.
type Segment struct {
	Marker Marker
	Data   []byte
}

// ReadAll reads a JPEG stream up to and including the SOS marker and
// returns a slice with marker and segment data. The SOS marker isn't
// included in the slice.
func ReadAll(reader io.Reader) (*Scanner, []Segment, error) {
	var segments = make([]Segment, 0, 20)
	scanner, err := NewScanner(reader)
	if err != nil {
		return nil, segments, err
	}
	for {
		marker, buf, err := scanner.Scan()
		if err != nil {
			return scanner, segments, err
		}
		if marker == SOS {
			return scanner, segments, nil
		}
		cpy := make([]byte, len(buf))
		copy(cpy, buf)
		segments = append(segments, Segment{marker, cpy})
	}
}

// WriteAll writes a JPEG stream up to and including the SOS marker, given
// a slice with marker and segment data. The SOS marker is written
// automatically, it should not be included in the slice.
func WriteAll(writer io.Writer, segments []Segment) (*Dumper, error) {
	dumper, err := NewDumper(writer)
	if err != nil {
		return nil, err
	}
	for i := range segments {
		if err := dumper.Dump(segments[i].Marker, segments[i].Data); err != nil {
			return dumper, err
		}
	}
	return dumper, dumper.Dump(SOS, nil)
}
<commit_msg>define and export MPFHeaderSize, GetMPFHeader, PutMPFHeader<commit_after>package jpegsegs

import (
	"bytes"
	"errors"
	"fmt"
	"io"
)

const (
	TEM  = 0x01
	SOF0 = 0xC0 // SOFn = SOF0+n, n = 0-15 excluding 4, 8 and 12
	DHT  = 0xC4
	JPG  = 0xC8
	DAC  = 0xCC
	RST0 = 0xD0 // RSTn = RST0+n, n = 0-7
	SOI  = 0xD8
	EOI  = 0xD9
	SOS  = 0xDA
	DQT  = 0xDB
	DNL  = 0xDC
	DRI  = 0xDD
	DHP  = 0xDE
	EXP  = 0xDF
	APP0 = 0xE0 // APPn = APP0+n, n = 0-15
	JPG0 = 0xF0 // JPGn = JPG0+n  n = 0-13
	COM  = 0xFE
)

// Marker represents a JPEG marker, which usually indicates the start of a
// segment.
type Marker uint8

var markerNames [256]string

// Initialize markerNames
func init() {
	markerNames[0] = "NUL"
	markerNames[TEM] = "TEM"
	markerNames[DHT] = "DHT"
	markerNames[JPG] = "JPG"
	markerNames[DAC] = "DAC"
	markerNames[SOI] = "SOI"
	markerNames[EOI] = "EOI"
	markerNames[SOS] = "SOS"
	markerNames[DQT] = "DQT"
	markerNames[DNL] = "DNL"
	markerNames[DRI] = "DRI"
	markerNames[DHP] = "DHP"
	markerNames[EXP] = "EXP"
	markerNames[COM] = "COM"
	markerNames[0xFF] = "FILL"

	var i Marker
	for i = 0x02; i <= 0xBF; i++ {
		markerNames[i] = fmt.Sprintf("RES%.2X", i) // Reserved
	}
	for i = SOF0; i <= SOF0+0xF; i++ {
		if i == SOF0+4 || i == SOF0+8 || i == SOF0+12 {
			continue
		}
		markerNames[i] = fmt.Sprintf("SOF%d", i-SOF0)
	}
	for i = RST0; i <= RST0+7; i++ {
		markerNames[i] = fmt.Sprintf("RST%d", i-RST0)
	}
	for i = APP0; i <= APP0+0xF; i++ {
		markerNames[i] = fmt.Sprintf("APP%d", i-APP0)
	}
	for i = JPG0; i <= JPG0+0xD; i++ {
		markerNames[i] = fmt.Sprintf("JPG%d", i-JPG0)
	}
}

// Name returns the name of a marker value.
func (m Marker) Name() string {
	return markerNames[m]
}

// Size of a JPEG file header.
const HeaderSize = 2

// Indicate if buffer contains a JPEG header.
func IsJPEGHeader(buf []byte) bool {
	return buf[0] == 0xFF && buf[1] == SOI
}

// The JPEG header is a SOI marker. Filler bytes aren't allowed.
func ReadHeader(reader io.Reader, buf []byte) error {
	buf = buf[0:2]
	if _, err := io.ReadFull(reader, buf); err != nil {
		return err
	}
	if !IsJPEGHeader(buf) {
		return errors.New("SOI marker not found")
	}
	return nil
}

func ReadMarker(reader io.Reader, buf []byte) (Marker, error) {
	buf = buf[0:2]
	if _, err := io.ReadFull(reader, buf); err != nil {
		return 0, err
	}
	if buf[0] != 0xFF {
		return 0, errors.New("0xFF expected in marker")
	}
	buf = buf[1:2] // Look at the 2nd byte only.
	for {
		// Skip 0xFF fill bytes. Fill bytes don't seem to have
		// any purpose, so can be discarded.
		if buf[0] != 0xFF {
			break
		}
		if _, err := reader.Read(buf); err != nil {
			return 0, err
		}
	}
	if buf[0] == 0 {
		return 0, errors.New("Invalid marker 0")
	}
	return Marker(buf[0]), nil
}

func WriteMarker(writer io.Writer, marker Marker, buf []byte) error {
	buf = buf[0:2]
	buf[0] = 0xFF
	buf[1] = byte(marker)
	_, err := writer.Write(buf)
	return err
}

func ReadData(reader io.Reader, buf []byte) ([]byte, error) {
	buf = buf[0:2]
	if _, err := io.ReadFull(reader, buf); err != nil {
		return nil, err
	}
	length := int(buf[0])<<8 + int(buf[1]) - 2
	buf = buf[0:length]
	_, err := io.ReadFull(reader, buf)
	return buf, err
}

func WriteData(writer io.Writer, buf []byte, lenbuf []byte) error {
	len := len(buf) + 2
	if len >= 2<<15 {
		return errors.New(fmt.Sprintf("writeData: data is too long (%d), max 2^16 - 3 (%d)", len-2, 2<<15-3))
	}
	lenbuf[0] = byte(len / 256)
	lenbuf[1] = byte(len % 256)
	if _, err := writer.Write(lenbuf); err != nil {
		return err
	}
	_, err := writer.Write(buf)
	return err
}

// Read image scan data up to the next marker. 'buf' is either a
// buffer to read into, which will be reallocated if required, or nil
// to allocate a new buffer. Returns a buffer with the image data and
// the following marker.
func ReadImageData(reader io.ByteReader, buf []byte) ([]byte, Marker, error) {
	pos := 0
	ff := false
	if buf == nil {
		buf = make([]byte, 10000)
	} else {
		buf = buf[:cap(buf)]
	}
	for {
		var err error
		if buf[pos], err = reader.ReadByte(); err != nil {
			return buf[:pos], 0, err
		}
		if ff {
			if buf[pos] == 0 {
				// Escaped 0xFF in data stream, delete
				// the 0 by not incrementing pos.
				ff = false
				continue
			}
			// Marker
			return buf[:pos-1], Marker(buf[pos]), nil
		} else if buf[pos] == 0xFF {
			ff = true
		}
		pos++
		if cap(buf) < pos+1 {
			newbuf := make([]byte, 2*cap(buf))
			copy(newbuf, buf)
			buf = newbuf
		}
	}

}

// Write a block of image data.
func WriteImageData(writer io.ByteWriter, buf []byte) error {
	for pos := range buf {
		if err := writer.WriteByte(buf[pos]); err != nil {
			return err
		}
		if buf[pos] == 0xFF {
			if err := writer.WriteByte(0); err != nil {
				return err
			}
		}
		pos++
	}
	return nil
}

// Scanner represents a reader for JPEG markers and segments up to the
// SOS marker.
type Scanner struct {
	reader io.Reader
	buf    []byte // buffer of size 2^16 - 3
}

// NewScanner creates a new Scanner and checks the JPEG header.
func NewScanner(reader io.Reader) (*Scanner, error) {
	scanner := new(Scanner)
	scanner.reader = reader
	scanner.buf = make([]byte, 2<<15-3)
	if err := ReadHeader(reader, scanner.buf); err != nil {
		return nil, err
	}
	return scanner, nil
}

// Scan reads the next JPEG marker and its data segment if it has one.
// All markers are expected to have data except for SOS, which indicates
// the start of scan data. Scan doesn't work past that point. The data
// buffer is only valid until Scan is called again.
func (scanner *Scanner) Scan() (Marker, []byte, error) {
	marker, err := ReadMarker(scanner.reader, scanner.buf)
	if err != nil {
		return 0, nil, err
	}
	if marker == SOS {
		return marker, nil, nil
	}
	segment, err := ReadData(scanner.reader, scanner.buf)
	return marker, segment, err

}

// Dumper represents a writer for JPEG markers and segments up to the SOS
// marker.
type Dumper struct {
	writer io.Writer
	buf    []byte // buffer of size 2
}

// NewDumper creates a new Dumper and writes the JPEG header.
func NewDumper(writer io.Writer) (*Dumper, error) {
	dumper := new(Dumper)
	dumper.writer = writer
	dumper.buf = make([]byte, 2)
	if err := WriteMarker(writer, SOI, dumper.buf); err != nil {
		return nil, err
	}
	return dumper, nil
}

// Dump writes a marker and its data segment from buf. buf should be nil if
// it's the SOS marker (start of scan).
func (dumper *Dumper) Dump(marker Marker, buf []byte) error {
	if err := WriteMarker(dumper.writer, marker, dumper.buf); err != nil {
		return err
	}
	if buf == nil {
		return nil
	}
	return WriteData(dumper.writer, buf, dumper.buf)
}

// Copy reads all remaining data from a Scanner and lets the Dumper write it.
func (dumper *Dumper) Copy(scanner *Scanner) error {
	_, err := io.Copy(dumper.writer, scanner.reader)
	return err
}

// Segment represents a marker and its segment data.
type Segment struct {
	Marker Marker
	Data   []byte
}

// ReadAll reads a JPEG stream up to and including the SOS marker and
// returns a slice with marker and segment data. The SOS marker isn't
// included in the slice.
func ReadAll(reader io.Reader) (*Scanner, []Segment, error) {
	var segments = make([]Segment, 0, 20)
	scanner, err := NewScanner(reader)
	if err != nil {
		return nil, segments, err
	}
	for {
		marker, buf, err := scanner.Scan()
		if err != nil {
			return scanner, segments, err
		}
		if marker == SOS {
			return scanner, segments, nil
		}
		cpy := make([]byte, len(buf))
		copy(cpy, buf)
		segments = append(segments, Segment{marker, cpy})
	}
}

// WriteAll writes a JPEG stream up to and including the SOS marker, given
// a slice with marker and segment data. The SOS marker is written
// automatically, it should not be included in the slice.
func WriteAll(writer io.Writer, segments []Segment) (*Dumper, error) {
	dumper, err := NewDumper(writer)
	if err != nil {
		return nil, err
	}
	for i := range segments {
		if err := dumper.Dump(segments[i].Marker, segments[i].Data); err != nil {
			return dumper, err
		}
	}
	return dumper, dumper.Dump(SOS, nil)
}

// MPF header, as found i a JPEG APP2 segment.
var mpfheader = []byte("MPF\000")

// Size of a MPF header.
const MPFHeaderSize = 4

// Check if a slice starts with a Multi-Picture Format (MPF) header,
// as found in a JPEG APP2 segment.  Returns a flag and the position
// of the next byte.
func GetMPFHeader(buf []byte) (bool, uint32) {
	if uint32(len(buf)) >= MPFHeaderSize && bytes.Compare(buf[:MPFHeaderSize], mpfheader) == 0 {
		return true, MPFHeaderSize
	} else {
		return false, 0
	}
}

// Put a MPF header, as for a JPEG APP2 segment, at the start of a slice,
// returning the position of the next byte.
func PutMPFHeader(buf []byte) uint32 {
	copy(buf, mpfheader)
	return MPFHeaderSize
}
<|endoftext|>


<commit_before>package test

import (
	"encoding/json"
	"flag"
	"io/ioutil"
	"os"
	"testing"

	"github.com/livegrep/livegrep/client"
	"gopkg.in/check.v1"
)

func Test(t *testing.T) { check.TestingT(t) }

var repo = flag.String("test.repo", "", "Git repository to run integration tests against.")

type IntegrationSuite struct {
	config *os.File
	index  *os.File
	client client.Client
}

var _ = check.Suite(&IntegrationSuite{})

type J map[string]interface{}

func (i *IntegrationSuite) SetUpSuite(c *check.C) {
	if *repo == "" {
		c.Skip("No test.root specified.")
		return
	}

	var err error

	i.config, err = ioutil.TempFile("", "livegrep")
	if err != nil {
		c.Fatalf("TempFile: %s", err)
	}

	i.index, err = ioutil.TempFile("", "livegrep")
	if err != nil {
		c.Fatalf("TempFile: %s", err)
	}
	enc := json.NewEncoder(i.config)
	config := J{
		"name": "livegrep",
		"repositories": []J{
			J{
				"name":      "test",
				"path":      *repo,
				"revisions": []string{"HEAD"},
			},
		},
	}
	if err = enc.Encode(config); err != nil {
		c.Fatalf("Encode: %s", err)
	}

	i.client, err = NewClient("--dump_index", i.index.Name(), i.config.Name())
	if err != nil {
		c.Fatalf("NewClient: %s", err)
	}
}

func (i *IntegrationSuite) TearDownSuite(*check.C) {
	if i.client != nil {
		i.client.Close()
	}
	if i.index != nil {
		os.Remove(i.index.Name())
		i.index.Close()
	}
	if i.config != nil {
		os.Remove(i.config.Name())
		i.config.Close()
	}
}

func (i *IntegrationSuite) TestBasic(c *check.C) {
	search, err := i.client.Query(&client.Query{Line: "."})
	if err != nil {
		c.Fatalf("Query: %s", err)
	}
	for _ = range search.Results() {
	}
	stats, err := search.Close()
	c.Check(stats, check.Not(check.IsNil))
	c.Check(err, check.IsNil)
}
<commit_msg>basic crosscheck against `git grep`<commit_after>package test

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"sort"
	"strconv"
	"testing"

	"github.com/livegrep/livegrep/client"
	"gopkg.in/check.v1"
)

func Test(t *testing.T) { check.TestingT(t) }

var repo = flag.String("test.repo", "", "Git repository to run integration tests against.")

type IntegrationSuite struct {
	config *os.File
	index  *os.File
	client client.Client
}

var _ = check.Suite(&IntegrationSuite{})

type J map[string]interface{}

func (i *IntegrationSuite) SetUpSuite(c *check.C) {
	if *repo == "" {
		c.Skip("No test.root specified.")
		return
	}

	var err error

	i.config, err = ioutil.TempFile("", "livegrep")
	if err != nil {
		c.Fatalf("TempFile: %s", err)
	}

	i.index, err = ioutil.TempFile("", "livegrep")
	if err != nil {
		c.Fatalf("TempFile: %s", err)
	}
	enc := json.NewEncoder(i.config)
	config := J{
		"name": "livegrep",
		"repositories": []J{
			J{
				"name":      "test",
				"path":      *repo,
				"revisions": []string{"HEAD"},
			},
		},
	}
	if err = enc.Encode(config); err != nil {
		c.Fatalf("Encode: %s", err)
	}

	i.client, err = NewClient(
		"--dump_index", i.index.Name(),
		"--max_matches", "100000",
		"--timeout", "100000",
		i.config.Name())
	if err != nil {
		c.Fatalf("NewClient: %s", err)
	}
}

func (i *IntegrationSuite) TearDownSuite(*check.C) {
	if i.client != nil {
		i.client.Close()
	}
	if i.index != nil {
		os.Remove(i.index.Name())
		i.index.Close()
	}
	if i.config != nil {
		os.Remove(i.config.Name())
		i.config.Close()
	}
}

func (i *IntegrationSuite) TestBasic(c *check.C) {
	search, err := i.client.Query(&client.Query{Line: "."})
	if err != nil {
		c.Fatalf("Query: %s", err)
	}
	for _ = range search.Results() {
	}
	stats, err := search.Close()
	c.Check(stats, check.Not(check.IsNil))
	c.Check(err, check.IsNil)
}

type Match struct {
	Path string
	Line int
}

type SortMatches []Match

func (s SortMatches) Less(i, j int) bool {
	l, r := s[i], s[j]
	if l.Path != r.Path {
		return l.Path < r.Path
	}
	return l.Line < r.Line
}
func (s SortMatches) Len() int {
	return len(s)
}

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

var grepRE = regexp.MustCompile(`^HEAD:([^:]+):(\d+):`)

func gitGrep(path, regex string) ([]Match, error) {
	cmd := exec.Command("git", "grep", "-n", "-E", "-e", regex, "HEAD")
	cmd.Dir = path
	out, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	var matches []Match

	cmd.Start()
	scan := bufio.NewScanner(out)
	for scan.Scan() {
		line := scan.Bytes()
		m := grepRE.FindSubmatch(line)
		if m == nil {
			return nil, fmt.Errorf("unparsable: `%s'", line)
		}
		lno, _ := strconv.Atoi(string(m[2]))
		matches = append(matches, Match{string(m[1]), lno})
	}
	cmd.Wait()

	return matches, nil
}

func (i *IntegrationSuite) crosscheck(c *check.C, regex string) {
	gitMatches, err := gitGrep(*repo, regex)
	if err != nil {
		c.Fatalf("git grep: %s", err)
	}

	var livegrepMatches []Match
	search, err := i.client.Query(&client.Query{Line: regex})
	if err != nil {
		c.Fatalf("Query: %s", err)
	}
	for m := range search.Results() {
		livegrepMatches = append(livegrepMatches, Match{m.Path, m.LineNumber})
	}
	search.Close()

	sort.Sort(SortMatches(gitMatches))
	sort.Sort(SortMatches(livegrepMatches))

	c.Check(livegrepMatches, check.DeepEquals, gitMatches)
}

func (i *IntegrationSuite) TestCrosscheck(c *check.C) {
	i.crosscheck(c, `hello`)
}
<|endoftext|>


<commit_before>package chatblast

/*
 * It really knows how to control a room...
 */

import (
	"log"
	"sync"
)

type RoomController struct {
	// Map of room Ids to rooms
	Rooms     map[string]*Room
	RoomsLock sync.RWMutex
	// Map of user Ids to users
	Users     map[string]*User
	UsersLock sync.RWMutex
	Channel   chan *Message
}

func MakeRoomController() *RoomController {
	rc := &RoomController{
		Rooms:     make(map[string]*Room),
		RoomsLock: sync.RWMutex{},
		Users:     make(map[string]*User),
		UsersLock: sync.RWMutex{},
		Channel:   make(chan *Message),
	}
	log.Println("New RoomController created")
	// Add global (ownerless) room
	globalRoom := &Room{
		Name:        "Global",
		Id:          "global",
		Subscribers: map[string]*User{},
		Channel:     make(chan *Message),
	}
	log.Println("New room Global created")
	go globalRoom.Listen()

	rc.SetRoom("global", globalRoom)
	go rc.Listen()
	return rc
}

func (rc *RoomController) GetRoom(roomId string) (*Room, bool) {
	rc.RoomsLock.RLock()
	defer rc.RoomsLock.RUnlock()
	room, ok := rc.Rooms[roomId]
	return room, ok
}

func (rc *RoomController) SetRoom(roomId string, room *Room) {
	rc.RoomsLock.Lock()
	defer rc.RoomsLock.Unlock()
	rc.Rooms[roomId] = room
}

func (rc *RoomController) DeleteRoom(roomId string) {
	rc.RoomsLock.Lock()
	defer rc.RoomsLock.Unlock()
	delete(rc.Rooms, roomId)
}

func (rc *RoomController) GetUser(userId string) (*User, bool) {
	rc.UsersLock.RLock()
	defer rc.UsersLock.RUnlock()
	user, ok := rc.Users[userId]
	return user, ok
}

func (rc *RoomController) SetUser(userId string, u *User) {
	rc.UsersLock.Lock()
	defer rc.UsersLock.Unlock()
	rc.Users[userId] = u
}

func (rc *RoomController) DeleteUser(userId string) {
	rc.UsersLock.Lock()
	defer rc.UsersLock.Unlock()
	delete(rc.Users, userId)
}

func (rc *RoomController) SendGlobal(msg *Message) {
	if global, ok := rc.GetRoom("global"); ok {
		global.Broadcast(msg)
	}
}

func (rc *RoomController) AddRoom(name string, u *User) {
	// TODO Only allow uniquely named rooms?
	newRoom := NewRoom(name, u)
	rc.SetRoom(newRoom.Id, newRoom)
	rc.Subscribe(u, newRoom.Id)
	msg := &Message{
		RoomId: "global",
		Cmd:    "newrm",
		Msg: []messageData{
			{Name: "rid", Value: newRoom.Id},
			{Name: "name", Value: name},
		},
	}
	rc.Dispatch(msg)
	log.Println("New room", name, "created")
}

func (rc *RoomController) RemoveRoom(rid string, u *User) {
	if room, ok := rc.GetRoom(rid); ok {
		closed, err := room.CloseRoom(u)
		if err != nil {
			log.Println("Uh oh!")
		} else if closed {
			rc.DeleteRoom(rid)
			log.Println("Room", room.Name, "deleted")
		} else {
			msg := &Message{
				Cmd:    "err",
				Text:   "Room could not be closed.",
				RoomId: rid,
			}
			u.Tell(msg)
		}
	}
}

func (rc *RoomController) AddUser(u *User) {
	log.Println("New user", u.Name, "added to RoomController")
	rc.SetUser(u.Id, u)
	msg := &Message{
		Cmd:    "login",
		User: u,
		RoomId: "global",
	}
	rc.SendGlobal(msg)
	rc.Subscribe(u, "global")
}

func (rc *RoomController) RemoveUser(u *User) {
	rc.DeleteUser(u.Id)
	for _, room := range rc.Rooms {
		if room.Owner == u {
			rc.RemoveRoom(room.Id, u)
		} else if _, ok := room.GetSubscriber(u.Id); ok {
			rc.Unsubscribe(u, room.Id)
		}
	}
	msg := &Message{
		Cmd:    "logoff",
		UserId: u.Id,
		RoomId: "global",
	}
	rc.SendGlobal(msg)
	log.Println("User", u.Name, "removed from RoomController")
}

func (rc *RoomController) Subscribe(u *User, roomId string) {
	if room, ok := rc.GetRoom(roomId); ok {
		room.Join(u)
		log.Println("User", u.Name, "subscribed to", room.Name)
	} else {
		msg := &Message{
			Cmd:  "err",
			Text: "Room does not exist",
		}
		u.Tell(msg)
	}
}

func (rc *RoomController) Unsubscribe(u *User, roomId string) {
	if room, ok := rc.GetRoom(roomId); ok {
		room.Leave(u)
		log.Println("User", u.Name, "unsubscribed from", room.Name)
	} else {
		msg := &Message{
			Cmd:  "err",
			Text: "Room does not seem to exist",
		}
		u.Tell(msg)
	}
}

func (rc *RoomController) Broadcast(msg *Message) {
	// TODO Message dispatch
	rc.Channel <- msg
}

//func (rc *RoomController) GetUsers(roomId string) []*User {
//	var userArray []*User
//	if room, ok := rc.GetRoom(roomId); ok {
//		for _, user := range room.Subscribers {
//			userArray = append(userArray, user)
//		}
//	}
//	return userArray
//}

//func (rc *RoomController) GetRooms(roomId string) []*Room {
//	var userArray []*User
//	if room, ok := rc.GetRoom(roomId); ok {
//		for _, user := range room.Subscribers {
//			userArray = append(userArray, user)
//		}
//	}
//	return userArray
//}

func (rc *RoomController) Dispatch(incoming *Message) {
	switch incoming.Cmd {
	case "sub":
		if incoming.RoomId != "" {
			rc.Subscribe(incoming.User, incoming.RoomId)
		} else {
			log.Println("Bad sub request")
		}
		// TODO error handling
	case "unsub":
		if incoming.RoomId != "" {
			rc.Unsubscribe(incoming.User, incoming.RoomId)
		} else {
			log.Println("Bad unsub request")
		}
		// TODO error handling
	case "addrm":
		if incoming.Text != "" {
			rc.AddRoom(incoming.Text, incoming.User)
		} else {
			log.Println("Bad addrm request")
		}
		// TODO error handling
	case "remvrm":
		if incoming.RoomId != "" {
			rc.RemoveRoom(incoming.RoomId, incoming.User)
		} else {
			log.Println("Bad remvrm request")
		}
		// TODO error handling
		//	case "wspr":
		//		TODO add whispers
	default:
		if room, ok := rc.GetRoom(incoming.RoomId); ok {
			// Remove user info to reduce payload size
			incoming.User = nil
			room.Broadcast(incoming)
		}
	}
}

func (rc *RoomController) Listen() {
	for incoming := range rc.Channel {
		rc.Dispatch(incoming)
	}
}
<commit_msg>Revert<commit_after>package chatblast

/*
 * It really knows how to control a room...
 */

import (
	"log"
	"sync"
)

type RoomController struct {
	// Map of room Ids to rooms
	Rooms     map[string]*Room
	RoomsLock sync.RWMutex
	// Map of user Ids to users
	Users     map[string]*User
	UsersLock sync.RWMutex
	Channel   chan *Message
}

func MakeRoomController() *RoomController {
	rc := &RoomController{
		Rooms:     make(map[string]*Room),
		RoomsLock: sync.RWMutex{},
		Users:     make(map[string]*User),
		UsersLock: sync.RWMutex{},
		Channel:   make(chan *Message),
	}
	log.Println("New RoomController created")
	// Add global (ownerless) room
	globalRoom := &Room{
		Name:        "Global",
		Id:          "global",
		Subscribers: map[string]*User{},
		Channel:     make(chan *Message),
	}
	log.Println("New room Global created")
	go globalRoom.Listen()

	rc.SetRoom("global", globalRoom)
	go rc.Listen()
	return rc
}

func (rc *RoomController) GetRoom(roomId string) (*Room, bool) {
	rc.RoomsLock.RLock()
	defer rc.RoomsLock.RUnlock()
	room, ok := rc.Rooms[roomId]
	return room, ok
}

func (rc *RoomController) SetRoom(roomId string, room *Room) {
	rc.RoomsLock.Lock()
	defer rc.RoomsLock.Unlock()
	rc.Rooms[roomId] = room
}

func (rc *RoomController) DeleteRoom(roomId string) {
	rc.RoomsLock.Lock()
	defer rc.RoomsLock.Unlock()
	delete(rc.Rooms, roomId)
}

func (rc *RoomController) GetUser(userId string) (*User, bool) {
	rc.UsersLock.RLock()
	defer rc.UsersLock.RUnlock()
	user, ok := rc.Users[userId]
	return user, ok
}

func (rc *RoomController) SetUser(userId string, u *User) {
	rc.UsersLock.Lock()
	defer rc.UsersLock.Unlock()
	rc.Users[userId] = u
}

func (rc *RoomController) DeleteUser(userId string) {
	rc.UsersLock.Lock()
	defer rc.UsersLock.Unlock()
	delete(rc.Users, userId)
}

func (rc *RoomController) SendGlobal(msg *Message) {
	if global, ok := rc.GetRoom("global"); ok {
		global.Broadcast(msg)
	}
}

func (rc *RoomController) AddRoom(name string, u *User) {
	// TODO Only allow uniquely named rooms?
	newRoom := NewRoom(name, u)
	rc.SetRoom(newRoom.Id, newRoom)
	rc.Subscribe(u, newRoom.Id)
	msg := &Message{
		RoomId: "global",
		Cmd:    "newrm",
		Msg: []messageData{
			{Name: "rid", Value: newRoom.Id},
			{Name: "name", Value: name},
		},
	}
	rc.Dispatch(msg)
	log.Println("New room", name, "created")
}

func (rc *RoomController) RemoveRoom(rid string, u *User) {
	if room, ok := rc.GetRoom(rid); ok {
		closed, err := room.CloseRoom(u)
		if err != nil {
			log.Println("Uh oh!")
		} else if closed {
			rc.DeleteRoom(rid)
			log.Println("Room", room.Name, "deleted")
		} else {
			msg := &Message{
				Cmd:    "err",
				Text:   "Room could not be closed.",
				RoomId: rid,
			}
			u.Tell(msg)
		}
	}
}

func (rc *RoomController) AddUser(u *User) {
	log.Println("New user", u.Name, "added to RoomController")
	rc.SetUser(u.Id, u)
	msg := &Message{
		Cmd:    "login",
		User: u,
		RoomId: "global",
	}
	rc.SendGlobal(msg)
	rc.Subscribe(u, "global")
}

func (rc *RoomController) RemoveUser(u *User) {
	rc.DeleteUser(u.Id)
	for _, room := range rc.Rooms {
		if room.Owner == u {
			rc.RemoveRoom(room.Id, u)
		} else if _, ok := room.GetSubscriber(u.Id); ok {
			rc.Unsubscribe(u, room.Id)
		}
	}
	msg := &Message{
		Cmd:    "logoff",
		UserId: u.Id,
		RoomId: "global",
	}
	rc.SendGlobal(msg)
	log.Println("User", u.Name, "removed from RoomController")
}

func (rc *RoomController) Subscribe(u *User, roomId string) {
	if room, ok := rc.GetRoom(roomId); ok {
		room.Join(u)
		log.Println("User", u.Name, "subscribed to", room.Name)
	} else {
		msg := &Message{
			Cmd:  "err",
			Text: "Room does not exist",
		}
		u.Tell(msg)
	}
}

func (rc *RoomController) Unsubscribe(u *User, roomId string) {
	if room, ok := rc.GetRoom(roomId); ok {
		room.Leave(u)
		log.Println("User", u.Name, "unsubscribed from", room.Name)
	} else {
		msg := &Message{
			Cmd:  "err",
			Text: "Room does not seem to exist",
		}
		u.Tell(msg)
	}
}

func (rc *RoomController) Broadcast(msg *Message) {
	// TODO Message dispatch
	rc.Channel <- msg
}

//func (rc *RoomController) GetUsers(roomId string) []*User {
//	var userArray []*User
//	if room, ok := rc.GetRoom(roomId); ok {
//		for _, user := range room.Subscribers {
//			userArray = append(userArray, user)
//		}
//	}
//	return userArray
//}

//func (rc *RoomController) GetRooms(roomId string) []*Room {
//	var userArray []*User
//	if room, ok := rc.GetRoom(roomId); ok {
//		for _, user := range room.Subscribers {
//			userArray = append(userArray, user)
//		}
//	}
//	return userArray
//}

func (rc *RoomController) Dispatch(incoming *Message) {
	switch incoming.Cmd {
	case "sub":
		if incoming.RoomId != "" {
			rc.Subscribe(incoming.User, incoming.RoomId)
		} else {
			log.Println("Bad sub request")
		}
		// TODO error handling
	case "unsub":
		if incoming.RoomId != "" {
			rc.Unsubscribe(incoming.User, incoming.RoomId)
		} else {
			log.Println("Bad unsub request")
		}
		// TODO error handling
	case "addrm":
		if incoming.Text != "" {
			rc.AddRoom(incoming.Text, incoming.User)
		} else {
			log.Println("Bad addrm request")
		}
		// TODO error handling
	case "remvrm":
		if incoming.RoomId != "" {
			rc.RemoveRoom(incoming.RoomId, incoming.User)
		} else {
			log.Println("Bad remvrm request")
		}
		// TODO error handling
		//	case "wspr":
		//		TODO add whispers
	default:
		if room, ok := rc.GetRoom(incoming.RoomId); ok {
			room.Broadcast(incoming)
		}
	}
}

func (rc *RoomController) Listen() {
	for incoming := range rc.Channel {
		rc.Dispatch(incoming)
	}
}
<|endoftext|>


<commit_before>package main_test

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"

	"github.com/concourse/concourse/atc/postgresrunner"
	"github.com/onsi/gomega/gbytes"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
	"golang.org/x/crypto/ssh"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Web Command", func() {

	var (
		hostKeyFile    string
		hostPubKeyFile string

		concourseCommand *exec.Cmd
		concourseProcess ifrit.Process
		concourseRunner  *ginkgomon.Runner
		postgresRunner   postgresrunner.Runner
		dbProcess        ifrit.Process
	)

	BeforeEach(func() {
		hostKeyFile, hostPubKeyFile, _, _ = generateSSHKeypair()
		postgresRunner = postgresrunner.Runner{
			Port: 5433 + GinkgoParallelNode(),
		}
		dbProcess = ifrit.Invoke(postgresRunner)
		postgresRunner.CreateTestDB()

		concourseCommand = exec.Command(
			concoursePath,
			"web",
			"--tsa-host-key", hostKeyFile,
			"--postgres-user", "postgres",
			"--postgres-database", "testdb",
			"--postgres-port", strconv.Itoa(5433+GinkgoParallelNode()),
			"--main-team-local-user", "test",
			"--add-local-user", "test:test",
			"--debug-bind-port", strconv.Itoa(8000+GinkgoParallelNode()),
			"--bind-port", strconv.Itoa(8080+GinkgoParallelNode()),
			"--tsa-bind-port", strconv.Itoa(2222+GinkgoParallelNode()),
			"--client-id", "client-id",
			"--client-secret", "client-secret",
			"--tsa-client-id", "tsa-client-id",
			"--tsa-client-secret", "tsa-client-secret",
			"--tsa-token-url", "http://localhost/token",
		)
	})

	JustBeforeEach(func() {
		concourseRunner = ginkgomon.New(ginkgomon.Config{
			Command:       concourseCommand,
			Name:          "web",
			AnsiColorCode: "32m",
		})

		concourseProcess = ifrit.Background(concourseRunner)

		// workaround to avoid panic due to registering http handlers multiple times
		http.DefaultServeMux = new(http.ServeMux)
	})

	AfterEach(func() {
		ginkgomon.Interrupt(concourseProcess)
		<-concourseProcess.Wait()
		postgresRunner.DropTestDB()

		dbProcess.Signal(os.Interrupt)
		err := <-dbProcess.Wait()
		Expect(err).NotTo(HaveOccurred())
		os.Remove(hostKeyFile)
		os.Remove(hostPubKeyFile)
		os.Remove(filepath.Dir(hostPubKeyFile))
	})

	It("starts atc", func() {
		Eventually(concourseRunner.Buffer(), "30s", "2s").Should(gbytes.Say("atc.listening"))
	})

	It("starts tsa", func() {
		Eventually(concourseRunner.Buffer(), "30s", "2s").Should(gbytes.Say("tsa.listening"))
	})

	Context("when CONCOURSE_CONCURRENT_REQUEST_LIMIT is invalid", func() {
		BeforeEach(func() {
			concourseCommand.Env = append(concourseCommand.Env, "CONCOURSE_CONCURRENT_REQUEST_LIMIT=InvalidAction:0")
		})

		It("prints an error and exits", func() {
			Eventually(concourseRunner.Err()).Should(gbytes.Say("'InvalidAction' is not a valid action"))
		})
	})

	Context("with CONCOURSE_TSA_CLIENT_ID specified", func() {
		BeforeEach(func() {
			concourseCommand.Env = append(concourseCommand.Env, "CONCOURSE_TSA_CLIENT_ID=tsa-client-id")
		})

		It("starts atc", func() {
			Eventually(concourseRunner.Buffer(), "30s", "2s").Should(gbytes.Say("atc.listening"))
		})

		It("starts tsa", func() {
			Eventually(concourseRunner.Buffer(), "30s", "2s").Should(gbytes.Say("tsa.listening"))
		})

		Context("with CONCOURSE_SYSTEM_CLAIM_KEY is not set to 'aud'", func() {
			BeforeEach(func() {
				concourseCommand.Env = append(concourseCommand.Env, "CONCOURSE_SYSTEM_CLAIM_KEY=not-aud")
			})

			It("starts atc", func() {
				Eventually(concourseRunner.Buffer(), "30s", "2s").Should(gbytes.Say("atc.listening"))
			})

			It("starts tsa", func() {
				Eventually(concourseRunner.Buffer(), "30s", "2s").Should(gbytes.Say("tsa.listening"))
			})
		})

		Context("with CONCOURSE_SYSTEM_CLAIM_KEY set to 'aud'", func() {
			BeforeEach(func() {
				concourseCommand.Env = append(concourseCommand.Env, "CONCOURSE_SYSTEM_CLAIM_KEY=aud")
			})

			Context("when the system claim values does not contain the client id", func() {
				BeforeEach(func() {
					concourseCommand.Env = append(concourseCommand.Env, "CONCOURSE_SYSTEM_CLAIM_VALUE=system-claim-value-1,system-claim-value-2")
				})

				It("errors", func() {
					Eventually(concourseRunner.Err()).Should(
						gbytes.Say("at least one systemClaimValue must be equal to tsa-client-id"),
					)
				})
			})

			Context("when the system claim values contain the client id", func() {
				BeforeEach(func() {
					concourseCommand.Env = append(concourseCommand.Env, "CONCOURSE_SYSTEM_CLAIM_VALUE=system-claim-value-1,tsa-client-id")
				})

				It("starts atc", func() {
					Eventually(concourseRunner.Buffer(), "30s", "2s").Should(gbytes.Say("atc.listening"))
				})

				It("starts tsa", func() {
					Eventually(concourseRunner.Buffer(), "30s", "2s").Should(gbytes.Say("tsa.listening"))
				})
			})
		})
	})
})

func generateSSHKeypair() (string, string, *rsa.PrivateKey, ssh.PublicKey) {
	path, err := ioutil.TempDir("", "tsa-key")
	Expect(err).NotTo(HaveOccurred())

	privateKeyPath := filepath.Join(path, "id_rsa")
	publicKeyPath := privateKeyPath + ".pub"

	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	Expect(err).NotTo(HaveOccurred())

	privateKeyBytes := pem.EncodeToMemory(&pem.Block{
		Type:    "RSA PRIVATE KEY",
		Headers: nil,
		Bytes:   x509.MarshalPKCS1PrivateKey(privateKey),
	})

	publicKeyRsa, err := ssh.NewPublicKey(&privateKey.PublicKey)
	Expect(err).NotTo(HaveOccurred())

	publicKeyBytes := ssh.MarshalAuthorizedKey(publicKeyRsa)

	err = ioutil.WriteFile(privateKeyPath, privateKeyBytes, 0600)
	Expect(err).NotTo(HaveOccurred())

	err = ioutil.WriteFile(publicKeyPath, publicKeyBytes, 0600)
	Expect(err).NotTo(HaveOccurred())

	return privateKeyPath, publicKeyPath, privateKey, publicKeyRsa
}
<commit_msg>increase timeout for system claim startup error<commit_after>package main_test

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"time"

	"github.com/concourse/concourse/atc/postgresrunner"
	"github.com/onsi/gomega/gbytes"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
	"golang.org/x/crypto/ssh"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Web Command", func() {

	var (
		hostKeyFile    string
		hostPubKeyFile string

		concourseCommand *exec.Cmd
		concourseProcess ifrit.Process
		concourseRunner  *ginkgomon.Runner
		postgresRunner   postgresrunner.Runner
		dbProcess        ifrit.Process
	)

	BeforeEach(func() {
		hostKeyFile, hostPubKeyFile, _, _ = generateSSHKeypair()
		postgresRunner = postgresrunner.Runner{
			Port: 5433 + GinkgoParallelNode(),
		}
		dbProcess = ifrit.Invoke(postgresRunner)
		postgresRunner.CreateTestDB()

		concourseCommand = exec.Command(
			concoursePath,
			"web",
			"--tsa-host-key", hostKeyFile,
			"--postgres-user", "postgres",
			"--postgres-database", "testdb",
			"--postgres-port", strconv.Itoa(5433+GinkgoParallelNode()),
			"--main-team-local-user", "test",
			"--add-local-user", "test:test",
			"--debug-bind-port", strconv.Itoa(8000+GinkgoParallelNode()),
			"--bind-port", strconv.Itoa(8080+GinkgoParallelNode()),
			"--tsa-bind-port", strconv.Itoa(2222+GinkgoParallelNode()),
			"--client-id", "client-id",
			"--client-secret", "client-secret",
			"--tsa-client-id", "tsa-client-id",
			"--tsa-client-secret", "tsa-client-secret",
			"--tsa-token-url", "http://localhost/token",
		)
	})

	JustBeforeEach(func() {
		concourseRunner = ginkgomon.New(ginkgomon.Config{
			Command:       concourseCommand,
			Name:          "web",
			AnsiColorCode: "32m",
		})

		concourseProcess = ifrit.Background(concourseRunner)

		// workaround to avoid panic due to registering http handlers multiple times
		http.DefaultServeMux = new(http.ServeMux)
	})

	AfterEach(func() {
		ginkgomon.Interrupt(concourseProcess)
		<-concourseProcess.Wait()
		postgresRunner.DropTestDB()

		dbProcess.Signal(os.Interrupt)
		err := <-dbProcess.Wait()
		Expect(err).NotTo(HaveOccurred())
		os.Remove(hostKeyFile)
		os.Remove(hostPubKeyFile)
		os.Remove(filepath.Dir(hostPubKeyFile))
	})

	It("starts atc", func() {
		Eventually(concourseRunner.Buffer(), "30s", "2s").Should(gbytes.Say("atc.listening"))
	})

	It("starts tsa", func() {
		Eventually(concourseRunner.Buffer(), "30s", "2s").Should(gbytes.Say("tsa.listening"))
	})

	Context("when CONCOURSE_CONCURRENT_REQUEST_LIMIT is invalid", func() {
		BeforeEach(func() {
			concourseCommand.Env = append(concourseCommand.Env, "CONCOURSE_CONCURRENT_REQUEST_LIMIT=InvalidAction:0")
		})

		It("prints an error and exits", func() {
			Eventually(concourseRunner.Err()).Should(gbytes.Say("'InvalidAction' is not a valid action"))
		})
	})

	Context("with CONCOURSE_TSA_CLIENT_ID specified", func() {
		BeforeEach(func() {
			concourseCommand.Env = append(concourseCommand.Env, "CONCOURSE_TSA_CLIENT_ID=tsa-client-id")
		})

		It("starts atc", func() {
			Eventually(concourseRunner.Buffer(), "30s", "2s").Should(gbytes.Say("atc.listening"))
		})

		It("starts tsa", func() {
			Eventually(concourseRunner.Buffer(), "30s", "2s").Should(gbytes.Say("tsa.listening"))
		})

		Context("with CONCOURSE_SYSTEM_CLAIM_KEY is not set to 'aud'", func() {
			BeforeEach(func() {
				concourseCommand.Env = append(concourseCommand.Env, "CONCOURSE_SYSTEM_CLAIM_KEY=not-aud")
			})

			It("starts atc", func() {
				Eventually(concourseRunner.Buffer(), "30s", "2s").Should(gbytes.Say("atc.listening"))
			})

			It("starts tsa", func() {
				Eventually(concourseRunner.Buffer(), "30s", "2s").Should(gbytes.Say("tsa.listening"))
			})
		})

		Context("with CONCOURSE_SYSTEM_CLAIM_KEY set to 'aud'", func() {
			BeforeEach(func() {
				concourseCommand.Env = append(concourseCommand.Env, "CONCOURSE_SYSTEM_CLAIM_KEY=aud")
			})

			Context("when the system claim values does not contain the client id", func() {
				BeforeEach(func() {
					concourseCommand.Env = append(concourseCommand.Env, "CONCOURSE_SYSTEM_CLAIM_VALUE=system-claim-value-1,system-claim-value-2")
				})

				It("errors", func() {
					Eventually(concourseRunner.Err(), 5*time.Second).Should(
						gbytes.Say("at least one systemClaimValue must be equal to tsa-client-id"),
					)
				})
			})

			Context("when the system claim values contain the client id", func() {
				BeforeEach(func() {
					concourseCommand.Env = append(concourseCommand.Env, "CONCOURSE_SYSTEM_CLAIM_VALUE=system-claim-value-1,tsa-client-id")
				})

				It("starts atc", func() {
					Eventually(concourseRunner.Buffer(), "30s", "2s").Should(gbytes.Say("atc.listening"))
				})

				It("starts tsa", func() {
					Eventually(concourseRunner.Buffer(), "30s", "2s").Should(gbytes.Say("tsa.listening"))
				})
			})
		})
	})
})

func generateSSHKeypair() (string, string, *rsa.PrivateKey, ssh.PublicKey) {
	path, err := ioutil.TempDir("", "tsa-key")
	Expect(err).NotTo(HaveOccurred())

	privateKeyPath := filepath.Join(path, "id_rsa")
	publicKeyPath := privateKeyPath + ".pub"

	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	Expect(err).NotTo(HaveOccurred())

	privateKeyBytes := pem.EncodeToMemory(&pem.Block{
		Type:    "RSA PRIVATE KEY",
		Headers: nil,
		Bytes:   x509.MarshalPKCS1PrivateKey(privateKey),
	})

	publicKeyRsa, err := ssh.NewPublicKey(&privateKey.PublicKey)
	Expect(err).NotTo(HaveOccurred())

	publicKeyBytes := ssh.MarshalAuthorizedKey(publicKeyRsa)

	err = ioutil.WriteFile(privateKeyPath, privateKeyBytes, 0600)
	Expect(err).NotTo(HaveOccurred())

	err = ioutil.WriteFile(publicKeyPath, publicKeyBytes, 0600)
	Expect(err).NotTo(HaveOccurred())

	return privateKeyPath, publicKeyPath, privateKey, publicKeyRsa
}
<|endoftext|>


<commit_before>package main

var WorkerVersion = "0.0.0-dev"
<commit_msg>set worker version default to empty<commit_after>package main

var WorkerVersion = ""
<|endoftext|>


<commit_before>package main

import (
	"github.com/Symantec/Dominator/lib/fsutil"
	libprober "github.com/Symantec/health-agent/lib/prober"
	"github.com/Symantec/health-agent/lib/proberlist"
	dnsprober "github.com/Symantec/health-agent/probers/dns"
	ldapprober "github.com/Symantec/health-agent/probers/ldap"
	pidprober "github.com/Symantec/health-agent/probers/pidfile"
	testprogprober "github.com/Symantec/health-agent/probers/testprog"
	urlprober "github.com/Symantec/health-agent/probers/url"
	"github.com/Symantec/tricorder/go/tricorder"
	"github.com/Symantec/tricorder/go/tricorder/units"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"log"
	"os"
	"path"
	"strings"
	"time"
)

type testConfig struct {
	Testtype  string `yaml:"type"`
	Probefreq uint8  `yaml:"probe-freq"`
	Specs     testSpecs
}

type testSpecs struct {
	Pathname   string
	Urlpath    string `yaml:"url-path"`
	Urlport    int    `yaml:"url-port"`
	Hostname   string `yaml:"hostname"`
	SssdConfig string `yaml:"sssd-config"`
}

func setupHealthchecks(configDir string, pl *proberlist.ProberList,
	logger *log.Logger) error {
	topDir := "/health-checks"
	configdir, err := os.Open(path.Join(configDir, "tests.d"))
	defer configdir.Close()
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}
	configfiles, err := configdir.Readdir(0)
	if err != nil {
		return err
	}
	healthCheckers := make(map[string]libprober.HealthChecker)
	for _, configfile := range configfiles {
		if configfile.IsDir() {
			continue
		}
		data, err := ioutil.ReadFile(path.Join(configdir.Name(),
			configfile.Name()))
		if err != nil {
			logger.Printf("Unable to read file %q: %s",
				configfile.Name(), err)
			return err
		}
		c := testConfig{}
		if err := yaml.Unmarshal([]byte(data), &c); err != nil {
			logger.Printf("Error unmarshalling file %s: %q",
				configfile.Name(), err)
			return err
		}
		testname := strings.Split(configfile.Name(), ".")[0]
		if prober := makeProber(testname, &c, logger); prober != nil {
			pl.Add(prober, path.Join(topDir, testname), c.Probefreq)
			if hc, ok := prober.(libprober.HealthChecker); ok {
				healthCheckers[testname] = hc
			}
		}
	}
	var allHealthy bool
	list := tricorder.NewList([]string{}, tricorder.ImmutableSlice)
	group := tricorder.NewGroup()
	err = tricorder.RegisterMetricInGroup(path.Join(topDir, "*", "healthy"),
		&allHealthy, group, units.None,
		"If true, all health checks are healthy")
	if err != nil {
		return err
	}
	err = tricorder.RegisterMetricInGroup(path.Join(topDir, "*",
		"unhealthy-list"), list, group, units.None,
		"List of failed health checks")
	if err != nil {
		return err
	}
	group.RegisterUpdateFunc(func() time.Time {
		healthy := true
		var unhealthyList []string
		for testname, healthChecker := range healthCheckers {
			if !healthChecker.HealthCheck() {
				healthy = false
				unhealthyList = append(unhealthyList, testname)
			}
		}
		allHealthy = healthy
		list.Change(unhealthyList, tricorder.ImmutableSlice)
		return time.Now()
	})
	return nil
}

func makeProber(testname string, c *testConfig,
	logger *log.Logger) proberlist.RegisterProber {
	switch c.Testtype {
	case "dns":
		hostname := c.Specs.Hostname
		return dnsprober.New(testname, hostname)
	case "ldap":
		sssd := c.Specs.SssdConfig
		file, err := fsutil.WaitFile(sssd, time.Minute*5)
		if err != nil {
			logger.Println(err)
			return nil
		}
		return ldapprober.Makeldapprober(testname, file, c.Probefreq)
	case "pid":
		pidpath := c.Specs.Pathname
		if pidpath == "" {
			return nil
		}
		return pidprober.Makepidprober(testname, pidpath)
	case "testprog":
		testprogpath := c.Specs.Pathname
		if testprogpath == "" {
			return nil
		}
		return testprogprober.Maketestprogprober(testname, testprogpath)
	case "url":
		urlpath := c.Specs.Urlpath
		urlport := c.Specs.Urlport
		return urlprober.Makeurlprober(testname, urlpath, urlport)
	default:
		logger.Printf("Test type %s not supported", c.Testtype)
		return nil
	}
}
<commit_msg>Switch to tricorder/Group.RegisterMetric() method.<commit_after>package main

import (
	"github.com/Symantec/Dominator/lib/fsutil"
	libprober "github.com/Symantec/health-agent/lib/prober"
	"github.com/Symantec/health-agent/lib/proberlist"
	dnsprober "github.com/Symantec/health-agent/probers/dns"
	ldapprober "github.com/Symantec/health-agent/probers/ldap"
	pidprober "github.com/Symantec/health-agent/probers/pidfile"
	testprogprober "github.com/Symantec/health-agent/probers/testprog"
	urlprober "github.com/Symantec/health-agent/probers/url"
	"github.com/Symantec/tricorder/go/tricorder"
	"github.com/Symantec/tricorder/go/tricorder/units"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"log"
	"os"
	"path"
	"strings"
	"time"
)

type testConfig struct {
	Testtype  string `yaml:"type"`
	Probefreq uint8  `yaml:"probe-freq"`
	Specs     testSpecs
}

type testSpecs struct {
	Pathname   string
	Urlpath    string `yaml:"url-path"`
	Urlport    int    `yaml:"url-port"`
	Hostname   string `yaml:"hostname"`
	SssdConfig string `yaml:"sssd-config"`
}

func setupHealthchecks(configDir string, pl *proberlist.ProberList,
	logger *log.Logger) error {
	topDir := "/health-checks"
	configdir, err := os.Open(path.Join(configDir, "tests.d"))
	defer configdir.Close()
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}
	configfiles, err := configdir.Readdir(0)
	if err != nil {
		return err
	}
	healthCheckers := make(map[string]libprober.HealthChecker)
	for _, configfile := range configfiles {
		if configfile.IsDir() {
			continue
		}
		data, err := ioutil.ReadFile(path.Join(configdir.Name(),
			configfile.Name()))
		if err != nil {
			logger.Printf("Unable to read file %q: %s",
				configfile.Name(), err)
			return err
		}
		c := testConfig{}
		if err := yaml.Unmarshal([]byte(data), &c); err != nil {
			logger.Printf("Error unmarshalling file %s: %q",
				configfile.Name(), err)
			return err
		}
		testname := strings.Split(configfile.Name(), ".")[0]
		if prober := makeProber(testname, &c, logger); prober != nil {
			pl.Add(prober, path.Join(topDir, testname), c.Probefreq)
			if hc, ok := prober.(libprober.HealthChecker); ok {
				healthCheckers[testname] = hc
			}
		}
	}
	var allHealthy bool
	list := tricorder.NewList([]string{}, tricorder.ImmutableSlice)
	group := tricorder.NewGroup()
	err = group.RegisterMetric(path.Join(topDir, "*", "healthy"), &allHealthy,
		units.None, "If true, all health checks are healthy")
	if err != nil {
		return err
	}
	err = tricorder.RegisterMetricInGroup(path.Join(topDir, "*",
		"unhealthy-list"), list, group, units.None,
		"List of failed health checks")
	if err != nil {
		return err
	}
	group.RegisterUpdateFunc(func() time.Time {
		healthy := true
		var unhealthyList []string
		for testname, healthChecker := range healthCheckers {
			if !healthChecker.HealthCheck() {
				healthy = false
				unhealthyList = append(unhealthyList, testname)
			}
		}
		allHealthy = healthy
		list.Change(unhealthyList, tricorder.ImmutableSlice)
		return time.Now()
	})
	return nil
}

func makeProber(testname string, c *testConfig,
	logger *log.Logger) proberlist.RegisterProber {
	switch c.Testtype {
	case "dns":
		hostname := c.Specs.Hostname
		return dnsprober.New(testname, hostname)
	case "ldap":
		sssd := c.Specs.SssdConfig
		file, err := fsutil.WaitFile(sssd, time.Minute*5)
		if err != nil {
			logger.Println(err)
			return nil
		}
		return ldapprober.Makeldapprober(testname, file, c.Probefreq)
	case "pid":
		pidpath := c.Specs.Pathname
		if pidpath == "" {
			return nil
		}
		return pidprober.Makepidprober(testname, pidpath)
	case "testprog":
		testprogpath := c.Specs.Pathname
		if testprogpath == "" {
			return nil
		}
		return testprogprober.Maketestprogprober(testname, testprogpath)
	case "url":
		urlpath := c.Specs.Urlpath
		urlport := c.Specs.Urlport
		return urlprober.Makeurlprober(testname, urlpath, urlport)
	default:
		logger.Printf("Test type %s not supported", c.Testtype)
		return nil
	}
}
<|endoftext|>