<commit_before>// 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.

// A CLI application to monitor one or more stocks.
package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/mattn/go-runewidth"
	"github.com/nsf/termbox-go"
)

const TIMEOUT = time.Duration(time.Second * 10)                                   // how long to wait on a call
const URL = "http://finance.yahoo.com/webservice/v1/symbols/%s/quote?format=json" // where we're getting our data from
const UP = "↑"                                                                    // rune 8593
const DOWN = "↓"                                                                  // rune 8595

var re = regexp.MustCompile(`^\d.+\.\d{2}`) // this is to have only 2 decimal places
var signalChan = make(chan os.Signal, 1)    // channel to catch ctrl-c

// Flag variables to hold CLI arguments
var (
	symbolFlag   = flag.String("s", "", "Symbols for ticker, comma seperate (no spaces)")
	intervalFlag = flag.Int("i", 1, "Interval for stock data to be updated in seconds")
)

// Stock is the top level of the returned JSON
type Stock struct {
	List List `json:"list"`
}

// List hold the metadata and list of returned symbol data
type List struct {
	Meta      Meta        `json:"meta"`
	Resources []Resources `json:"resources"`
}

// Meta is the calls metadata
type Meta struct {
	Type  string `json:"type"`
	Start uint   `json:"start"`
	Count uint   `json:"count"`
}

// Resources holds a JSON obj with the symbol data
type Resources struct {
	Resource Resource `json:"resource"`
}

// Resource contains the actual JSON obj with the symbol data
type Resource struct {
	Classname string `json:"classname"`
	Fields    Fields `json:"fields"`
}

// Fields holds all of the retrieved data from the API call
type Fields struct {
	Name    string `json:"name"`    // name of company
	Price   string `json:"price"`   // current price
	Symbol  string `json:"symbol"`  // stock symbol
	TS      string `json:"ts"`      //
	Type    string `json:"type"`    // type of stock (equity, etc...)
	UTCTime string `json:"utctime"` // time in UTC
	Volume  string `json:"volume"`  // shares traded
}

// stockwatcher holds the relevant data for the running instance
type stockwatcher struct {
	quotes   map[string]map[string]float64
	interval time.Duration
	m        *sync.Mutex
}

// NewStockWatcher returns a new instance of stockwatcher with the given parameters
func NewStockWatcher(i time.Duration) *stockwatcher {
	return &stockwatcher{
		quotes:   make(map[string]map[string]float64),
		interval: i,
		m:        &sync.Mutex{},
	}
}

// add takes the given symbol and populates a key in the quotes map
func (t *stockwatcher) add(symbol string) {
	t.m.Lock()
	defer t.m.Unlock()
	if _, ok := t.quotes[symbol]; !ok {
		t.quotes[symbol] = map[string]float64{}
	}
}

// updateStock populates stockwatcher struct with stock prices
func (t *stockwatcher) updateStock(symbol string, price float64) {
	t.m.Lock()
	defer t.m.Unlock()
	t.quotes[symbol] = map[string]float64{
		"previous": t.quotes[symbol]["current"],
		"current":  price,
	}
}

// query will retrieve data for a given symbol
func query(symbol string) (*Stock, error) {
	data := &Stock{}
	client := http.Client{
		Timeout: TIMEOUT,
	}

	resp, err := client.Get(fmt.Sprintf(URL, symbol))
	if err != nil {
		return nil, errors.New("unable to retrive symbol data")
	}
	defer resp.Body.Close()

	if err := json.NewDecoder(resp.Body).Decode(&data); err != nil {
		log.Fatalln(err)
	}
	return data, nil
}

// convertPrice converts the given string to a float64 value
func convertPrice(p string) float64 {
	price, err := strconv.ParseFloat(p, 64)
	if err != nil {
		log.Fatalln(err)
		os.Exit(1)
	}
	return price
}

// runner goes through and gets the data for each symbol
func (t *stockwatcher) runner() {
	var wg sync.WaitGroup
	for k, _ := range t.quotes {
		wg.Add(1)
		go func(k string) {
			defer wg.Done()
			stock, err := query(k)
			if err != nil { // if we can't get a response from the API, put 0.00 in and keep going
				t.updateStock(k, 0.00)
				return
			}
			t.updateStock(stock.List.Resources[0].Resource.Fields.Symbol,
				convertPrice(re.FindString(stock.List.Resources[0].Resource.Fields.Price)),
			)
		}(k)
	}
	wg.Wait()
}

// formatData formats the given data for printing
func (t *stockwatcher) formatData() {
	// populate list with keys from quote map
	var keys []string
	for k := range t.quotes {
		keys = append(keys, k)
	}
	sort.Strings(keys) // alphabetize keys

	pos := 1
	for _, k := range keys {
		// print format for first run or if not change detected from previous run
		if t.quotes[k]["previous"] == 0.00 || t.quotes[k]["previous"] == t.quotes[k]["current"] {
			printTb(1,
				pos,
				fmt.Sprintf("%-6s %-7v %11s %-4s\n", k, t.quotes[k]["current"], "%", "-"),
				termbox.ColorWhite, termbox.ColorDefault,
			)
			pos++
			// print format in green if current price being is greater than previous price
		} else if t.quotes[k]["current"] > t.quotes[k]["previous"] {
			printTb(1,
				pos,
				fmt.Sprintf("%-6s %-7v +%-.6f %% %-4s\n", k, t.quotes[k]["current"], t.quotes[k]["current"]/t.quotes[k]["previous"], UP),
				termbox.ColorGreen,
				termbox.ColorDefault,
			)
			pos++
			// print format in red if current price being is lesser than previous price
		} else {
			printTb(1,
				pos,
				fmt.Sprintf("%-6s %-7v -%-.6f %% %-4s\n", k, t.quotes[k]["current"], t.quotes[k]["current"]/t.quotes[k]["previous"], DOWN),
				termbox.ColorRed,
				termbox.ColorDefault,
			)
			pos++
		}
	}
}

// printTb prints the given data out to the screen
func printTb(x, y int, msg string, fg, bg termbox.Attribute) {
	for _, c := range []rune(msg) {
		termbox.SetCell(x, y, c, fg, bg)
		x += runewidth.RuneWidth(c)
	}
	termbox.Flush()
}

func main() {
	flag.Parse()

	// make sure we got what was expected from the CLI
	if flag.NFlag() != 2 || *symbolFlag == "" {
		flag.Usage()
		os.Exit(1)
	}

	t := NewStockWatcher(time.Duration(*intervalFlag) * time.Second)

	// check if more than one symbol has been given
	switch {
	case strings.Contains(*symbolFlag, ","):
		for _, a := range strings.Split(*symbolFlag, ",") {
			t.add(a)
		}
	default:
		t.add(*symbolFlag)
	}

	// initialize termbox
	err := termbox.Init()
	if err != nil {
		log.Fatal(err)
	}
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)

	event := make(chan termbox.Event)
	go func() {
		for {
			// Post events to channel
			event <- termbox.PollEvent()
		}
	}()

loop:
	for {
		t.runner()
		t.formatData()

		// Poll key event or timeout (maybe)
		select {
		case <-event:
			break loop
		case <-time.After(t.interval):
			continue loop
		}
	}
	close(event)
	time.Sleep(1 * time.Second)
	termbox.Close()
	os.Exit(0)
}
<commit_msg>small comment update<commit_after>// 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.

// A CLI application to monitor one or more stocks.
package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/mattn/go-runewidth"
	"github.com/nsf/termbox-go"
)

const TIMEOUT = time.Duration(time.Second * 10)                                   // how long to wait on a call
const URL = "http://finance.yahoo.com/webservice/v1/symbols/%s/quote?format=json" // where we're getting our data from
const UP = "↑"                                                                    // rune 8593
const DOWN = "↓"                                                                  // rune 8595

var re = regexp.MustCompile(`^\d.+\.\d{2}`) // this is to have only 2 decimal places
var signalChan = make(chan os.Signal, 1)    // channel to catch ctrl-c

// Flag variables to hold CLI arguments
var (
	symbolFlag   = flag.String("s", "", "Symbols for ticker, comma seperate (no spaces)")
	intervalFlag = flag.Int("i", 1, "Interval for stock data to be updated in seconds")
)

// Stock is the top level of the returned JSON
type Stock struct {
	List List `json:"list"`
}

// List hold the metadata and list of returned symbol data
type List struct {
	Meta      Meta        `json:"meta"`
	Resources []Resources `json:"resources"`
}

// Meta is the calls metadata
type Meta struct {
	Type  string `json:"type"`
	Start uint   `json:"start"`
	Count uint   `json:"count"`
}

// Resources holds a JSON obj with the symbol data
type Resources struct {
	Resource Resource `json:"resource"`
}

// Resource contains the actual JSON obj with the symbol data
type Resource struct {
	Classname string `json:"classname"`
	Fields    Fields `json:"fields"`
}

// Fields holds all of the retrieved data from the API call
type Fields struct {
	Name    string `json:"name"`    // name of company
	Price   string `json:"price"`   // current price
	Symbol  string `json:"symbol"`  // stock symbol
	TS      string `json:"ts"`      //
	Type    string `json:"type"`    // type of stock (equity, etc...)
	UTCTime string `json:"utctime"` // time in UTC
	Volume  string `json:"volume"`  // shares traded
}

// stockwatcher holds the relevant data for the running instance
type stockwatcher struct {
	quotes   map[string]map[string]float64
	interval time.Duration
	m        *sync.Mutex
}

// NewStockWatcher returns a new instance of stockwatcher with the given parameters
func NewStockWatcher(i time.Duration) *stockwatcher {
	return &stockwatcher{
		quotes:   make(map[string]map[string]float64),
		interval: i,
		m:        &sync.Mutex{},
	}
}

// add takes the given symbol and populates a key in the quotes map
func (t *stockwatcher) add(symbol string) {
	t.m.Lock()
	defer t.m.Unlock()
	if _, ok := t.quotes[symbol]; !ok {
		t.quotes[symbol] = map[string]float64{}
	}
}

// updateStock populates stockwatcher struct with stock prices
func (t *stockwatcher) updateStock(symbol string, price float64) {
	t.m.Lock()
	defer t.m.Unlock()
	t.quotes[symbol] = map[string]float64{
		"previous": t.quotes[symbol]["current"],
		"current":  price,
	}
}

// query will retrieve data for a given symbol
func query(symbol string) (*Stock, error) {
	data := &Stock{}
	client := http.Client{
		Timeout: TIMEOUT,
	}

	resp, err := client.Get(fmt.Sprintf(URL, symbol))
	if err != nil {
		return nil, errors.New("unable to retrive symbol data")
	}
	defer resp.Body.Close()

	if err := json.NewDecoder(resp.Body).Decode(&data); err != nil {
		log.Fatalln(err)
	}
	return data, nil
}

// convertPrice converts the given string to a float64 value
func convertPrice(p string) float64 {
	price, err := strconv.ParseFloat(p, 64)
	if err != nil {
		log.Fatalln(err)
		os.Exit(1)
	}
	return price
}

// runner goes through and gets the data for each symbol
func (t *stockwatcher) runner() {
	var wg sync.WaitGroup
	for k, _ := range t.quotes {
		wg.Add(1)
		go func(k string) {
			defer wg.Done()
			stock, err := query(k)
			if err != nil { // if we can't get a response from the API, put 0.00's in and keep going
				t.updateStock(k, 0.00)
				return
			}
			t.updateStock(stock.List.Resources[0].Resource.Fields.Symbol,
				convertPrice(re.FindString(stock.List.Resources[0].Resource.Fields.Price)),
			)
		}(k)
	}
	wg.Wait()
}

// formatData formats the given data for printing
func (t *stockwatcher) formatData() {
	// populate list with keys from quote map
	var keys []string
	for k := range t.quotes {
		keys = append(keys, k)
	}
	sort.Strings(keys) // alphabetize keys

	pos := 1
	for _, k := range keys {
		// print format for first run or if not change detected from previous run
		if t.quotes[k]["previous"] == 0.00 || t.quotes[k]["previous"] == t.quotes[k]["current"] {
			printTb(1,
				pos,
				fmt.Sprintf("%-6s %-7v %11s %-4s\n", k, t.quotes[k]["current"], "%", "-"),
				termbox.ColorWhite, termbox.ColorDefault,
			)
			pos++
			// print format in green if current price being is greater than previous price
		} else if t.quotes[k]["current"] > t.quotes[k]["previous"] {
			printTb(1,
				pos,
				fmt.Sprintf("%-6s %-7v +%-.6f %% %-4s\n", k, t.quotes[k]["current"], t.quotes[k]["current"]/t.quotes[k]["previous"], UP),
				termbox.ColorGreen,
				termbox.ColorDefault,
			)
			pos++
			// print format in red if current price being is lesser than previous price
		} else {
			printTb(1,
				pos,
				fmt.Sprintf("%-6s %-7v -%-.6f %% %-4s\n", k, t.quotes[k]["current"], t.quotes[k]["current"]/t.quotes[k]["previous"], DOWN),
				termbox.ColorRed,
				termbox.ColorDefault,
			)
			pos++
		}
	}
}

// printTb prints the given data out to the screen
func printTb(x, y int, msg string, fg, bg termbox.Attribute) {
	for _, c := range []rune(msg) {
		termbox.SetCell(x, y, c, fg, bg)
		x += runewidth.RuneWidth(c)
	}
	termbox.Flush()
}

func main() {
	flag.Parse()

	// make sure we got what was expected from the CLI
	if flag.NFlag() != 2 || *symbolFlag == "" {
		flag.Usage()
		os.Exit(1)
	}

	t := NewStockWatcher(time.Duration(*intervalFlag) * time.Second)

	// check if more than one symbol has been given
	switch {
	case strings.Contains(*symbolFlag, ","):
		for _, a := range strings.Split(*symbolFlag, ",") {
			t.add(a)
		}
	default:
		t.add(*symbolFlag)
	}

	// initialize termbox
	err := termbox.Init()
	if err != nil {
		log.Fatal(err)
	}
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)

	event := make(chan termbox.Event)
	go func() {
		for {
			// Post events to channel
			event <- termbox.PollEvent()
		}
	}()

loop:
	for {
		t.runner()
		t.formatData()

		// Poll key event or timeout (maybe)
		select {
		case <-event:
			break loop
		case <-time.After(t.interval):
			continue loop
		}
	}
	close(event)
	time.Sleep(1 * time.Second)
	termbox.Close()
	os.Exit(0)
}
<|endoftext|>


<commit_before>package storage

import (
	"encoding/binary"
	"errors"
	"hash/crc32"
	"io"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
)

var (
	// ErrIsDirectory means that directory is found where file assumed.
	ErrIsDirectory = errors.New("Got directory, need file")
	// ErrBadHeader means that blob is unable to initialize due header corruption or wrong file format.
	ErrBadHeader = errors.New("Bad header magic bytes, file corrupted or in wrong format")
	// ErrBadHeaderCapacity means that decoded capacity from header is less than actual file size.
	ErrBadHeaderCapacity = errors.New("Capacity in header is less than actual file size, file can be corrupted")
	// ErrBadHeaderCRC means that header crc check failed.
	ErrBadHeaderCRC = errors.New("Header CRC missmatch")
)

// Allocator wraps Allocate method for allocating slices. Should be goroutine-safe.
type Allocator interface {
	Allocate(size int64) (offset int64, err error)
}

// StatBackend wraps Stat method for retrieving file stats.
type StatBackend interface {
	Stat() (os.FileInfo, error)
}

// TruncateSyncer is the interface that groups Sync and Truncate methods.
type TruncateSyncer interface {
	// Sync commits the current contents of the file to stable storage.
	Sync() error
	// Truncate changes the size of the storage.
	Truncate(size int64) error
}

// BlobBackend is the interface that groups basic methods for consistent storage.
type BlobBackend interface {
	io.ReaderAt
	io.WriterAt
	io.Closer
	TruncateSyncer
}

// Blob represents set of data slices on top of BlobBackend.
type Blob struct {
	sync.RWMutex
	Backend    BlobBackend
	Size       int64
	Capacity   int64
	headerBuff [blobHeaderSize]byte
}

// Sync commits the current state of blob.
func (b *Blob) Sync() (err error) {
	b.Lock()
	if err = b.writeHeader(); err == nil {
		err = b.Backend.Sync()
	}
	b.Unlock()
	return err
}

// Truncate changes the capacity of the blob.
func (b *Blob) Truncate(size int64) (err error) {
	b.Lock()
	if err = b.Backend.Truncate(size); err == nil {
		b.Capacity = size
		// rendering capacity changes to header
		err = b.writeHeader()
	}
	b.Unlock()
	return err
}

// Allocate returns offset to atomically allocated slice of provided size and error if any.
// After allocation it is safe to call WriteAt(b, offset) with len(b) = size.
//
// Example:
//     data := make([]byte, size)
//     offset, _ := b.Allocate(size)
//     b.WriteAt(data, offset)
func (b *Blob) Allocate(size int64) (int64, error) {
	newSize := atomic.AddInt64(&b.Size, size)
	return newSize - size, nil
}

// readHeader encodes header to start of backend and returns error if any.
// It is not goroutine-safe. Can return errors from backend.
// Uses b.headerBuff as write buffer.
func (b *Blob) writeHeader() error {
	header := BlobHeader{
		Size:     b.Size,
		Capacity: b.Capacity,
	}
	header.Put(b.headerBuff[:])
	_, err := b.Backend.WriteAt(b.headerBuff[:], 0)
	if b.Size == 0 {
		b.Size = blobHeaderSize
	}
	return err
}

// readHeader decodes header from start of backend and returns error if any.
// It is not goroutine-safe.
// Can return ErrBadHeaderCapacity, ErrBadHeader and errors from backend.
// Uses b.headerBuff as read buffer.
func (b *Blob) readHeader() error {
	if _, err := b.Backend.ReadAt(b.headerBuff[:], 0); err != nil {
		return err
	}
	h := BlobHeader{}
	if err := h.Read(b.headerBuff[:]); err != nil {
		return err
	}
	if b.Capacity < h.Capacity {
		return ErrBadHeaderCapacity
	}
	b.Size = h.Size
	return nil
}

const (
	// DefaultBlobSize is initial capacity for newly created blob.
	DefaultBlobSize = 1024
)

// BlobConfig is configuration for blob processing.
type BlobConfig struct {
	InitialSize int64
}

// GetInitialSize returns initial size for the blob used upon creation.
func (i *BlobConfig) GetInitialSize() int64 {
	if i == nil || i.InitialSize == 0 {
		return DefaultBlobSize
	}
	return i.InitialSize
}

// OpenBlob opens or creates a Blob for the given path.
//
// The returned Blob instance is goroutine-safe.
// The Blob must be closed after use, by calling Close method.
func OpenBlob(path string, cfg *BlobConfig) (*Blob, error) {
	f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0644)
	if err != nil {
		return nil, err
	}
	stat, err := f.Stat()
	if err != nil {
		return nil, err
	}
	if stat.IsDir() {
		return nil, ErrIsDirectory
	}
	b := &Blob{
		Capacity: stat.Size(),
		Size:     0,
		Backend:  f,
	}
	if b.Capacity == 0 {
		err = b.Truncate(cfg.GetInitialSize())
	} else {
		err = b.readHeader()
	}
	runtime.SetFinalizer(b, (*Blob).Close)
	return b, err
}

// Close closes the Blob, rendering it unusable for changes.
// It returns an error, if any.
func (b *Blob) Close() error {
	if err := b.Sync(); err != nil {
		return err
	}
	return b.Backend.Close()
}

// BlobHeader contains info about Blob size and capacity.
type BlobHeader struct {
	Size     int64
	Capacity int64
}

// blobHeaderSize = magic + size + capacity + crc
const blobHeaderSize = 8 + 8 + 8 + 8

var blobHeaderMagic = [...]byte{
	0xbb,
	0xba,
	0xbd,
	0xbb,
	0x13,
	0x37,
	0x20,
	0x16,
}

// Put encodes BlobHeader into buf and returns the number of bytes written.
// If the buffer is too small, Put will panic.
func (h BlobHeader) Put(buf []byte) int {
	var offset = 8
	copy(buf[:offset], blobHeaderMagic[:])
	binary.PutVarint(buf[offset:], h.Size)
	offset += 8
	binary.PutVarint(buf[offset:offset+8], h.Capacity)
	offset += 8
	crc := crc32.ChecksumIEEE(buf[:offset])
	binary.PutUvarint(buf[offset:offset+8], uint64(crc))
	offset += 8
	return offset
}

// Read decodes BlobHeader from buf and returns ErrBadHeader if it fails.
func (h *BlobHeader) Read(buf []byte) error {
	// checking header magic
	for i, v := range blobHeaderMagic {
		if v != buf[i] {
			return ErrBadHeader
		}
	}
	offset := 8
	h.Size, _ = binary.Varint(buf[offset:])
	offset += 8
	h.Capacity, _ = binary.Varint(buf[offset:])
	offset += 8
	// calculating crc
	expectedCRC := crc32.ChecksumIEEE(buf[:offset])
	// checking crc
	if crc, _ := binary.Uvarint(buf[offset:]); crc != uint64(expectedCRC) {
		return ErrBadHeaderCRC
	}
	return nil
}
<commit_msg>storage: add comment to blobHeaderMagic<commit_after>package storage

import (
	"encoding/binary"
	"errors"
	"hash/crc32"
	"io"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
)

var (
	// ErrIsDirectory means that directory is found where file assumed.
	ErrIsDirectory = errors.New("Got directory, need file")
	// ErrBadHeader means that blob is unable to initialize due header corruption or wrong file format.
	ErrBadHeader = errors.New("Bad header magic bytes, file corrupted or in wrong format")
	// ErrBadHeaderCapacity means that decoded capacity from header is less than actual file size.
	ErrBadHeaderCapacity = errors.New("Capacity in header is less than actual file size, file can be corrupted")
	// ErrBadHeaderCRC means that header crc check failed.
	ErrBadHeaderCRC = errors.New("Header CRC missmatch")
)

// Allocator wraps Allocate method for allocating slices. Should be goroutine-safe.
type Allocator interface {
	Allocate(size int64) (offset int64, err error)
}

// StatBackend wraps Stat method for retrieving file stats.
type StatBackend interface {
	Stat() (os.FileInfo, error)
}

// TruncateSyncer is the interface that groups Sync and Truncate methods.
type TruncateSyncer interface {
	// Sync commits the current contents of the file to stable storage.
	Sync() error
	// Truncate changes the size of the storage.
	Truncate(size int64) error
}

// BlobBackend is the interface that groups basic methods for consistent storage.
type BlobBackend interface {
	io.ReaderAt
	io.WriterAt
	io.Closer
	TruncateSyncer
}

// Blob represents set of data slices on top of BlobBackend.
type Blob struct {
	sync.RWMutex
	Backend    BlobBackend
	Size       int64
	Capacity   int64
	headerBuff [blobHeaderSize]byte
}

// Sync commits the current state of blob.
func (b *Blob) Sync() (err error) {
	b.Lock()
	if err = b.writeHeader(); err == nil {
		err = b.Backend.Sync()
	}
	b.Unlock()
	return err
}

// Truncate changes the capacity of the blob.
func (b *Blob) Truncate(size int64) (err error) {
	b.Lock()
	if err = b.Backend.Truncate(size); err == nil {
		b.Capacity = size
		// rendering capacity changes to header
		err = b.writeHeader()
	}
	b.Unlock()
	return err
}

// Allocate returns offset to atomically allocated slice of provided size and error if any.
// After allocation it is safe to call WriteAt(b, offset) with len(b) = size.
//
// Example:
//     data := make([]byte, size)
//     offset, _ := b.Allocate(size)
//     b.WriteAt(data, offset)
func (b *Blob) Allocate(size int64) (int64, error) {
	newSize := atomic.AddInt64(&b.Size, size)
	return newSize - size, nil
}

// readHeader encodes header to start of backend and returns error if any.
// It is not goroutine-safe. Can return errors from backend.
// Uses b.headerBuff as write buffer.
func (b *Blob) writeHeader() error {
	header := BlobHeader{
		Size:     b.Size,
		Capacity: b.Capacity,
	}
	header.Put(b.headerBuff[:])
	_, err := b.Backend.WriteAt(b.headerBuff[:], 0)
	if b.Size == 0 {
		b.Size = blobHeaderSize
	}
	return err
}

// readHeader decodes header from start of backend and returns error if any.
// It is not goroutine-safe.
// Can return ErrBadHeaderCapacity, ErrBadHeader and errors from backend.
// Uses b.headerBuff as read buffer.
func (b *Blob) readHeader() error {
	if _, err := b.Backend.ReadAt(b.headerBuff[:], 0); err != nil {
		return err
	}
	h := BlobHeader{}
	if err := h.Read(b.headerBuff[:]); err != nil {
		return err
	}
	if b.Capacity < h.Capacity {
		return ErrBadHeaderCapacity
	}
	b.Size = h.Size
	return nil
}

const (
	// DefaultBlobSize is initial capacity for newly created blob.
	DefaultBlobSize = 1024
)

// BlobConfig is configuration for blob processing.
type BlobConfig struct {
	InitialSize int64
}

// GetInitialSize returns initial size for the blob used upon creation.
func (i *BlobConfig) GetInitialSize() int64 {
	if i == nil || i.InitialSize == 0 {
		return DefaultBlobSize
	}
	return i.InitialSize
}

// OpenBlob opens or creates a Blob for the given path.
//
// The returned Blob instance is goroutine-safe.
// The Blob must be closed after use, by calling Close method.
func OpenBlob(path string, cfg *BlobConfig) (*Blob, error) {
	f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0644)
	if err != nil {
		return nil, err
	}
	stat, err := f.Stat()
	if err != nil {
		return nil, err
	}
	if stat.IsDir() {
		return nil, ErrIsDirectory
	}
	b := &Blob{
		Capacity: stat.Size(),
		Size:     0,
		Backend:  f,
	}
	if b.Capacity == 0 {
		err = b.Truncate(cfg.GetInitialSize())
	} else {
		err = b.readHeader()
	}
	runtime.SetFinalizer(b, (*Blob).Close)
	return b, err
}

// Close closes the Blob, rendering it unusable for changes.
// It returns an error, if any.
func (b *Blob) Close() error {
	if err := b.Sync(); err != nil {
		return err
	}
	return b.Backend.Close()
}

// BlobHeader contains info about Blob size and capacity.
type BlobHeader struct {
	Size     int64
	Capacity int64
}

// blobHeaderSize = magic + size + capacity + crc
const blobHeaderSize = 8 + 8 + 8 + 8

// blobHeaderMagic are magic bytes at start of blob header.
var blobHeaderMagic = [...]byte{
	0xbb,
	0xba,
	0xbd,
	0xbb,
	0x13,
	0x37,
	0x20,
	0x16,
}

// Put encodes BlobHeader into buf and returns the number of bytes written.
// If the buffer is too small, Put will panic.
func (h BlobHeader) Put(buf []byte) int {
	var offset = 8
	copy(buf[:offset], blobHeaderMagic[:])
	binary.PutVarint(buf[offset:], h.Size)
	offset += 8
	binary.PutVarint(buf[offset:offset+8], h.Capacity)
	offset += 8
	crc := crc32.ChecksumIEEE(buf[:offset])
	binary.PutUvarint(buf[offset:offset+8], uint64(crc))
	offset += 8
	return offset
}

// Read decodes BlobHeader from buf and returns ErrBadHeader if it fails.
func (h *BlobHeader) Read(buf []byte) error {
	// checking header magic
	for i, v := range blobHeaderMagic {
		if v != buf[i] {
			return ErrBadHeader
		}
	}
	offset := 8
	h.Size, _ = binary.Varint(buf[offset:])
	offset += 8
	h.Capacity, _ = binary.Varint(buf[offset:])
	offset += 8
	// calculating crc
	expectedCRC := crc32.ChecksumIEEE(buf[:offset])
	// checking crc
	if crc, _ := binary.Uvarint(buf[offset:]); crc != uint64(expectedCRC) {
		return ErrBadHeaderCRC
	}
	return nil
}
<|endoftext|>


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

package registry

import (
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"github.com/jacobsa/aws/sdb"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/crypto"
	"math/rand"
	"regexp"
	"strconv"
	"time"
	"unicode/utf8"
)

const (
	// The item name we use for "domain is already in use" markers.
	markerItemName = "comeback_marker"

	// The attribute name we use for storing crypto-compatibility data.
	cryptoMarkerAttributeName = "encrypted_data"

	// A time format that works properly with range queries.
	iso8601TimeFormat = "2006-01-02T15:04:05Z"
)

// A regexp for selected item names.
var itemNameRegexp = regexp.MustCompile(`^backup_([0-9a-f]{16})$`)

type Registry interface {
	// Record that the named backup job has completed.
	RecordBackup(j CompletedJob) (err error)

	// Return a list of the most recent completed backups.
	ListRecentBackups() (jobs []CompletedJob, err error)

	// Find a particular completed job by ID.
	FindBackup(jobId uint64) (job CompletedJob, err error)
}

func verifyCompatible(
	markerAttr sdb.Attribute,
	crypter crypto.Crypter) (err error) {
	// The ciphertext should be base64-encoded.
	encoded := markerAttr.Value
	ciphertext, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		err = fmt.Errorf("base64.DecodeString(%s): %v", encoded, err)
		return
	}

	// Attempt to decrypt the ciphertext.
	if _, err = crypter.Decrypt(ciphertext); err != nil {
		// Special case: Did the crypter signal that the key was wrong?
		if _, ok := err.(*crypto.NotAuthenticError); ok {
			err = fmt.Errorf("The supplied password is incorrect.")
			return
		}

		// Generic error.
		err = fmt.Errorf("Decrypt: %v", err)
		return
	}

	return
}

// Create a registry that stores data in the supplied SimpleDB domain,
// encrypting using a key derived from the supplied password using the supplied
// function.
//
// Before doing so, check to see whether this domain has been used as a
// registry before. If not, write an encrypted marker. If it has been used
// before, make sure that it was used with the same password. This prevents
// accidentally writing data with the wrong key when the user enters the wrong
// password.
//
// The crypter must be set up such that it is guaranteed to return an error if
// it is used to decrypt ciphertext encrypted with a different key.
func NewRegistry(
	domain sdb.Domain,
	cryptoPassword string,
	deriveKey func (password string, salt []byte) []byte,
	randSrc *rand.Rand,
) (r Registry, err error) {
	return newRegistry(
		domain,
		cryptoPassword,
		deriveKey,
		crypto.NewCrypter,
		randSrc)
}

// A version split out for testability.
func newRegistry(
	domain sdb.Domain,
	cryptoPassword string,
	deriveKey func (password string, salt []byte) []byte,
	createCrypter func (key []byte) (crypto.Crypter, error),
	randSrc *rand.Rand,
) (r Registry, err error) {
	// Set up a tentative result.
	r = &registry{crypter, domain.Db(), domain}

	// Ask for the data that will tell us whether the crypter is compatible with
	// the previous one used in this domain, if any.
	attrs, err := domain.GetAttributes(
		markerItemName,
		false, // No need to ask for a consistent read
		[]string{cryptoMarkerAttributeName},
	)

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

	// If we got back an attribute, we must verify that it is compatible.
	if len(attrs) > 0 {
		err = verifyCompatible(attrs[0], crypter)
		return
	}

	// Otherwise, we want to claim this domain. Encrypt some random data, base64
	// encode it, then write it out. Make sure to use a precondition to defeat
	// the race condition where another machine is doing the same simultaneously.
	plaintext := get8RandBytes(randSrc)
	ciphertext, err := crypter.Encrypt(plaintext)
	if err != nil {
		err = fmt.Errorf("Encrypt: %v", err)
		return
	}

	encoded := base64.StdEncoding.EncodeToString(ciphertext)

	err = domain.PutAttributes(
		markerItemName,
		[]sdb.PutUpdate{
			sdb.PutUpdate{Name: cryptoMarkerAttributeName, Value: encoded},
		},
		&sdb.Precondition{Name: cryptoMarkerAttributeName, Value: nil},
	)

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

	// All is good.
	return
}

func get8RandBytes(src *rand.Rand) []byte {
	a := src.Uint32()
	b := src.Uint32()

	return []byte{
		byte(a),
		byte(a >> 8),
		byte(a >> 16),
		byte(a >> 24),
		byte(b),
		byte(b >> 8),
		byte(b >> 16),
		byte(b >> 24),
	}
}

type registry struct {
	crypter crypto.Crypter
	db      sdb.SimpleDB
	domain  sdb.Domain
}

func (r *registry) RecordBackup(job CompletedJob) (err error) {
	// Job names must be valid SimpleDB attribute values.
	if len(job.Name) == 0 || len(job.Name) > 1024 || !utf8.ValidString(job.Name) {
		err = fmt.Errorf(
			"Job names must be non-empty UTF-8 no more than 1024 bytes long.")
		return
	}

	// Format the item name.
	itemName := sdb.ItemName(fmt.Sprintf("backup_%016x", job.Id))

	// Use a time format that works correctly with range queries. Make sure to
	// standardize on UTC.
	formattedTime := job.StartTime.UTC().Format(iso8601TimeFormat)

	// Call the domain.
	updates := []sdb.PutUpdate{
		sdb.PutUpdate{Name: "job_name", Value: job.Name},
		sdb.PutUpdate{Name: "start_time", Value: formattedTime},
		sdb.PutUpdate{Name: "score", Value: job.Score.Hex()},
	}

	precond := sdb.Precondition{Name: "score", Value: nil}

	if err = r.domain.PutAttributes(itemName, updates, &precond); err != nil {
		err = fmt.Errorf("PutAttributes: %v", err)
		return
	}

	return
}

func convertAttributes(attrs []sdb.Attribute) (j CompletedJob, err error) {
	for _, attr := range attrs {
		switch attr.Name {
		case "job_name":
			j.Name = attr.Value

		case "start_time":
			if j.StartTime, err = time.Parse(iso8601TimeFormat, attr.Value); err != nil {
				err = fmt.Errorf("Invalid start_time value: %v", err)
				return
			}

		case "score":
			var decoded []byte
			decoded, err = hex.DecodeString(attr.Value)
			if err != nil || len(decoded) != 20 {
				err = fmt.Errorf("Invalid score: %s", attr.Value)
				return
			}

			j.Score = blob.Score(decoded)
		}
	}

	// Everything must have been returned.
	if j.Name == "" {
		err = fmt.Errorf("Missing job_name attribute.")
		return
	}

	if j.StartTime.IsZero() {
		err = fmt.Errorf("Missing start_time attribute.")
		return
	}

	if j.Score == nil {
		err = fmt.Errorf("Missing score attribute.")
		return
	}

	return
}

func convertSelectedItem(item sdb.SelectedItem) (j CompletedJob, err error) {
	// Convert the item's attributes.
	if j, err = convertAttributes(item.Attributes); err != nil {
		return
	}

	// Convert the item name.
	subMatches := itemNameRegexp.FindStringSubmatch(string(item.Name))
	if subMatches == nil {
		err = fmt.Errorf("Invalid item name: %s", item.Name)
		return
	}

	if j.Id, err = strconv.ParseUint(subMatches[1], 16, 64); err != nil {
		panic(fmt.Sprintf("Unexpected result for name: %s", item.Name))
	}

	return
}

func (r *registry) ListRecentBackups() (jobs []CompletedJob, err error) {
	// Call the database.
	query := fmt.Sprintf(
		"select job_name, start_time, score from `%s` where "+
			"start_time is not null order by start_time desc",
		r.domain.Name(),
	)

	results, _, err := r.db.Select(
		query,
		false, // No need for consistent reads.
		nil,
	)

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

	// Convert each result.
	for _, item := range results {
		var job CompletedJob
		job, err = convertSelectedItem(item)
		if err != nil {
			err = fmt.Errorf("Item %s is invalid: %v", item.Name, err)
			return
		}

		jobs = append(jobs, job)
	}

	return
}

func (r *registry) FindBackup(jobId uint64) (job CompletedJob, err error) {
	// Call the domain.
	attrs, err := r.domain.GetAttributes(
		sdb.ItemName(fmt.Sprintf("backup_%016x", jobId)),
		false, // Consistent read unnecessary
		[]string{"job_name", "start_time", "score"},
	)

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

	// Convert the results.
	if job, err = convertAttributes(attrs); err != nil {
		err = fmt.Errorf("Returned attributes invalid: %v", err)
		return
	}

	job.Id = jobId
	return
}
<commit_msg>Use a key deriver.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package registry

import (
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"github.com/jacobsa/aws/sdb"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/crypto"
	"math/rand"
	"regexp"
	"strconv"
	"time"
	"unicode/utf8"
)

const (
	// The item name we use for "domain is already in use" markers.
	markerItemName = "comeback_marker"

	// The attribute name we use for storing crypto-compatibility data.
	cryptoMarkerAttributeName = "encrypted_data"

	// A time format that works properly with range queries.
	iso8601TimeFormat = "2006-01-02T15:04:05Z"
)

// A regexp for selected item names.
var itemNameRegexp = regexp.MustCompile(`^backup_([0-9a-f]{16})$`)

type Registry interface {
	// Record that the named backup job has completed.
	RecordBackup(j CompletedJob) (err error)

	// Return a list of the most recent completed backups.
	ListRecentBackups() (jobs []CompletedJob, err error)

	// Find a particular completed job by ID.
	FindBackup(jobId uint64) (job CompletedJob, err error)
}

func verifyCompatible(
	markerAttr sdb.Attribute,
	crypter crypto.Crypter) (err error) {
	// The ciphertext should be base64-encoded.
	encoded := markerAttr.Value
	ciphertext, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		err = fmt.Errorf("base64.DecodeString(%s): %v", encoded, err)
		return
	}

	// Attempt to decrypt the ciphertext.
	if _, err = crypter.Decrypt(ciphertext); err != nil {
		// Special case: Did the crypter signal that the key was wrong?
		if _, ok := err.(*crypto.NotAuthenticError); ok {
			err = fmt.Errorf("The supplied password is incorrect.")
			return
		}

		// Generic error.
		err = fmt.Errorf("Decrypt: %v", err)
		return
	}

	return
}

// Create a registry that stores data in the supplied SimpleDB domain,
// encrypting using a key derived from the supplied password.
//
// Before doing so, check to see whether this domain has been used as a
// registry before. If not, write an encrypted marker. If it has been used
// before, make sure that it was used with the same password. This prevents
// accidentally writing data with the wrong key when the user enters the wrong
// password.
//
// The crypter must be set up such that it is guaranteed to return an error if
// it is used to decrypt ciphertext encrypted with a different key.
func NewRegistry(
	domain sdb.Domain,
	cryptoPassword string,
	deriver crypto.KeyDeriver,
	randSrc *rand.Rand,
) (r Registry, err error) {
	return newRegistry(
		domain,
		cryptoPassword,
		deriver,
		crypto.NewCrypter,
		randSrc)
}

// A version split out for testability.
func newRegistry(
	domain sdb.Domain,
	cryptoPassword string,
	deriver crypto.KeyDeriver,
	createCrypter func (key []byte) (crypto.Crypter, error),
	randSrc *rand.Rand,
) (r Registry, err error) {
	// Set up a tentative result.
	r = &registry{crypter, domain.Db(), domain}

	// Ask for the data that will tell us whether the crypter is compatible with
	// the previous one used in this domain, if any.
	attrs, err := domain.GetAttributes(
		markerItemName,
		false, // No need to ask for a consistent read
		[]string{cryptoMarkerAttributeName},
	)

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

	// If we got back an attribute, we must verify that it is compatible.
	if len(attrs) > 0 {
		err = verifyCompatible(attrs[0], crypter)
		return
	}

	// Otherwise, we want to claim this domain. Encrypt some random data, base64
	// encode it, then write it out. Make sure to use a precondition to defeat
	// the race condition where another machine is doing the same simultaneously.
	plaintext := get8RandBytes(randSrc)
	ciphertext, err := crypter.Encrypt(plaintext)
	if err != nil {
		err = fmt.Errorf("Encrypt: %v", err)
		return
	}

	encoded := base64.StdEncoding.EncodeToString(ciphertext)

	err = domain.PutAttributes(
		markerItemName,
		[]sdb.PutUpdate{
			sdb.PutUpdate{Name: cryptoMarkerAttributeName, Value: encoded},
		},
		&sdb.Precondition{Name: cryptoMarkerAttributeName, Value: nil},
	)

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

	// All is good.
	return
}

func get8RandBytes(src *rand.Rand) []byte {
	a := src.Uint32()
	b := src.Uint32()

	return []byte{
		byte(a),
		byte(a >> 8),
		byte(a >> 16),
		byte(a >> 24),
		byte(b),
		byte(b >> 8),
		byte(b >> 16),
		byte(b >> 24),
	}
}

type registry struct {
	crypter crypto.Crypter
	db      sdb.SimpleDB
	domain  sdb.Domain
}

func (r *registry) RecordBackup(job CompletedJob) (err error) {
	// Job names must be valid SimpleDB attribute values.
	if len(job.Name) == 0 || len(job.Name) > 1024 || !utf8.ValidString(job.Name) {
		err = fmt.Errorf(
			"Job names must be non-empty UTF-8 no more than 1024 bytes long.")
		return
	}

	// Format the item name.
	itemName := sdb.ItemName(fmt.Sprintf("backup_%016x", job.Id))

	// Use a time format that works correctly with range queries. Make sure to
	// standardize on UTC.
	formattedTime := job.StartTime.UTC().Format(iso8601TimeFormat)

	// Call the domain.
	updates := []sdb.PutUpdate{
		sdb.PutUpdate{Name: "job_name", Value: job.Name},
		sdb.PutUpdate{Name: "start_time", Value: formattedTime},
		sdb.PutUpdate{Name: "score", Value: job.Score.Hex()},
	}

	precond := sdb.Precondition{Name: "score", Value: nil}

	if err = r.domain.PutAttributes(itemName, updates, &precond); err != nil {
		err = fmt.Errorf("PutAttributes: %v", err)
		return
	}

	return
}

func convertAttributes(attrs []sdb.Attribute) (j CompletedJob, err error) {
	for _, attr := range attrs {
		switch attr.Name {
		case "job_name":
			j.Name = attr.Value

		case "start_time":
			if j.StartTime, err = time.Parse(iso8601TimeFormat, attr.Value); err != nil {
				err = fmt.Errorf("Invalid start_time value: %v", err)
				return
			}

		case "score":
			var decoded []byte
			decoded, err = hex.DecodeString(attr.Value)
			if err != nil || len(decoded) != 20 {
				err = fmt.Errorf("Invalid score: %s", attr.Value)
				return
			}

			j.Score = blob.Score(decoded)
		}
	}

	// Everything must have been returned.
	if j.Name == "" {
		err = fmt.Errorf("Missing job_name attribute.")
		return
	}

	if j.StartTime.IsZero() {
		err = fmt.Errorf("Missing start_time attribute.")
		return
	}

	if j.Score == nil {
		err = fmt.Errorf("Missing score attribute.")
		return
	}

	return
}

func convertSelectedItem(item sdb.SelectedItem) (j CompletedJob, err error) {
	// Convert the item's attributes.
	if j, err = convertAttributes(item.Attributes); err != nil {
		return
	}

	// Convert the item name.
	subMatches := itemNameRegexp.FindStringSubmatch(string(item.Name))
	if subMatches == nil {
		err = fmt.Errorf("Invalid item name: %s", item.Name)
		return
	}

	if j.Id, err = strconv.ParseUint(subMatches[1], 16, 64); err != nil {
		panic(fmt.Sprintf("Unexpected result for name: %s", item.Name))
	}

	return
}

func (r *registry) ListRecentBackups() (jobs []CompletedJob, err error) {
	// Call the database.
	query := fmt.Sprintf(
		"select job_name, start_time, score from `%s` where "+
			"start_time is not null order by start_time desc",
		r.domain.Name(),
	)

	results, _, err := r.db.Select(
		query,
		false, // No need for consistent reads.
		nil,
	)

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

	// Convert each result.
	for _, item := range results {
		var job CompletedJob
		job, err = convertSelectedItem(item)
		if err != nil {
			err = fmt.Errorf("Item %s is invalid: %v", item.Name, err)
			return
		}

		jobs = append(jobs, job)
	}

	return
}

func (r *registry) FindBackup(jobId uint64) (job CompletedJob, err error) {
	// Call the domain.
	attrs, err := r.domain.GetAttributes(
		sdb.ItemName(fmt.Sprintf("backup_%016x", jobId)),
		false, // Consistent read unnecessary
		[]string{"job_name", "start_time", "score"},
	)

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

	// Convert the results.
	if job, err = convertAttributes(attrs); err != nil {
		err = fmt.Errorf("Returned attributes invalid: %v", err)
		return
	}

	job.Id = jobId
	return
}
<|endoftext|>


<commit_before>package registry

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

	docker "github.com/fsouza/go-dockerclient"
	"github.com/litl/galaxy/log"
	"github.com/litl/galaxy/utils"
)

/*
All config opbects in redis will be stored in a hash with an id key.
Services will have id, version and environment keys; while Hosts will have id
and location keys.

TODO: IMPORTANT: make an atomic compare-and-swap script to save configs, or
      switch to ORDERED SETS and log changes
*/

const (
	DefaultTTL = 60
)

type ServiceRegistry struct {
	backend      RegistryBackend
	Hostname     string
	TTL          uint64
	OutputBuffer *utils.OutputBuffer
	pollCh       chan bool
	redisHost    string
}

func NewServiceRegistry(ttl uint64) *ServiceRegistry {
	return &ServiceRegistry{
		TTL:    ttl,
		pollCh: make(chan bool),
	}

}

// Build the Redis Pool
func (r *ServiceRegistry) Connect(redisHost string) {

	r.redisHost = redisHost
	r.backend = &RedisBackend{
		RedisHost: redisHost,
	}
	r.backend.Connect()
}

func (r *ServiceRegistry) newServiceRegistration(container *docker.Container, hostIP string) *ServiceRegistration {
	//FIXME: We're using the first found port and assuming it's tcp.
	//How should we handle a service that exposes multiple ports
	//as well as tcp vs udp ports.
	var externalPort, internalPort string
	for k, v := range container.NetworkSettings.Ports {
		if len(v) > 0 {
			externalPort = v[0].HostPort
			internalPort = k.Port()
			break
		}
	}

	serviceRegistration := ServiceRegistration{
		ContainerName: container.Name,
		ContainerID:   container.ID,
		StartedAt:     container.Created,
		Image:         container.Config.Image,
	}

	if externalPort != "" && internalPort != "" {
		serviceRegistration.ExternalIP = hostIP
		serviceRegistration.InternalIP = container.NetworkSettings.IPAddress
		serviceRegistration.ExternalPort = externalPort
		serviceRegistration.InternalPort = internalPort
	}
	return &serviceRegistration
}

type ServiceRegistration struct {
	Name          string            `json:"NAME,omitempty"`
	ExternalIP    string            `json:"EXTERNAL_IP,omitempty"`
	ExternalPort  string            `json:"EXTERNAL_PORT,omitempty"`
	InternalIP    string            `json:"INTERNAL_IP,omitempty"`
	InternalPort  string            `json:"INTERNAL_PORT,omitempty"`
	ContainerID   string            `json:"CONTAINER_ID"`
	ContainerName string            `json:"CONTAINER_NAME"`
	Image         string            `json:"IMAGE,omitempty"`
	ImageId       string            `json:"IMAGE_ID,omitempty"`
	StartedAt     time.Time         `json:"STARTED_AT"`
	Expires       time.Time         `json:"-"`
	Path          string            `json:"-"`
	VirtualHosts  []string          `json:"VIRTUAL_HOSTS"`
	Port          string            `json:"PORT"`
	ErrorPages    map[string]string `json:"ERROR_PAGES,omitempty"`
}

func (s *ServiceRegistration) Equals(other ServiceRegistration) bool {
	return s.ExternalIP == other.ExternalIP &&
		s.ExternalPort == other.ExternalPort &&
		s.InternalIP == other.InternalIP &&
		s.InternalPort == other.InternalPort
}

func (s *ServiceRegistration) addr(ip, port string) string {
	if ip != "" && port != "" {
		return fmt.Sprint(ip, ":", port)
	}
	return ""

}
func (s *ServiceRegistration) ExternalAddr() string {
	return s.addr(s.ExternalIP, s.ExternalPort)
}

func (s *ServiceRegistration) InternalAddr() string {
	return s.addr(s.InternalIP, s.InternalPort)
}

func (r *ServiceRegistry) RegisterService(env, pool, hostIP string, container *docker.Container) (*ServiceRegistration, error) {
	environment := r.EnvFor(container)

	name := environment["GALAXY_APP"]
	if name == "" {
		return nil, fmt.Errorf("GALAXY_APP not set on container %s", container.ID[0:12])
	}

	registrationPath := path.Join(env, pool, "hosts", hostIP, name, container.ID[0:12])

	serviceRegistration := r.newServiceRegistration(container, hostIP)
	serviceRegistration.Name = name
	serviceRegistration.ImageId = container.Config.Image

	vhosts := environment["VIRTUAL_HOST"]
	serviceRegistration.VirtualHosts = strings.Split(vhosts, ",")

	errorPages := make(map[string]string)

	// scan environment variables for the VIRTUAL_HOST_%d pattern
	// but save the original variable and url.
	for vhostCode, url := range environment {
		code := 0
		n, err := fmt.Sscanf(vhostCode, "VIRTUAL_HOST_%d", &code)
		if err != nil || n == 0 {
			continue
		}

		errorPages[vhostCode] = url
	}

	if len(errorPages) > 0 {
		serviceRegistration.ErrorPages = errorPages
	}

	serviceRegistration.VirtualHosts = strings.Split(vhosts, ",")

	serviceRegistration.Port = environment["GALAXY_PORT"]

	jsonReg, err := json.Marshal(serviceRegistration)
	if err != nil {
		return nil, err
	}

	// TODO: use a compare-and-swap SCRIPT
	_, err = r.backend.Set(registrationPath, "location", string(jsonReg))
	if err != nil {
		return nil, err
	}

	_, err = r.backend.Expire(registrationPath, r.TTL)
	if err != nil {
		return nil, err
	}
	serviceRegistration.Expires = time.Now().UTC().Add(time.Duration(r.TTL) * time.Second)

	return serviceRegistration, nil
}

func (r *ServiceRegistry) UnRegisterService(env, pool, hostIP string, container *docker.Container) (*ServiceRegistration, error) {

	environment := r.EnvFor(container)

	name := environment["GALAXY_APP"]
	if name == "" {
		return nil, fmt.Errorf("GALAXY_APP not set on container %s", container.ID[0:12])
	}

	registrationPath := path.Join(env, pool, "hosts", hostIP, name, container.ID[0:12])

	registration, err := r.GetServiceRegistration(env, pool, hostIP, container)
	if err != nil || registration == nil {
		return registration, err
	}

	if registration.ContainerID != container.ID {
		return nil, nil
	}

	_, err = r.backend.Delete(registrationPath)
	if err != nil {
		return registration, err
	}

	return registration, nil
}

func (r *ServiceRegistry) GetServiceRegistration(env, pool, hostIP string, container *docker.Container) (*ServiceRegistration, error) {

	environment := r.EnvFor(container)

	name := environment["GALAXY_APP"]
	if name == "" {
		return nil, fmt.Errorf("GALAXY_APP not set on container %s", container.ID[0:12])
	}

	regPath := path.Join(env, pool, "hosts", hostIP, name, container.ID[0:12])

	existingRegistration := ServiceRegistration{
		Path: regPath,
	}

	location, err := r.backend.Get(regPath, "location")

	if err != nil {
		return nil, err
	}

	if location != "" {
		err = json.Unmarshal([]byte(location), &existingRegistration)
		if err != nil {
			return nil, err
		}

		expires, err := r.backend.Ttl(regPath)
		if err != nil {
			return nil, err
		}
		existingRegistration.Expires = time.Now().UTC().Add(time.Duration(expires) * time.Second)
		return &existingRegistration, nil
	}

	return nil, nil
}

func (r *ServiceRegistry) IsRegistered(env, pool, hostIP string, container *docker.Container) (bool, error) {

	reg, err := r.GetServiceRegistration(env, pool, hostIP, container)
	return reg != nil, err
}

// TODO: get all ServiceRegistrations
func (r *ServiceRegistry) ListRegistrations(env string) ([]ServiceRegistration, error) {

	// TODO: convert to scan
	keys, err := r.backend.Keys(path.Join(env, "*", "hosts", "*", "*", "*"))
	if err != nil {
		return nil, err
	}

	var regList []ServiceRegistration
	for _, key := range keys {

		val, err := r.backend.Get(key, "location")
		if err != nil {
			log.Warnf("WARN: Unable to get location for %s: %s", key, err)
			continue
		}

		svcReg := ServiceRegistration{
			Name: path.Base(key),
		}
		err = json.Unmarshal([]byte(val), &svcReg)
		if err != nil {
			log.Warnf("WARN: Unable to unmarshal JSON for %s: %s", key, err)
			continue
		}

		svcReg.Path = key

		regList = append(regList, svcReg)
	}

	return regList, nil
}

func (s *ServiceRegistry) EnvFor(container *docker.Container) map[string]string {
	env := map[string]string{}
	for _, item := range container.Config.Env {
		sep := strings.Index(item, "=")
		k := item[0:sep]
		v := item[sep+1:]
		env[k] = v
	}
	return env
}
<commit_msg>commander: don't send [""] to shuttle for virtual hosts value<commit_after>package registry

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

	docker "github.com/fsouza/go-dockerclient"
	"github.com/litl/galaxy/log"
	"github.com/litl/galaxy/utils"
)

/*
All config opbects in redis will be stored in a hash with an id key.
Services will have id, version and environment keys; while Hosts will have id
and location keys.

TODO: IMPORTANT: make an atomic compare-and-swap script to save configs, or
      switch to ORDERED SETS and log changes
*/

const (
	DefaultTTL = 60
)

type ServiceRegistry struct {
	backend      RegistryBackend
	Hostname     string
	TTL          uint64
	OutputBuffer *utils.OutputBuffer
	pollCh       chan bool
	redisHost    string
}

func NewServiceRegistry(ttl uint64) *ServiceRegistry {
	return &ServiceRegistry{
		TTL:    ttl,
		pollCh: make(chan bool),
	}

}

// Build the Redis Pool
func (r *ServiceRegistry) Connect(redisHost string) {

	r.redisHost = redisHost
	r.backend = &RedisBackend{
		RedisHost: redisHost,
	}
	r.backend.Connect()
}

func (r *ServiceRegistry) newServiceRegistration(container *docker.Container, hostIP string) *ServiceRegistration {
	//FIXME: We're using the first found port and assuming it's tcp.
	//How should we handle a service that exposes multiple ports
	//as well as tcp vs udp ports.
	var externalPort, internalPort string
	for k, v := range container.NetworkSettings.Ports {
		if len(v) > 0 {
			externalPort = v[0].HostPort
			internalPort = k.Port()
			break
		}
	}

	serviceRegistration := ServiceRegistration{
		ContainerName: container.Name,
		ContainerID:   container.ID,
		StartedAt:     container.Created,
		Image:         container.Config.Image,
	}

	if externalPort != "" && internalPort != "" {
		serviceRegistration.ExternalIP = hostIP
		serviceRegistration.InternalIP = container.NetworkSettings.IPAddress
		serviceRegistration.ExternalPort = externalPort
		serviceRegistration.InternalPort = internalPort
	}
	return &serviceRegistration
}

type ServiceRegistration struct {
	Name          string            `json:"NAME,omitempty"`
	ExternalIP    string            `json:"EXTERNAL_IP,omitempty"`
	ExternalPort  string            `json:"EXTERNAL_PORT,omitempty"`
	InternalIP    string            `json:"INTERNAL_IP,omitempty"`
	InternalPort  string            `json:"INTERNAL_PORT,omitempty"`
	ContainerID   string            `json:"CONTAINER_ID"`
	ContainerName string            `json:"CONTAINER_NAME"`
	Image         string            `json:"IMAGE,omitempty"`
	ImageId       string            `json:"IMAGE_ID,omitempty"`
	StartedAt     time.Time         `json:"STARTED_AT"`
	Expires       time.Time         `json:"-"`
	Path          string            `json:"-"`
	VirtualHosts  []string          `json:"VIRTUAL_HOSTS"`
	Port          string            `json:"PORT"`
	ErrorPages    map[string]string `json:"ERROR_PAGES,omitempty"`
}

func (s *ServiceRegistration) Equals(other ServiceRegistration) bool {
	return s.ExternalIP == other.ExternalIP &&
		s.ExternalPort == other.ExternalPort &&
		s.InternalIP == other.InternalIP &&
		s.InternalPort == other.InternalPort
}

func (s *ServiceRegistration) addr(ip, port string) string {
	if ip != "" && port != "" {
		return fmt.Sprint(ip, ":", port)
	}
	return ""

}
func (s *ServiceRegistration) ExternalAddr() string {
	return s.addr(s.ExternalIP, s.ExternalPort)
}

func (s *ServiceRegistration) InternalAddr() string {
	return s.addr(s.InternalIP, s.InternalPort)
}

func (r *ServiceRegistry) RegisterService(env, pool, hostIP string, container *docker.Container) (*ServiceRegistration, error) {
	environment := r.EnvFor(container)

	name := environment["GALAXY_APP"]
	if name == "" {
		return nil, fmt.Errorf("GALAXY_APP not set on container %s", container.ID[0:12])
	}

	registrationPath := path.Join(env, pool, "hosts", hostIP, name, container.ID[0:12])

	serviceRegistration := r.newServiceRegistration(container, hostIP)
	serviceRegistration.Name = name
	serviceRegistration.ImageId = container.Config.Image

	vhosts := environment["VIRTUAL_HOST"]
	if strings.TrimSpace(vhosts) != "" {
		serviceRegistration.VirtualHosts = strings.Split(vhosts, ",")
	}

	errorPages := make(map[string]string)

	// scan environment variables for the VIRTUAL_HOST_%d pattern
	// but save the original variable and url.
	for vhostCode, url := range environment {
		code := 0
		n, err := fmt.Sscanf(vhostCode, "VIRTUAL_HOST_%d", &code)
		if err != nil || n == 0 {
			continue
		}

		errorPages[vhostCode] = url
	}

	if len(errorPages) > 0 {
		serviceRegistration.ErrorPages = errorPages
	}

	serviceRegistration.Port = environment["GALAXY_PORT"]

	jsonReg, err := json.Marshal(serviceRegistration)
	if err != nil {
		return nil, err
	}

	// TODO: use a compare-and-swap SCRIPT
	_, err = r.backend.Set(registrationPath, "location", string(jsonReg))
	if err != nil {
		return nil, err
	}

	_, err = r.backend.Expire(registrationPath, r.TTL)
	if err != nil {
		return nil, err
	}
	serviceRegistration.Expires = time.Now().UTC().Add(time.Duration(r.TTL) * time.Second)

	return serviceRegistration, nil
}

func (r *ServiceRegistry) UnRegisterService(env, pool, hostIP string, container *docker.Container) (*ServiceRegistration, error) {

	environment := r.EnvFor(container)

	name := environment["GALAXY_APP"]
	if name == "" {
		return nil, fmt.Errorf("GALAXY_APP not set on container %s", container.ID[0:12])
	}

	registrationPath := path.Join(env, pool, "hosts", hostIP, name, container.ID[0:12])

	registration, err := r.GetServiceRegistration(env, pool, hostIP, container)
	if err != nil || registration == nil {
		return registration, err
	}

	if registration.ContainerID != container.ID {
		return nil, nil
	}

	_, err = r.backend.Delete(registrationPath)
	if err != nil {
		return registration, err
	}

	return registration, nil
}

func (r *ServiceRegistry) GetServiceRegistration(env, pool, hostIP string, container *docker.Container) (*ServiceRegistration, error) {

	environment := r.EnvFor(container)

	name := environment["GALAXY_APP"]
	if name == "" {
		return nil, fmt.Errorf("GALAXY_APP not set on container %s", container.ID[0:12])
	}

	regPath := path.Join(env, pool, "hosts", hostIP, name, container.ID[0:12])

	existingRegistration := ServiceRegistration{
		Path: regPath,
	}

	location, err := r.backend.Get(regPath, "location")

	if err != nil {
		return nil, err
	}

	if location != "" {
		err = json.Unmarshal([]byte(location), &existingRegistration)
		if err != nil {
			return nil, err
		}

		expires, err := r.backend.Ttl(regPath)
		if err != nil {
			return nil, err
		}
		existingRegistration.Expires = time.Now().UTC().Add(time.Duration(expires) * time.Second)
		return &existingRegistration, nil
	}

	return nil, nil
}

func (r *ServiceRegistry) IsRegistered(env, pool, hostIP string, container *docker.Container) (bool, error) {

	reg, err := r.GetServiceRegistration(env, pool, hostIP, container)
	return reg != nil, err
}

// TODO: get all ServiceRegistrations
func (r *ServiceRegistry) ListRegistrations(env string) ([]ServiceRegistration, error) {

	// TODO: convert to scan
	keys, err := r.backend.Keys(path.Join(env, "*", "hosts", "*", "*", "*"))
	if err != nil {
		return nil, err
	}

	var regList []ServiceRegistration
	for _, key := range keys {

		val, err := r.backend.Get(key, "location")
		if err != nil {
			log.Warnf("WARN: Unable to get location for %s: %s", key, err)
			continue
		}

		svcReg := ServiceRegistration{
			Name: path.Base(key),
		}
		err = json.Unmarshal([]byte(val), &svcReg)
		if err != nil {
			log.Warnf("WARN: Unable to unmarshal JSON for %s: %s", key, err)
			continue
		}

		svcReg.Path = key

		regList = append(regList, svcReg)
	}

	return regList, nil
}

func (s *ServiceRegistry) EnvFor(container *docker.Container) map[string]string {
	env := map[string]string{}
	for _, item := range container.Config.Env {
		sep := strings.Index(item, "=")
		k := item[0:sep]
		v := item[sep+1:]
		env[k] = v
	}
	return env
}
<|endoftext|>


<commit_before>/*
   Copyright (c) 2014, Percona LLC and/or its affiliates. All rights reserved.

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>
*/

package pct

import (
	"fmt"
	"github.com/percona/cloud-protocol/proto"
	"time"
)

type Logger struct {
	logChan chan *proto.LogEntry
	service string
	cmd     *proto.Cmd
}

func NewLogger(logChan chan *proto.LogEntry, service string) *Logger {
	l := &Logger{
		logChan: logChan,
		service: service,
	}
	return l
}

func (l *Logger) LogChan() chan *proto.LogEntry {
	return l.logChan
}

func (l *Logger) InResponseTo(cmd *proto.Cmd) {
	l.cmd = cmd
}

func (l *Logger) Debug(entry ...interface{}) {
	l.log(proto.LOG_DEBUG, entry)
}

func (l *Logger) Info(entry ...interface{}) {
	l.log(proto.LOG_INFO, entry)
}

func (l *Logger) Warn(entry ...interface{}) {
	l.log(proto.LOG_WARNING, entry)
}

func (l *Logger) Error(entry ...interface{}) {
	l.log(proto.LOG_ERROR, entry)
}

func (l *Logger) Fatal(entry ...interface{}) {
	l.log(proto.LOG_CRITICAL, entry)
}

func (l *Logger) log(level byte, entry ...interface{}) {
	fullMsg := ""
	for i, str := range entry {
		if i > 0 {
			fullMsg += " "
		}
		fullMsg += fmt.Sprintf("%v", str)
	}
	logEntry := &proto.LogEntry{
		Ts:      time.Now().UTC(),
		Level:   level,
		Service: l.service,
		Msg:     fullMsg,
	}
	if l.cmd != nil {
		logEntry.Cmd = fmt.Sprintf("%s", l.cmd)
	}
	select {
	case l.logChan <- logEntry:
	default:
		// todo: lot.Println()?
		// This happens when LogRelay.LogChan() is full, which means the log relay
		// is receiving log entries faster than it can buffer and send them.
	}
}
<commit_msg>Fix logger.<commit_after>/*
   Copyright (c) 2014, Percona LLC and/or its affiliates. All rights reserved.

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>
*/

package pct

import (
	"fmt"
	"github.com/percona/cloud-protocol/proto"
	"time"
)

type Logger struct {
	logChan chan *proto.LogEntry
	service string
	cmd     *proto.Cmd
}

func NewLogger(logChan chan *proto.LogEntry, service string) *Logger {
	l := &Logger{
		logChan: logChan,
		service: service,
	}
	return l
}

func (l *Logger) LogChan() chan *proto.LogEntry {
	return l.logChan
}

func (l *Logger) InResponseTo(cmd *proto.Cmd) {
	l.cmd = cmd
}

func (l *Logger) Debug(entry ...interface{}) {
	l.log(proto.LOG_DEBUG, entry)
}

func (l *Logger) Info(entry ...interface{}) {
	l.log(proto.LOG_INFO, entry)
}

func (l *Logger) Warn(entry ...interface{}) {
	l.log(proto.LOG_WARNING, entry)
}

func (l *Logger) Error(entry ...interface{}) {
	l.log(proto.LOG_ERROR, entry)
}

func (l *Logger) Fatal(entry ...interface{}) {
	l.log(proto.LOG_CRITICAL, entry)
}

func (l *Logger) log(level byte, entry []interface{}) {
	fullMsg := ""
	for i, str := range entry {
		if i > 0 {
			fullMsg += " "
		}
		fullMsg += fmt.Sprintf("%v", str)
	}
	logEntry := &proto.LogEntry{
		Ts:      time.Now().UTC(),
		Level:   level,
		Service: l.service,
		Msg:     fullMsg,
	}
	if l.cmd != nil {
		logEntry.Cmd = fmt.Sprintf("%s", l.cmd)
	}
	select {
	case l.logChan <- logEntry:
	default:
		// todo: lot.Println()?
		// This happens when LogRelay.LogChan() is full, which means the log relay
		// is receiving log entries faster than it can buffer and send them.
	}
}
<|endoftext|>


<commit_before>package seccomp

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/subgraph/oz"
	"golang.org/x/sys/unix"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"strconv"
	"syscall"
	"unsafe"
)

const (
	STRINGARG = iota + 1
	PTRARG
	INTARG
)

const MAXARGS int = 6
const MAXSTRING int = 100

type SystemCallArgs [MAXARGS]int

type SystemCall struct {
	prefix string
	name   string
	num    int
	args   SystemCallArgs
}

func Tracer() {

	p := new(oz.Profile)
	if err := json.NewDecoder(os.Stdin).Decode(&p); err != nil {
		log.Error("unable to decode profile data: %v", err)
		os.Exit(1)
	}

	log := createLogger()

	var proc_attr syscall.ProcAttr
	var sys_attr syscall.SysProcAttr

	sys_attr.Ptrace = true
	done := false
	proc_attr.Sys = &sys_attr

	cmd := os.Args[1]
	cmdArgs := os.Args[2:]
	log.Info("Tracer running command (%v) arguments (%v)\n", cmd, cmdArgs)
	c := exec.Command(cmd)
	c.SysProcAttr = &syscall.SysProcAttr{Ptrace: true}
	c.Env = os.Environ()
	c.Args = append(c.Args, cmdArgs...)

	pi, err := c.StdinPipe()
	if err != nil {
		fmt.Errorf("error creating stdin pipe for tracer process: %v", err)
		os.Exit(1)
	}
	jdata, err := json.Marshal(p)
	if err != nil {
		fmt.Errorf("Unable to marshal seccomp state: %+v", err)
		os.Exit(1)
	}
	io.Copy(pi, bytes.NewBuffer(jdata))
	log.Info(string(jdata))
	pi.Close()

	children := make(map[int]bool)

	if err := c.Start(); err == nil {
		children[c.Process.Pid] = true
		var s syscall.WaitStatus
		pid, err := syscall.Wait4(-1, &s, syscall.WALL, nil)
		children[pid] = true
		if err != nil {
			log.Error("Error (wait4): %v", err)
		}
		log.Info("Tracing child pid: %v\n", pid)
		for done == false {
			syscall.PtraceSetOptions(pid, unix.PTRACE_O_TRACESECCOMP|unix.PTRACE_O_TRACEFORK|unix.PTRACE_O_TRACEVFORK|unix.PTRACE_O_TRACECLONE|unix.PTRACE_O_TRACEEXIT)
			syscall.PtraceCont(pid, 0)
			pid, err = syscall.Wait4(-1, &s, syscall.WALL, nil)
			if err != nil {
				log.Error("Error (wait4): %v\n", err)
				if len(children) == 0 {
					done = true
				}
				continue
			}
			children[pid] = true
			if s.Exited() == true {
				delete(children, pid)
				log.Info("Child pid %v finished.\n", pid)
				if len(children) == 0 {
					done = true
				}
				continue
			}
			if uint32(s)>>8 == (uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_SECCOMP << 8)) {
				if err != nil {
					log.Error("Error (ptrace): %v", err)
					continue
				}
				var regs syscall.PtraceRegs
				err = syscall.PtraceGetRegs(pid, &regs)
				if err != nil {
					log.Error("Error (ptrace): %v", err)
				}
				systemcall, err := syscallByNum(int(regs.Orig_rax))
				if err != nil {
					log.Error("Error: %v", err)
					continue
				}
				var callrep string = fmt.Sprintf("%s(", systemcall.name)
				var reg uint64 = 0

				for arg := range systemcall.args {

					if systemcall.args[arg] == 0 {
						break
					}

					if arg > 0 {
						callrep += fmt.Sprintf(",")
					}

					switch arg {
					case 0:
						reg = regs.Rdi
					case 1:
						reg = regs.Rsi
					case 2:
						reg = regs.Rdx
					case 3:
						reg = regs.Rcx
					case 4:
						reg = regs.R8
					case 5:
						reg = regs.R9
					}
					if systemcall.args[arg] == STRINGARG {
						str, err := readStringArg(pid, uintptr(reg))
						if err != nil {
							log.Error("Error: %v", err)
						} else {
							callrep += fmt.Sprintf("\"%s\"", str)
						}
					} else if systemcall.args[arg] == INTARG {
						callrep += fmt.Sprintf("%d", uint64(reg))
					} else {
						/* Stringify pointers in writes to stdout/stderr */
						write, err := syscallByName("write")
						if err != nil {
							log.Error("Error: %v", err)
						}
						if systemcall.num == write.num && (regs.Rdi == uint64(syscall.Stdout) || regs.Rdi == uint64(syscall.Stderr)) {
							str, err := readStringArg(pid, uintptr(reg))
							if err != nil {
								log.Error("Error %v", err)
							} else {
								if isPrintableASCII(str) == true {
									callrep += fmt.Sprintf("\"%s\"", str)
								} else {
									callrep += fmt.Sprintf("0x%X", uintptr(reg))
								}
							}
						} else {
							callrep += fmt.Sprintf("0x%X", uintptr(reg))
						}
					}

				}
				callrep += ")"
				log.Info("==============================================\nseccomp hit on sandbox pid %v (%v) syscall %v (%v): \n\n%s\nI ==============================================\n\n", pid, getProcessCmdLine(pid), systemcall.name, regs.Orig_rax, callrep)
			}
		}
	} else {
		log.Error("Error: %v", err)
	}

}

func readStringArg(pid int, addr uintptr) (s string, err error) {
	buf := make([]byte, unsafe.Sizeof(addr))
	done := false
	err = nil
	for done == false {
		_, err := syscall.PtracePeekText(pid, addr, buf)
		if err != nil {
			fmt.Printf("Error (ptrace): %v\n", err)
		} else {
			for b := range buf {
				if buf[b] == 0 {
					done = true
					break
				} else {
					s += string(buf[b])
					if len(s) > 90 {
						s += "..."
						break
					}
				}
			}
		}
		addr += unsafe.Sizeof(addr)
	}
	return s, nil
}

func readUIntArg(pid int, addr uintptr) (uint64, error) {
	buf := make([]byte, unsafe.Sizeof(addr))

	_, err := syscall.PtracePeekText(pid, addr, buf)
	if err != nil {
		fmt.Printf("Error (ptrace): %v\n", err)
		return 0, err
	} else {
		i := binary.LittleEndian.Uint64(buf)
		return i, nil
	}
	return 0, errors.New("Error.")
}

func readPtrArg(pid int, addr uintptr) (uintptr, error) {

	buf := make([]byte, unsafe.Sizeof(addr))

	_, err := syscall.PtracePeekText(pid, addr, buf)
	if err != nil {
		fmt.Printf("Error (ptrace): %v\n", err)
		return 0, err
	} else {
		i := binary.LittleEndian.Uint64(buf)
		return uintptr(i), nil
	}
	return 0, nil
}

func syscallByNum(num int) (s SystemCall, err error) {
	var q SystemCall = SystemCall{"", "", -1, [6]int{0, 0, 0, 0, 0, 0}}
	for i := range syscalls {
		if syscalls[i].num == num {
			q = syscalls[i]
			return q, nil
		}
	}
	return q, errors.New("System call not found.\n")
}

func syscallByName(name string) (s SystemCall, err error) {
	var q SystemCall = SystemCall{"", "", -1, [6]int{0, 0, 0, 0, 0, 0}}
	for i := range syscalls {
		if syscalls[i].name == name {
			q = syscalls[i]
			return q, nil
		}
	}
	return q, errors.New("System call not found.\n")
}

func isPrintableASCII(s string) bool {
	for _, x := range s {
		if x > 127 {
			return false
		}
	}
	return true
}

func getProcessCmdLine(pid int) string {
	path := "/proc/" + strconv.Itoa(pid) + "/cmdline"
	cmdline, err := ioutil.ReadFile(path)
	for b := range cmdline {
		if b < (len(cmdline) - 1) {
			if cmdline[b] == 0x00 {
				cmdline[b] = 0x20
			}
		}
	}
	if err != nil {
		log.Error("Error (read): %v", err)
		return "unknown"
	}
	return string(cmdline)
}
<commit_msg>Removed, not necessary.<commit_after>package seccomp

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/subgraph/oz"
	"golang.org/x/sys/unix"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"strconv"
	"syscall"
	"unsafe"
)

const (
	STRINGARG = iota + 1
	PTRARG
	INTARG
)

const MAXARGS int = 6
const MAXSTRING int = 100

type SystemCallArgs [MAXARGS]int

type SystemCall struct {
	prefix string
	name   string
	num    int
	args   SystemCallArgs
}

func Tracer() {

	p := new(oz.Profile)
	if err := json.NewDecoder(os.Stdin).Decode(&p); err != nil {
		log.Error("unable to decode profile data: %v", err)
		os.Exit(1)
	}

	var proc_attr syscall.ProcAttr
	var sys_attr syscall.SysProcAttr

	sys_attr.Ptrace = true
	done := false
	proc_attr.Sys = &sys_attr

	cmd := os.Args[1]
	cmdArgs := os.Args[2:]
	log.Info("Tracer running command (%v) arguments (%v)\n", cmd, cmdArgs)
	c := exec.Command(cmd)
	c.SysProcAttr = &syscall.SysProcAttr{Ptrace: true}
	c.Env = os.Environ()
	c.Args = append(c.Args, cmdArgs...)

	pi, err := c.StdinPipe()
	if err != nil {
		fmt.Errorf("error creating stdin pipe for tracer process: %v", err)
		os.Exit(1)
	}
	jdata, err := json.Marshal(p)
	if err != nil {
		fmt.Errorf("Unable to marshal seccomp state: %+v", err)
		os.Exit(1)
	}
	io.Copy(pi, bytes.NewBuffer(jdata))
	log.Info(string(jdata))
	pi.Close()

	children := make(map[int]bool)

	if err := c.Start(); err == nil {
		children[c.Process.Pid] = true
		var s syscall.WaitStatus
		pid, err := syscall.Wait4(-1, &s, syscall.WALL, nil)
		children[pid] = true
		if err != nil {
			log.Error("Error (wait4): %v", err)
		}
		log.Info("Tracing child pid: %v\n", pid)
		for done == false {
			syscall.PtraceSetOptions(pid, unix.PTRACE_O_TRACESECCOMP|unix.PTRACE_O_TRACEFORK|unix.PTRACE_O_TRACEVFORK|unix.PTRACE_O_TRACECLONE|unix.PTRACE_O_TRACEEXIT)
			syscall.PtraceCont(pid, 0)
			pid, err = syscall.Wait4(-1, &s, syscall.WALL, nil)
			if err != nil {
				log.Error("Error (wait4): %v\n", err)
				if len(children) == 0 {
					done = true
				}
				continue
			}
			children[pid] = true
			if s.Exited() == true {
				delete(children, pid)
				log.Info("Child pid %v finished.\n", pid)
				if len(children) == 0 {
					done = true
				}
				continue
			}
			if uint32(s)>>8 == (uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_SECCOMP << 8)) {
				if err != nil {
					log.Error("Error (ptrace): %v", err)
					continue
				}
				var regs syscall.PtraceRegs
				err = syscall.PtraceGetRegs(pid, &regs)
				if err != nil {
					log.Error("Error (ptrace): %v", err)
				}
				systemcall, err := syscallByNum(int(regs.Orig_rax))
				if err != nil {
					log.Error("Error: %v", err)
					continue
				}
				var callrep string = fmt.Sprintf("%s(", systemcall.name)
				var reg uint64 = 0

				for arg := range systemcall.args {

					if systemcall.args[arg] == 0 {
						break
					}

					if arg > 0 {
						callrep += fmt.Sprintf(",")
					}

					switch arg {
					case 0:
						reg = regs.Rdi
					case 1:
						reg = regs.Rsi
					case 2:
						reg = regs.Rdx
					case 3:
						reg = regs.Rcx
					case 4:
						reg = regs.R8
					case 5:
						reg = regs.R9
					}
					if systemcall.args[arg] == STRINGARG {
						str, err := readStringArg(pid, uintptr(reg))
						if err != nil {
							log.Error("Error: %v", err)
						} else {
							callrep += fmt.Sprintf("\"%s\"", str)
						}
					} else if systemcall.args[arg] == INTARG {
						callrep += fmt.Sprintf("%d", uint64(reg))
					} else {
						/* Stringify pointers in writes to stdout/stderr */
						write, err := syscallByName("write")
						if err != nil {
							log.Error("Error: %v", err)
						}
						if systemcall.num == write.num && (regs.Rdi == uint64(syscall.Stdout) || regs.Rdi == uint64(syscall.Stderr)) {
							str, err := readStringArg(pid, uintptr(reg))
							if err != nil {
								log.Error("Error %v", err)
							} else {
								if isPrintableASCII(str) == true {
									callrep += fmt.Sprintf("\"%s\"", str)
								} else {
									callrep += fmt.Sprintf("0x%X", uintptr(reg))
								}
							}
						} else {
							callrep += fmt.Sprintf("0x%X", uintptr(reg))
						}
					}

				}
				callrep += ")"
				log.Info("==============================================\nseccomp hit on sandbox pid %v (%v) syscall %v (%v): \n\n%s\nI ==============================================\n\n", pid, getProcessCmdLine(pid), systemcall.name, regs.Orig_rax, callrep)
			}
		}
	} else {
		log.Error("Error: %v", err)
	}

}

func readStringArg(pid int, addr uintptr) (s string, err error) {
	buf := make([]byte, unsafe.Sizeof(addr))
	done := false
	err = nil
	for done == false {
		_, err := syscall.PtracePeekText(pid, addr, buf)
		if err != nil {
			fmt.Printf("Error (ptrace): %v\n", err)
		} else {
			for b := range buf {
				if buf[b] == 0 {
					done = true
					break
				} else {
					s += string(buf[b])
					if len(s) > 90 {
						s += "..."
						break
					}
				}
			}
		}
		addr += unsafe.Sizeof(addr)
	}
	return s, nil
}

func readUIntArg(pid int, addr uintptr) (uint64, error) {
	buf := make([]byte, unsafe.Sizeof(addr))

	_, err := syscall.PtracePeekText(pid, addr, buf)
	if err != nil {
		fmt.Printf("Error (ptrace): %v\n", err)
		return 0, err
	} else {
		i := binary.LittleEndian.Uint64(buf)
		return i, nil
	}
	return 0, errors.New("Error.")
}

func readPtrArg(pid int, addr uintptr) (uintptr, error) {

	buf := make([]byte, unsafe.Sizeof(addr))

	_, err := syscall.PtracePeekText(pid, addr, buf)
	if err != nil {
		fmt.Printf("Error (ptrace): %v\n", err)
		return 0, err
	} else {
		i := binary.LittleEndian.Uint64(buf)
		return uintptr(i), nil
	}
	return 0, nil
}

func syscallByNum(num int) (s SystemCall, err error) {
	var q SystemCall = SystemCall{"", "", -1, [6]int{0, 0, 0, 0, 0, 0}}
	for i := range syscalls {
		if syscalls[i].num == num {
			q = syscalls[i]
			return q, nil
		}
	}
	return q, errors.New("System call not found.\n")
}

func syscallByName(name string) (s SystemCall, err error) {
	var q SystemCall = SystemCall{"", "", -1, [6]int{0, 0, 0, 0, 0, 0}}
	for i := range syscalls {
		if syscalls[i].name == name {
			q = syscalls[i]
			return q, nil
		}
	}
	return q, errors.New("System call not found.\n")
}

func isPrintableASCII(s string) bool {
	for _, x := range s {
		if x > 127 {
			return false
		}
	}
	return true
}

func getProcessCmdLine(pid int) string {
	path := "/proc/" + strconv.Itoa(pid) + "/cmdline"
	cmdline, err := ioutil.ReadFile(path)
	for b := range cmdline {
		if b < (len(cmdline) - 1) {
			if cmdline[b] == 0x00 {
				cmdline[b] = 0x20
			}
		}
	}
	if err != nil {
		log.Error("Error (read): %v", err)
		return "unknown"
	}
	return string(cmdline)
}
<|endoftext|>


<commit_before>package storage

import (
	"io"

	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/migration"
	"github.com/lxc/lxd/lxd/operations"
	"github.com/lxc/lxd/lxd/state"
	"github.com/lxc/lxd/lxd/storage/drivers"
	"github.com/lxc/lxd/shared/api"
)

// Instance represents the storage relevant subset of a LXD instance.
type Instance interface {
	Name() string
	Project() string
	Type() instancetype.Type
	Path() string

	IsRunning() bool
	IsSnapshot() bool
	DeferTemplateApply(trigger string) error
}

// Pool represents a LXD storage pool.
type Pool interface {
	// Internal.
	DaemonState() *state.State

	// Pool.
	ID() int64
	Name() string
	Driver() drivers.Driver

	GetResources() (*api.ResourcesStoragePool, error)
	Delete(op *operations.Operation) error

	Mount() (bool, error)
	Unmount() (bool, error)

	// Instances.
	CreateInstance(i Instance, op *operations.Operation) error
	CreateInstanceFromBackup(i Instance, sourcePath string, op *operations.Operation) error
	CreateInstanceFromCopy(i Instance, src Instance, snapshots bool, op *operations.Operation) error
	CreateInstanceFromImage(i Instance, fingerprint string, op *operations.Operation) error
	CreateInstanceFromMigration(i Instance, conn io.ReadWriteCloser, args migration.SinkArgs, op *operations.Operation) error
	RenameInstance(i Instance, newName string, op *operations.Operation) error
	DeleteInstance(i Instance, op *operations.Operation) error

	MigrateInstance(i Instance, snapshots bool, args migration.SourceArgs) (migration.StorageSourceDriver, error)
	RefreshInstance(i Instance, src Instance, snapshots bool, op *operations.Operation) error
	BackupInstance(i Instance, targetPath string, optimized bool, snapshots bool, op *operations.Operation) error

	GetInstanceUsage(i Instance) (int64, error)
	SetInstanceQuota(i Instance, quota uint64) error

	MountInstance(i Instance, op *operations.Operation) (bool, error)
	UnmountInstance(i Instance, op *operations.Operation) (bool, error)
	GetInstanceDisk(i Instance) (string, string, error)

	// Instance snapshots.
	CreateInstanceSnapshot(i Instance, name string, op *operations.Operation) error
	RenameInstanceSnapshot(i Instance, newName string, op *operations.Operation) error
	DeleteInstanceSnapshot(i Instance, op *operations.Operation) error
	RestoreInstanceSnapshot(i Instance, op *operations.Operation) error
	MountInstanceSnapshot(i Instance, op *operations.Operation) (bool, error)
	UnmountInstanceSnapshot(i Instance, op *operations.Operation) (bool, error)

	// Images.
	EnsureImage(fingerprint string, op *operations.Operation) error
	DeleteImage(fingerprint string, op *operations.Operation) error

	// Custom volumes.
	CreateCustomVolume(volName, desc string, config map[string]string, op *operations.Operation) error
	CreateCustomVolumeFromCopy(volName, desc string, config map[string]string, srcPoolName, srcVolName string, srcVolOnly bool, op *operations.Operation) error
	UpdateCustomVolume(volName, newDesc string, newConfig map[string]string, op *operations.Operation) error
	RenameCustomVolume(volName string, newVolName string, op *operations.Operation) error
	DeleteCustomVolume(volName string, op *operations.Operation) error
	GetCustomVolumeUsage(volName string) (int64, error)
	MountCustomVolume(volName string, op *operations.Operation) (bool, error)
	UnmountCustomVolume(volName string, op *operations.Operation) (bool, error)

	// Custom volume snapshots.
	CreateCustomVolumeSnapshot(volName string, newSnapshotName string, op *operations.Operation) error
	RenameCustomVolumeSnapshot(volName string, newSnapshotName string, op *operations.Operation) error
	DeleteCustomVolumeSnapshot(volName string, op *operations.Operation) error
	RestoreCustomVolume(volName string, snapshotName string, op *operations.Operation) error

	// Custom volume migration.
	MigrationTypes(contentType drivers.ContentType) []migration.Type
	CreateCustomVolumeFromMigration(conn io.ReadWriteCloser, args migration.VolumeTargetArgs, op *operations.Operation) error
	MigrateCustomVolume(conn io.ReadWriteCloser, args migration.VolumeSourceArgs, op *operations.Operation) error
}
<commit_msg>lxd/storage/interfaces: SetInstanceQuota signature<commit_after>package storage

import (
	"io"

	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/migration"
	"github.com/lxc/lxd/lxd/operations"
	"github.com/lxc/lxd/lxd/state"
	"github.com/lxc/lxd/lxd/storage/drivers"
	"github.com/lxc/lxd/shared/api"
)

// Instance represents the storage relevant subset of a LXD instance.
type Instance interface {
	Name() string
	Project() string
	Type() instancetype.Type
	Path() string

	IsRunning() bool
	IsSnapshot() bool
	DeferTemplateApply(trigger string) error
}

// Pool represents a LXD storage pool.
type Pool interface {
	// Internal.
	DaemonState() *state.State

	// Pool.
	ID() int64
	Name() string
	Driver() drivers.Driver

	GetResources() (*api.ResourcesStoragePool, error)
	Delete(op *operations.Operation) error

	Mount() (bool, error)
	Unmount() (bool, error)

	// Instances.
	CreateInstance(i Instance, op *operations.Operation) error
	CreateInstanceFromBackup(i Instance, sourcePath string, op *operations.Operation) error
	CreateInstanceFromCopy(i Instance, src Instance, snapshots bool, op *operations.Operation) error
	CreateInstanceFromImage(i Instance, fingerprint string, op *operations.Operation) error
	CreateInstanceFromMigration(i Instance, conn io.ReadWriteCloser, args migration.SinkArgs, op *operations.Operation) error
	RenameInstance(i Instance, newName string, op *operations.Operation) error
	DeleteInstance(i Instance, op *operations.Operation) error

	MigrateInstance(i Instance, snapshots bool, args migration.SourceArgs) (migration.StorageSourceDriver, error)
	RefreshInstance(i Instance, src Instance, snapshots bool, op *operations.Operation) error
	BackupInstance(i Instance, targetPath string, optimized bool, snapshots bool, op *operations.Operation) error

	GetInstanceUsage(i Instance) (int64, error)
	SetInstanceQuota(i Instance, size string, op *operations.Operation) error

	MountInstance(i Instance, op *operations.Operation) (bool, error)
	UnmountInstance(i Instance, op *operations.Operation) (bool, error)
	GetInstanceDisk(i Instance) (string, string, error)

	// Instance snapshots.
	CreateInstanceSnapshot(i Instance, name string, op *operations.Operation) error
	RenameInstanceSnapshot(i Instance, newName string, op *operations.Operation) error
	DeleteInstanceSnapshot(i Instance, op *operations.Operation) error
	RestoreInstanceSnapshot(i Instance, op *operations.Operation) error
	MountInstanceSnapshot(i Instance, op *operations.Operation) (bool, error)
	UnmountInstanceSnapshot(i Instance, op *operations.Operation) (bool, error)

	// Images.
	EnsureImage(fingerprint string, op *operations.Operation) error
	DeleteImage(fingerprint string, op *operations.Operation) error

	// Custom volumes.
	CreateCustomVolume(volName, desc string, config map[string]string, op *operations.Operation) error
	CreateCustomVolumeFromCopy(volName, desc string, config map[string]string, srcPoolName, srcVolName string, srcVolOnly bool, op *operations.Operation) error
	UpdateCustomVolume(volName, newDesc string, newConfig map[string]string, op *operations.Operation) error
	RenameCustomVolume(volName string, newVolName string, op *operations.Operation) error
	DeleteCustomVolume(volName string, op *operations.Operation) error
	GetCustomVolumeUsage(volName string) (int64, error)
	MountCustomVolume(volName string, op *operations.Operation) (bool, error)
	UnmountCustomVolume(volName string, op *operations.Operation) (bool, error)

	// Custom volume snapshots.
	CreateCustomVolumeSnapshot(volName string, newSnapshotName string, op *operations.Operation) error
	RenameCustomVolumeSnapshot(volName string, newSnapshotName string, op *operations.Operation) error
	DeleteCustomVolumeSnapshot(volName string, op *operations.Operation) error
	RestoreCustomVolume(volName string, snapshotName string, op *operations.Operation) error

	// Custom volume migration.
	MigrationTypes(contentType drivers.ContentType) []migration.Type
	CreateCustomVolumeFromMigration(conn io.ReadWriteCloser, args migration.VolumeTargetArgs, op *operations.Operation) error
	MigrateCustomVolume(conn io.ReadWriteCloser, args migration.VolumeSourceArgs, op *operations.Operation) error
}
<|endoftext|>


<commit_before>package p2p

import (
	"bytes"
	"testing"

	"github.com/tendermint/tendermint/libs/log"
)

func TestListener(t *testing.T) {
	// Create a listener
	l := NewDefaultListener("tcp://:8001", "", false, log.TestingLogger())

	// Dial the listener
	lAddr := l.ExternalAddress()
	connOut, err := lAddr.Dial()
	if err != nil {
		t.Fatalf("Could not connect to listener address %v", lAddr)
	} else {
		t.Logf("Created a connection to listener address %v", lAddr)
	}
	connIn, ok := <-l.Connections()
	if !ok {
		t.Fatalf("Could not get inbound connection from listener")
	}

	msg := []byte("hi!")
	go func() {
		_, err := connIn.Write(msg)
		if err != nil {
			t.Error(err)
		}
	}()
	b := make([]byte, 32)
	n, err := connOut.Read(b)
	if err != nil {
		t.Fatalf("Error reading off connection: %v", err)
	}

	b = b[:n]
	if !bytes.Equal(msg, b) {
		t.Fatalf("Got %s, expected %s", b, msg)
	}

	// Close the server, no longer needed.
	l.Stop()
}
<commit_msg>add test<commit_after>package p2p

import (
	"bytes"
	"net"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/tendermint/tendermint/libs/log"
)

func TestListener(t *testing.T) {
	// Create a listener
	l := NewDefaultListener("tcp://:8001", "", false, log.TestingLogger())

	// Dial the listener
	lAddr := l.ExternalAddress()
	connOut, err := lAddr.Dial()
	if err != nil {
		t.Fatalf("Could not connect to listener address %v", lAddr)
	} else {
		t.Logf("Created a connection to listener address %v", lAddr)
	}
	connIn, ok := <-l.Connections()
	if !ok {
		t.Fatalf("Could not get inbound connection from listener")
	}

	msg := []byte("hi!")
	go func() {
		_, err := connIn.Write(msg)
		if err != nil {
			t.Error(err)
		}
	}()
	b := make([]byte, 32)
	n, err := connOut.Read(b)
	if err != nil {
		t.Fatalf("Error reading off connection: %v", err)
	}

	b = b[:n]
	if !bytes.Equal(msg, b) {
		t.Fatalf("Got %s, expected %s", b, msg)
	}

	// Close the server, no longer needed.
	l.Stop()
}

func TestExternalAddress(t *testing.T) {
	{
		// Create a listener with no external addr. Should default
		// to local ipv4.
		l := NewDefaultListener("tcp://:8001", "", false, log.TestingLogger())
		lAddr := l.ExternalAddress().String()
		_, _, err := net.SplitHostPort(lAddr)
		require.Nil(t, err)
		spl := strings.Split(lAddr, ".")
		require.Equal(t, len(spl), 4)
		l.Stop()
	}

	{
		// Create a listener with set external ipv4 addr.
		setExAddr := "8.8.8.8:8080"
		l := NewDefaultListener("tcp://:8001", setExAddr, false, log.TestingLogger())
		lAddr := l.ExternalAddress().String()
		require.Equal(t, lAddr, setExAddr)
		l.Stop()
	}

	{
		// Invalid external addr causes panic
		setExAddr := "awrlsckjnal:8080"
		require.Panics(t, func() { NewDefaultListener("tcp://:8001", setExAddr, false, log.TestingLogger()) })
	}
}
<|endoftext|>


<commit_before>package command

import (
	"flag"
	"fmt"
	"strings"
	"time"

	"github.com/koding/kite"
	"github.com/koding/kite/kitekey"
	"github.com/mitchellh/cli"
)

const (
	defaultKontrolURL = "https://discover.koding.io/kite"
	defaultUsername   = "koding"
)

type Register struct {
	KiteClient *kite.Kite
	Ui         cli.Ui
}

func NewRegister() cli.CommandFactory {
	return func() (cli.Command, error) {
		return &Register{
			KiteClient: DefaultKiteClient,
			Ui:         DefaultUi,
		}, nil
	}
}

func (c *Register) Synopsis() string {
	return "Registers this host to a kite authority"
}

func (c *Register) Help() string {
	helpText := `
Usage: kitectl register [options]

  Registers your host to a kite authority.
  If no server is specified, "https://discovery.koding.io/kite" is the default.

Options:

  -to=https://discovery.koding.io/kite  Kontrol URL
  -username=koding                      Username
`
	return strings.TrimSpace(helpText)
}

func (c *Register) Run(args []string) int {
	var kontrolURL, username string
	var err error

	flags := flag.NewFlagSet("register", flag.ExitOnError)
	flags.StringVar(&kontrolURL, "to", defaultKontrolURL, "Kontrol URL")
	flags.StringVar(&username, "username", "", "Username")
	flags.Parse(args)

	// Open up a prompt
	if username == "" {
		username, err = c.Ui.Ask(fmt.Sprintf("Username (%s):", defaultUsername))
		if err != nil {
			c.Ui.Error(err.Error())
			return 1
		}
		// User can just press enter to use the default on the prompt
		if username == "" {
			username = defaultUsername
		}
	}

	c.KiteClient.Config.Username = username

	if _, err := kitekey.Read(); err == nil {
		c.Ui.Info("Already registered. Registering again...")
	}

	kontrol := c.KiteClient.NewClient(kontrolURL)
	if err := kontrol.Dial(); err != nil {
		c.Ui.Error(err.Error())
		return 1
	}

	result, err := kontrol.TellWithTimeout("registerMachine", 5*time.Minute, username)
	if err != nil {
		c.Ui.Error(err.Error())
		return 1
	}

	if err := kitekey.Write(result.MustString()); err != nil {
		c.Ui.Error(err.Error())
		return 1
	}

	c.Ui.Info("Registered successfully")

	return 0
}
<commit_msg>kitectl: Fix wrong default Kontrol URL.<commit_after>package command

import (
	"flag"
	"fmt"
	"strings"
	"time"

	"github.com/koding/kite"
	"github.com/koding/kite/kitekey"
	"github.com/mitchellh/cli"
)

const (
	defaultKontrolURL = "https://discovery.koding.io/kite"
	defaultUsername   = "koding"
)

type Register struct {
	KiteClient *kite.Kite
	Ui         cli.Ui
}

func NewRegister() cli.CommandFactory {
	return func() (cli.Command, error) {
		return &Register{
			KiteClient: DefaultKiteClient,
			Ui:         DefaultUi,
		}, nil
	}
}

func (c *Register) Synopsis() string {
	return "Registers this host to a kite authority"
}

func (c *Register) Help() string {
	helpText := `
Usage: kitectl register [options]

  Registers your host to a kite authority.
  If no server is specified, "https://discovery.koding.io/kite" is the default.

Options:

  -to=https://discovery.koding.io/kite  Kontrol URL
  -username=koding                      Username
`
	return strings.TrimSpace(helpText)
}

func (c *Register) Run(args []string) int {
	var kontrolURL, username string
	var err error

	flags := flag.NewFlagSet("register", flag.ExitOnError)
	flags.StringVar(&kontrolURL, "to", defaultKontrolURL, "Kontrol URL")
	flags.StringVar(&username, "username", "", "Username")
	flags.Parse(args)

	// Open up a prompt
	if username == "" {
		username, err = c.Ui.Ask(fmt.Sprintf("Username (%s):", defaultUsername))
		if err != nil {
			c.Ui.Error(err.Error())
			return 1
		}
		// User can just press enter to use the default on the prompt
		if username == "" {
			username = defaultUsername
		}
	}

	c.KiteClient.Config.Username = username

	if _, err := kitekey.Read(); err == nil {
		c.Ui.Info("Already registered. Registering again...")
	}

	kontrol := c.KiteClient.NewClient(kontrolURL)
	if err := kontrol.Dial(); err != nil {
		c.Ui.Error(err.Error())
		return 1
	}

	result, err := kontrol.TellWithTimeout("registerMachine", 5*time.Minute, username)
	if err != nil {
		c.Ui.Error(err.Error())
		return 1
	}

	if err := kitekey.Write(result.MustString()); err != nil {
		c.Ui.Error(err.Error())
		return 1
	}

	c.Ui.Info("Registered successfully")

	return 0
}
<|endoftext|>