<commit_before>package elasticsearch

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/compose/transporter/adaptor"
	"github.com/compose/transporter/adaptor/elasticsearch/clients"
	// used to call init function for each client to register itself
	_ "github.com/compose/transporter/adaptor/elasticsearch/clients/all"
	"github.com/compose/transporter/client"
	"github.com/compose/transporter/log"
	version "github.com/hashicorp/go-version"
)

const (
	description  = "an elasticsearch sink adaptor"
	sampleConfig = `    type: elasticsearch
    uri: https://username:password@hostname:port
    # timeout: 10s # defaults to 30s
    # aws_access_key: XXX # used for signing requests to AWS Elasticsearch service
    # aws_access_secret: XXX # used for signing requests to AWS Elasticsearch service`
)

var (
	_ adaptor.Adaptor = &Elasticsearch{}
)

// Elasticsearch is an adaptor to connect a pipeline to
// an elasticsearch cluster.
type Elasticsearch struct {
	adaptor.BaseConfig
	AWSAccessKeyID  string `json:"aws_access_key" doc:"credentials for use with AWS Elasticsearch service"`
	AWSAccessSecret string `json:"aws_access_secret" doc:"credentials for use with AWS Elasticsearch service"`
}

// Description for the Elasticsearcb adaptor
func (e *Elasticsearch) Description() string {
	return description
}

// SampleConfig for elasticsearch adaptor
func (e *Elasticsearch) SampleConfig() string {
	return sampleConfig
}

func init() {
	adaptor.Add(
		"elasticsearch",
		func() adaptor.Adaptor {
			return &Elasticsearch{}
		},
	)
}

// Client returns a client that doesn't do anything other than fulfill the client.Client interface.
func (e *Elasticsearch) Client() (client.Client, error) {
	return &client.Mock{}, nil
}

// Reader returns an error because this adaptor is currently not supported as a Source.
func (e *Elasticsearch) Reader() (client.Reader, error) {
	return nil, adaptor.ErrFuncNotSupported{Name: "Reader()", Func: "elasticsearch"}
}

// Writer determines the which underlying writer to used based on the cluster's version.
func (e *Elasticsearch) Writer(done chan struct{}, wg *sync.WaitGroup) (client.Writer, error) {
	index, _, _ := adaptor.CompileNamespace(e.Namespace)
	return setupWriter(index, e)
}

func setupWriter(index string, conf *Elasticsearch) (client.Writer, error) {
	uri, err := url.Parse(conf.URI)
	if err != nil {
		return nil, client.InvalidURIError{URI: conf.URI, Err: err.Error()}
	}

	hostsAndPorts := strings.Split(uri.Host, ",")
	stringVersion, err := determineVersion(fmt.Sprintf("%s://%s", uri.Scheme, hostsAndPorts[0]), uri.User)
	if err != nil {
		return nil, err
	}

	v, err := version.NewVersion(stringVersion)
	if err != nil {
		return nil, client.VersionError{URI: conf.URI, V: stringVersion, Err: err.Error()}
	}

	timeout, err := time.ParseDuration(conf.Timeout)
	if err != nil {
		log.Infof("failed to parse duration, %s, falling back to default timeout of 30s", conf.Timeout)
		timeout = 30 * time.Second
	}

	httpClient := &http.Client{
		Timeout:   timeout,
		Transport: newTransport(conf.AWSAccessKeyID, conf.AWSAccessSecret),
	}

	for _, vc := range clients.Clients {
		if vc.Constraint.Check(v) {
			urls := make([]string, len(hostsAndPorts))
			for i, hAndP := range hostsAndPorts {
				urls[i] = fmt.Sprintf("%s://%s", uri.Scheme, hAndP)
			}
			opts := &clients.ClientOptions{
				URLs:       urls,
				UserInfo:   uri.User,
				HTTPClient: httpClient,
				Index:      index,
			}
			versionedClient, _ := vc.Creator(opts)
			return versionedClient, nil
		}
	}

	return nil, client.VersionError{URI: conf.URI, V: stringVersion, Err: "unsupported client"}
}

func determineVersion(uri string, user *url.Userinfo) (string, error) {
	req, err := http.NewRequest(http.MethodGet, uri, nil)
	if err != nil {
		return "", err
	}
	if user != nil {
		if pwd, ok := user.Password(); ok {
			req.SetBasicAuth(user.Username(), pwd)
		}
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", client.ConnectError{Reason: uri}
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", client.VersionError{URI: uri, V: "", Err: "unable to read response body"}
	}
	defer resp.Body.Close()
	var r struct {
		Name    string `json:"name"`
		Version struct {
			Number string `json:"number"`
		} `json:"version"`
	}
	if resp.StatusCode != http.StatusOK {
		return "", client.VersionError{URI: uri, V: "", Err: fmt.Sprintf("bad status code: %d", resp.StatusCode)}
	}
	err = json.Unmarshal(body, &r)
	if err != nil {
		return "", client.VersionError{URI: uri, V: "", Err: fmt.Sprintf("malformed JSON: %s", body)}
	} else if r.Version.Number == "" {
		return "", client.VersionError{URI: uri, V: "", Err: fmt.Sprintf("missing version: %s", body)}
	}
	return r.Version.Number, nil
}
<commit_msg>Fix ES to use ELASTICSEARCH_URI (#305)<commit_after>package elasticsearch

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/compose/transporter/adaptor"
	"github.com/compose/transporter/adaptor/elasticsearch/clients"
	// used to call init function for each client to register itself
	_ "github.com/compose/transporter/adaptor/elasticsearch/clients/all"
	"github.com/compose/transporter/client"
	"github.com/compose/transporter/log"
	version "github.com/hashicorp/go-version"
)

const (
	description  = "an elasticsearch sink adaptor"
	sampleConfig = `    type: elasticsearch
    uri: ${ELASTICSEARCH_URI}
    # timeout: 10s # defaults to 30s
    # aws_access_key: XXX # used for signing requests to AWS Elasticsearch service
    # aws_access_secret: XXX # used for signing requests to AWS Elasticsearch service`
)

var (
	_ adaptor.Adaptor = &Elasticsearch{}
)

// Elasticsearch is an adaptor to connect a pipeline to
// an elasticsearch cluster.
type Elasticsearch struct {
	adaptor.BaseConfig
	AWSAccessKeyID  string `json:"aws_access_key" doc:"credentials for use with AWS Elasticsearch service"`
	AWSAccessSecret string `json:"aws_access_secret" doc:"credentials for use with AWS Elasticsearch service"`
}

// Description for the Elasticsearcb adaptor
func (e *Elasticsearch) Description() string {
	return description
}

// SampleConfig for elasticsearch adaptor
func (e *Elasticsearch) SampleConfig() string {
	return sampleConfig
}

func init() {
	adaptor.Add(
		"elasticsearch",
		func() adaptor.Adaptor {
			return &Elasticsearch{}
		},
	)
}

// Client returns a client that doesn't do anything other than fulfill the client.Client interface.
func (e *Elasticsearch) Client() (client.Client, error) {
	return &client.Mock{}, nil
}

// Reader returns an error because this adaptor is currently not supported as a Source.
func (e *Elasticsearch) Reader() (client.Reader, error) {
	return nil, adaptor.ErrFuncNotSupported{Name: "Reader()", Func: "elasticsearch"}
}

// Writer determines the which underlying writer to used based on the cluster's version.
func (e *Elasticsearch) Writer(done chan struct{}, wg *sync.WaitGroup) (client.Writer, error) {
	index, _, _ := adaptor.CompileNamespace(e.Namespace)
	return setupWriter(index, e)
}

func setupWriter(index string, conf *Elasticsearch) (client.Writer, error) {
	uri, err := url.Parse(conf.URI)
	if err != nil {
		return nil, client.InvalidURIError{URI: conf.URI, Err: err.Error()}
	}

	hostsAndPorts := strings.Split(uri.Host, ",")
	stringVersion, err := determineVersion(fmt.Sprintf("%s://%s", uri.Scheme, hostsAndPorts[0]), uri.User)
	if err != nil {
		return nil, err
	}

	v, err := version.NewVersion(stringVersion)
	if err != nil {
		return nil, client.VersionError{URI: conf.URI, V: stringVersion, Err: err.Error()}
	}

	timeout, err := time.ParseDuration(conf.Timeout)
	if err != nil {
		log.Infof("failed to parse duration, %s, falling back to default timeout of 30s", conf.Timeout)
		timeout = 30 * time.Second
	}

	httpClient := &http.Client{
		Timeout:   timeout,
		Transport: newTransport(conf.AWSAccessKeyID, conf.AWSAccessSecret),
	}

	for _, vc := range clients.Clients {
		if vc.Constraint.Check(v) {
			urls := make([]string, len(hostsAndPorts))
			for i, hAndP := range hostsAndPorts {
				urls[i] = fmt.Sprintf("%s://%s", uri.Scheme, hAndP)
			}
			opts := &clients.ClientOptions{
				URLs:       urls,
				UserInfo:   uri.User,
				HTTPClient: httpClient,
				Index:      index,
			}
			versionedClient, _ := vc.Creator(opts)
			return versionedClient, nil
		}
	}

	return nil, client.VersionError{URI: conf.URI, V: stringVersion, Err: "unsupported client"}
}

func determineVersion(uri string, user *url.Userinfo) (string, error) {
	req, err := http.NewRequest(http.MethodGet, uri, nil)
	if err != nil {
		return "", err
	}
	if user != nil {
		if pwd, ok := user.Password(); ok {
			req.SetBasicAuth(user.Username(), pwd)
		}
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", client.ConnectError{Reason: uri}
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", client.VersionError{URI: uri, V: "", Err: "unable to read response body"}
	}
	defer resp.Body.Close()
	var r struct {
		Name    string `json:"name"`
		Version struct {
			Number string `json:"number"`
		} `json:"version"`
	}
	if resp.StatusCode != http.StatusOK {
		return "", client.VersionError{URI: uri, V: "", Err: fmt.Sprintf("bad status code: %d", resp.StatusCode)}
	}
	err = json.Unmarshal(body, &r)
	if err != nil {
		return "", client.VersionError{URI: uri, V: "", Err: fmt.Sprintf("malformed JSON: %s", body)}
	} else if r.Version.Number == "" {
		return "", client.VersionError{URI: uri, V: "", Err: fmt.Sprintf("missing version: %s", body)}
	}
	return r.Version.Number, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"strconv"
	"strings"
)

const YNABDateFormat string = "02.01.2006"

type YNABOutput struct {
}

func (*YNABOutput) WriteHeader(w *bufio.Writer) {
	fmt.Fprintln(w, "Date,Payee,Category,Memo,Outflow,Inflow")
}

func (*YNABOutput) Process(w *bufio.Writer, t *Transaction) {
	date := t.Date.Format(YNABDateFormat)
	outflow := ""
	inflow := ""
	if t.ValueCent < 0 {
		outflow = formatValue(-t.ValueCent)
	} else {
		inflow = formatValue(t.ValueCent)
	}
	output := strings.Join([]string{date, t.Payee, t.Category, t.Comment, outflow, inflow}, ",")
	fmt.Fprintln(w, output)
}

func (y *YNABOutput) BatchProcess(w *bufio.Writer, t []*Transaction) {
	for _, trans := range t {
		y.Process(w, trans)
	}
}

func formatValue(v int) string {
	neg := v < 0
	if neg {
		v = -v
	}
	str := strconv.FormatInt(int64(v), 10)
	for len(str) < 3 {
		str = "0" + str
	}
	eur := str[:len(str)-2]
	cent := str[len(str)-2:]
	result := eur + "." + cent
	if neg {
		result = "-" + result
	}
	return result
}
<commit_msg>Fixed commata in output<commit_after>package main

import (
	"bufio"
	"fmt"
	"strconv"
	"strings"
)

const YNABDateFormat string = "02.01.2006"

type YNABOutput struct {
}

func (*YNABOutput) WriteHeader(w *bufio.Writer) {
	fmt.Fprintln(w, "Date,Payee,Category,Memo,Outflow,Inflow")
}

func (*YNABOutput) Process(w *bufio.Writer, t *Transaction) {
	date := t.Date.Format(YNABDateFormat)
	outflow := ""
	inflow := ""
	if t.ValueCent < 0 {
		outflow = formatValue(-t.ValueCent)
	} else {
		inflow = formatValue(t.ValueCent)
	}

	removeInvalidChars(t)

	output := strings.Join([]string{date, t.Payee, t.Category, t.Comment, outflow, inflow}, ",")
	fmt.Fprintln(w, output)
}

func (y *YNABOutput) BatchProcess(w *bufio.Writer, t []*Transaction) {
	for _, trans := range t {
		y.Process(w, trans)
	}
}

func formatValue(v int) string {
	neg := v < 0
	if neg {
		v = -v
	}
	str := strconv.FormatInt(int64(v), 10)
	for len(str) < 3 {
		str = "0" + str
	}
	eur := str[:len(str)-2]
	cent := str[len(str)-2:]
	result := eur + "." + cent
	if neg {
		result = "-" + result
	}
	return result
}

func removeInvalidChars(t *Transaction) {
	t.Category = strings.Replace(t.Category, ",", ";", -1)
	t.Comment = strings.Replace(t.Comment, ",", ";", -1)
	t.Payee = strings.Replace(t.Payee, ",", ";", -1)
}
<|endoftext|>


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

// This file defines the Null refactoring, which makes no changes to a program.
// It is for testing only (and can be used as a template for building new
// refactorings).

package doctor

// A nullRefactoring makes no changes to a program.
type nullRefactoring struct {
	refactoringBase
	optShowAST        bool
	optShowPackages   bool
	optShowReferences bool
}

func (r *nullRefactoring) Description() *Description {
	return &Description{
		Name:   "Null Refactoring",
		Params: []string{},
	}
}

func (r *nullRefactoring) Run(config *Config) *Result {
	r.refactoringBase.Run(config)
	r.Log.ChangeInitialErrorsToWarnings()
	if r.Log.ContainsErrors() {
		return &r.Result
	}

	if len(config.Args) != 0 {
		r.Log.Log(FATAL_ERROR, "This refactoring takes no arguments.")
		return &r.Result
	}

	return &r.Result
}
<commit_msg>Removed unnecessary fields from nullRefactoring<commit_after>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file defines the Null refactoring, which makes no changes to a program.
// It is for testing only (and can be used as a template for building new
// refactorings).

package doctor

// A nullRefactoring makes no changes to a program.
type nullRefactoring struct {
	refactoringBase
}

func (r *nullRefactoring) Description() *Description {
	return &Description{
		Name:   "Null Refactoring",
		Params: []string{},
	}
}

func (r *nullRefactoring) Run(config *Config) *Result {
	r.refactoringBase.Run(config)
	r.Log.ChangeInitialErrorsToWarnings()
	if r.Log.ContainsErrors() {
		return &r.Result
	}

	if len(config.Args) != 0 {
		r.Log.Log(FATAL_ERROR, "This refactoring takes no arguments.")
		return &r.Result
	}

	return &r.Result
}
<|endoftext|>


<commit_before>package jwd

import (
	"math"
	"strings"
)

// According to this tool: http://www.csun.edu/english/edit_distance.php
// parameter order should have no difference in the result. Therefore,
// to avoid panicing later on, we will order the strings according to
// their length.
func order(s1, s2 string) (string, string) {
	if strings.Count(s1, "")-1 <= strings.Count(s2, "")-1 {
		return s1, s2
	} else {
		return s2, s1
	}

}

// Calculates Jaro-Winkler distance of two strings. The function lowercases and sorts the parameters
// so that that the longest string is evaluated against the shorter one.
func Calculate(s1, s2 string) float64 {

	s1, s2 = order(strings.ToLower(s1), strings.ToLower(s2))

	// This avoids the function to return NaN.
	if strings.Count(s1, "") == 1 || strings.Count(s2, "") == 1 {
		return float64(0)
	}

	// m as `matching characters`
	// t as `transposition`
	// l as `the length of common prefix at the start of the string up to a maximum of 4 characters`.
	// See more: https://en.wikipedia.org/wiki/Jaro%E2%80%93Winkler_distance
	m := 0
	t := 0
	l := 0

	window := math.Floor(float64(math.Max(float64(len(s1)), float64(len(s2)))/2) - 1)

	//debug:
	//fmt.Println("s1:", s1, "s2:", s2)
	//fmt.Println("Match window:", window)
	//fmt.Println("len(s1):", len(s1), "len(s2):", len(s2))

	for i := 0; i < len(s1); i++ {
		// Exact match
		if s1[i] == s2[i] {
			m++
			// Common prefix limitter
			if i == l && i < 4 {
				l++
			}
		} else {
			if strings.Contains(s2, string(s1[i])) {
				// The character is also considered matching if the amount of characters between the occurances in s1 and s2
				// is less than match window
				gap := strings.Index(s2, string(s1[i])) - strings.Index(s1, string(s1[i]))
				if gap <= int(window) {
					m++
					// Check if transposition is in reach of window
					for k := i; k < len(s1); k++ {
						if strings.Index(s2, string(s1[k])) <= i {
							t++
						}
					}
				}
			}
		}
	}

	distance := (float64(m)/float64(len(s1)) + float64(m)/float64(len(s2)) + (float64(m)-math.Floor(float64(t)/float64(2)))/float64(m)) / float64(3)
	jwd := distance + (float64(l) * float64(0.1) * (float64(1) - distance))

	//debug:
	//fmt.Println("- transpositions:", t)
	//fmt.Println("- matches:", m)
	//fmt.Println("- l:", l)
	//fmt.Println(jwd)

	return jwd

}
<commit_msg>golint fixes suggested by godoc<commit_after>package jwd

import (
	"math"
	"strings"
)

// According to this tool: http://www.csun.edu/english/edit_distance.php
// parameter order should have no difference in the result. Therefore,
// to avoid panicing later on, we will order the strings according to
// their length.
func order(s1, s2 string) (string, string) {
	if strings.Count(s1, "")-1 <= strings.Count(s2, "")-1 {
		return s1, s2
	}
	return s2, s1
}

// Calculate calculates Jaro-Winkler distance of two strings. The function lowercases and sorts the parameters
// so that that the longest string is evaluated against the shorter one.
func Calculate(s1, s2 string) float64 {

	s1, s2 = order(strings.ToLower(s1), strings.ToLower(s2))

	// This avoids the function to return NaN.
	if strings.Count(s1, "") == 1 || strings.Count(s2, "") == 1 {
		return float64(0)
	}

	// m as `matching characters`
	// t as `transposition`
	// l as `the length of common prefix at the start of the string up to a maximum of 4 characters`.
	// See more: https://en.wikipedia.org/wiki/Jaro%E2%80%93Winkler_distance
	m := 0
	t := 0
	l := 0

	window := math.Floor(float64(math.Max(float64(len(s1)), float64(len(s2)))/2) - 1)

	//debug:
	//fmt.Println("s1:", s1, "s2:", s2)
	//fmt.Println("Match window:", window)
	//fmt.Println("len(s1):", len(s1), "len(s2):", len(s2))

	for i := 0; i < len(s1); i++ {
		// Exact match
		if s1[i] == s2[i] {
			m++
			// Common prefix limitter
			if i == l && i < 4 {
				l++
			}
		} else {
			if strings.Contains(s2, string(s1[i])) {
				// The character is also considered matching if the amount of characters between the occurances in s1 and s2
				// is less than match window
				gap := strings.Index(s2, string(s1[i])) - strings.Index(s1, string(s1[i]))
				if gap <= int(window) {
					m++
					// Check if transposition is in reach of window
					for k := i; k < len(s1); k++ {
						if strings.Index(s2, string(s1[k])) <= i {
							t++
						}
					}
				}
			}
		}
	}

	distance := (float64(m)/float64(len(s1)) + float64(m)/float64(len(s2)) + (float64(m)-math.Floor(float64(t)/float64(2)))/float64(m)) / float64(3)
	jwd := distance + (float64(l) * float64(0.1) * (float64(1) - distance))

	//debug:
	//fmt.Println("- transpositions:", t)
	//fmt.Println("- matches:", m)
	//fmt.Println("- l:", l)
	//fmt.Println(jwd)

	return jwd

}
<|endoftext|>


<commit_before>package redismq

import (
	"github.com/adeven/redis"
	"time"
)

func (self *BufferedQueue) Put(payload string) error {
	p := &Package{CreatedAt: time.Now(), Payload: payload, Queue: self}
	self.Buffer <- p
	return nil
}

func (self *BufferedQueue) FlushBuffer() {
	flushing := make(chan bool, 1)
	self.flushStatus <- flushing
	<-flushing
	return
}

func (self *BufferedQueue) startWritingBufferToRedis() {
	go func() {
		self.nextWrite = time.Now().Unix()
		for {
			if len(self.Buffer) == cap(self.Buffer) || time.Now().Unix() >= self.nextWrite {
				size := len(self.Buffer)
				self.redisClient.Pipelined(func(c *redis.PipelineClient) {
					for i := 0; i < size; i++ {
						p := <-self.Buffer
						c.LPush(self.InputName(), p.GetString())
					}
					c.IncrBy(self.InputCounterName(), int64(size))
				})
				for i := 0; i < len(self.flushStatus); i++ {
					c := <-self.flushStatus
					c <- true
				}
				self.nextWrite = time.Now().Unix() + 1
			}
		}
	}()
}

func (self *BufferedQueue) startHeartbeat() {
	go func() {
		for {
			self.redisClient.SetEx(self.HeartbeatName(), 1, "ping")
			time.Sleep(500 * time.Millisecond)
		}
	}()
}
<commit_msg>fix bussy wait<commit_after>package redismq

import (
	"github.com/adeven/redis"
	"time"
)

func (self *BufferedQueue) Put(payload string) error {
	p := &Package{CreatedAt: time.Now(), Payload: payload, Queue: self}
	self.Buffer <- p
	return nil
}

func (self *BufferedQueue) FlushBuffer() {
	flushing := make(chan bool, 1)
	self.flushStatus <- flushing
	<-flushing
	return
}

func (self *BufferedQueue) startWritingBufferToRedis() {
	go func() {
		self.nextWrite = time.Now().Unix()
		for {
			if len(self.Buffer) == cap(self.Buffer) || time.Now().Unix() >= self.nextWrite {
				size := len(self.Buffer)
				self.redisClient.Pipelined(func(c *redis.PipelineClient) {
					for i := 0; i < size; i++ {
						p := <-self.Buffer
						c.LPush(self.InputName(), p.GetString())
					}
					c.IncrBy(self.InputCounterName(), int64(size))
				})
				for i := 0; i < len(self.flushStatus); i++ {
					c := <-self.flushStatus
					c <- true
				}
				self.nextWrite = time.Now().Unix() + 1
			}
			time.Sleep(10 * time.Millisecond)
		}
	}()
}

func (self *BufferedQueue) startHeartbeat() {
	go func() {
		for {
			self.redisClient.SetEx(self.HeartbeatName(), 1, "ping")
			time.Sleep(500 * time.Millisecond)
		}
	}()
}
<|endoftext|>


<commit_before>package mcg

import "fmt"
import "bytes"
import "encoding/gob"
import "encoding/json"

import "github.com/streadway/amqp"

// -----------------------------------------------------------------------------
// -- AMQPAgent
// -----------------------------------------------------------------------------

// AMQPAgent DOC: ..
type AMQPAgent struct {
	URL   string
	Retry int

	conn *amqp.Connection
	done chan *amqp.Error
}

// NewAMQPAgent DOC: ..
func NewAMQPAgent(url string) (*AMQPAgent, error) {
	var a = &AMQPAgent{URL: url}
	return a, a.Connect()
}

// NewAMQPBroker DOC: ..
func NewAMQPBroker(url, topic string) (*Broker, error) {
	var err error
	var agent *AMQPAgent

	if agent, err = NewAMQPAgent(url); err != nil {
		return nil, err
	}

	return NewBroker(topic, agent), nil
}

// ---

// Connect DOC: ..
func (a *AMQPAgent) Connect() error {
	var err error

	if a.URL == "" {
		return fmt.Errorf("can't connect with empty url")
	}

	if a.conn, err = amqp.Dial(a.URL); err != nil {
		return err
	}

	a.conn.NotifyClose(a.done)

	return nil
}

// Close DOC: ..
func (a *AMQPAgent) Close() error {
	return a.conn.Close()
}

// Done DOC: ..
func (a *AMQPAgent) Done() error {
	var err = <-a.done
	return err
}

// ---

// Send DOC: ..
func (a *AMQPAgent) Send(topic, key string, message *Message) error {
	if a.conn == nil {
		return fmt.Errorf("can't send message to route `%v:%v` to nil connection", topic, key)
	}

	var err error
	var ch *amqp.Channel

	if ch, err = a.conn.Channel(); err != nil {
		return err
	}

	defer ch.Close()

	var route = topic + ":" + key

	if _, err = amqpQueueDeclare(ch, route); err != nil {
		return err
	}

	var body []byte

	if body, err = amqpEncodeJSON(message); err != nil {
		return err
	}

	var pub = amqp.Publishing{
		DeliveryMode: amqp.Persistent,
		ContentType:  "text/plain",
		Body:         body,
	}

	if err = ch.Publish("", route, false, false, pub); err != nil {
		return err
	}

	return nil
}

// Receive DOC: ..
// TODO: Response value should be error channel
func (a *AMQPAgent) Receive(topic, key string, limit int, handler HandlerFunc) error {
	var err error
	var ch *amqp.Channel

	if ch, err = a.conn.Channel(); err != nil {
		return err
	}

	defer ch.Close()

	if err = ch.Qos(limit, 0, false); err != nil {
		return err
	}

	var route = topic + ":" + key

	if _, err = amqpQueueDeclare(ch, route); err != nil {
		return err
	}

	var consumer = route
	var deliveries <-chan amqp.Delivery

	// TODO: Check that these arguments work for handling work queue messages
	// ARGS: queue, consumer, auto-ack, exclusive, no-local, no-wait, args
	if deliveries, err = ch.Consume(route, consumer, false, false, false, false, nil); err != nil {
		return err
	}

	for d := range deliveries {
		go a.procDelivery(d, handler)
	}

	return nil
}

func (a *AMQPAgent) procDelivery(delivery amqp.Delivery, handler HandlerFunc) {
	var err error
	var message *Message

	if message, err = amqpDecodeJSON(delivery.Body); err != nil {
		fmt.Printf("Can't decode: %v\n", err)
		delivery.Nack(false, false)
		return
	}

	if err = handler(message); err != nil {
		delivery.Nack(false, true)
		return
	}

	delivery.Ack(false)
}

// ---

func amqpEncodeJSON(message *Message) ([]byte, error) {
	var err error
	var body []byte

	if body, err = json.Marshal(message); err != nil {
		return nil, err
	}

	return body, nil
}

func amqpDecodeJSON(body []byte) (*Message, error) {
	var err error
	var message Message

	if err = json.Unmarshal(body, &message); err != nil {
		return nil, err
	}

	return &message, nil
}

func amqpEncodeGOB(message *Message) ([]byte, error) {
	var err error
	var buf bytes.Buffer

	if err = gob.NewEncoder(&buf).Encode(message); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

func amqpDecodeGOB(body []byte) (*Message, error) {
	var err error
	var message *Message
	var buf bytes.Buffer

	if _, err = buf.Write(body); err != nil {
		return nil, err
	}

	if err = gob.NewDecoder(&buf).Decode(message); err != nil {
		return nil, err
	}

	return message, nil
}

// ---

func amqpExchangeDeclare(ch *amqp.Channel, topic string) error {
	// TODO: Ensure this doesn't drop exchange topics at unexpected times.
	// Args: name, kind, durable, autoDelete, internal, noWait, args
	return ch.ExchangeDeclare(topic, "topic", true, false, false, false, nil)
}

func amqpQueueBind(ch *amqp.Channel, topic, key string) (*amqp.Queue, error) {
	var err error

	// NOTE: We're only doing this when we successfully instantiate connections.
	if err = amqpExchangeDeclare(ch, topic); err != nil {
		return nil, err
	}

	var queue amqp.Queue
	var name = topic + ":" + key

	// Args: name, durable, autoDelete, exclusive, noWait, args
	if queue, err = amqpQueueDeclare(ch, name); err != nil {
		return nil, err
	}

	// Args: name, key, exchange, noWait, args
	if err = ch.QueueBind(name, key, topic, false, nil); err != nil {
		return nil, err
	}

	return &queue, nil
}

// ---

func amqpQueueDeclare(ch *amqp.Channel, queue string) (amqp.Queue, error) {
	// Args: name, durable, autoDelete, exclusive, noWait, args
	return ch.QueueDeclare(queue, true, false, false, false, nil)
}
<commit_msg>Ensure URL before connect; give proper error message<commit_after>package mcg

import "fmt"
import "bytes"
import "encoding/gob"
import "encoding/json"

import "github.com/streadway/amqp"

// -----------------------------------------------------------------------------
// -- AMQPAgent
// -----------------------------------------------------------------------------

// AMQPAgent DOC: ..
type AMQPAgent struct {
	URL   string
	Retry int

	conn *amqp.Connection
	done chan *amqp.Error
}

// NewAMQPAgent DOC: ..
func NewAMQPAgent(url string) (*AMQPAgent, error) {
	var a = &AMQPAgent{URL: url}
	return a, a.Connect()
}

// NewAMQPBroker DOC: ..
func NewAMQPBroker(url, topic string) (*Broker, error) {
	var err error
	var agent *AMQPAgent

	if agent, err = NewAMQPAgent(url); err != nil {
		return nil, err
	}

	return NewBroker(topic, agent), nil
}

// ---

// Connect DOC: ..
func (a *AMQPAgent) Connect() error {
	var err error

	if a == nil {
		return fmt.Errorf("unable to connect with nil agent")
	}

	if a.URL == "" {
		return fmt.Errorf("unable to connect with empty url")
	}

	if a.conn, err = amqp.Dial(a.URL); err != nil {
		return err
	}

	a.conn.NotifyClose(a.done)

	return nil
}

// Close DOC: ..
func (a *AMQPAgent) Close() error {
	return a.conn.Close()
}

// Done DOC: ..
func (a *AMQPAgent) Done() error {
	var err = <-a.done
	return err
}

// ---

// Send DOC: ..
func (a *AMQPAgent) Send(topic, key string, message *Message) error {
	if a.conn == nil {
		return fmt.Errorf("can't send message to route `%v:%v` to nil connection", topic, key)
	}

	var err error
	var ch *amqp.Channel

	if ch, err = a.conn.Channel(); err != nil {
		return err
	}

	defer ch.Close()

	var route = topic + ":" + key

	if _, err = amqpQueueDeclare(ch, route); err != nil {
		return err
	}

	var body []byte

	if body, err = amqpEncodeJSON(message); err != nil {
		return err
	}

	var pub = amqp.Publishing{
		DeliveryMode: amqp.Persistent,
		ContentType:  "text/plain",
		Body:         body,
	}

	if err = ch.Publish("", route, false, false, pub); err != nil {
		return err
	}

	return nil
}

// Receive DOC: ..
// TODO: Response value should be error channel
func (a *AMQPAgent) Receive(topic, key string, limit int, handler HandlerFunc) error {
	var err error
	var ch *amqp.Channel

	if ch, err = a.conn.Channel(); err != nil {
		return err
	}

	defer ch.Close()

	if err = ch.Qos(limit, 0, false); err != nil {
		return err
	}

	var route = topic + ":" + key

	if _, err = amqpQueueDeclare(ch, route); err != nil {
		return err
	}

	var consumer = route
	var deliveries <-chan amqp.Delivery

	// TODO: Check that these arguments work for handling work queue messages
	// ARGS: queue, consumer, auto-ack, exclusive, no-local, no-wait, args
	if deliveries, err = ch.Consume(route, consumer, false, false, false, false, nil); err != nil {
		return err
	}

	for d := range deliveries {
		go a.procDelivery(d, handler)
	}

	return nil
}

func (a *AMQPAgent) procDelivery(delivery amqp.Delivery, handler HandlerFunc) {
	var err error
	var message *Message

	if message, err = amqpDecodeJSON(delivery.Body); err != nil {
		fmt.Printf("Can't decode: %v\n", err)
		delivery.Nack(false, false)
		return
	}

	if err = handler(message); err != nil {
		delivery.Nack(false, true)
		return
	}

	delivery.Ack(false)
}

// ---

func amqpEncodeJSON(message *Message) ([]byte, error) {
	var err error
	var body []byte

	if body, err = json.Marshal(message); err != nil {
		return nil, err
	}

	return body, nil
}

func amqpDecodeJSON(body []byte) (*Message, error) {
	var err error
	var message Message

	if err = json.Unmarshal(body, &message); err != nil {
		return nil, err
	}

	return &message, nil
}

func amqpEncodeGOB(message *Message) ([]byte, error) {
	var err error
	var buf bytes.Buffer

	if err = gob.NewEncoder(&buf).Encode(message); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

func amqpDecodeGOB(body []byte) (*Message, error) {
	var err error
	var message *Message
	var buf bytes.Buffer

	if _, err = buf.Write(body); err != nil {
		return nil, err
	}

	if err = gob.NewDecoder(&buf).Decode(message); err != nil {
		return nil, err
	}

	return message, nil
}

// ---

func amqpExchangeDeclare(ch *amqp.Channel, topic string) error {
	// TODO: Ensure this doesn't drop exchange topics at unexpected times.
	// Args: name, kind, durable, autoDelete, internal, noWait, args
	return ch.ExchangeDeclare(topic, "topic", true, false, false, false, nil)
}

func amqpQueueBind(ch *amqp.Channel, topic, key string) (*amqp.Queue, error) {
	var err error

	// NOTE: We're only doing this when we successfully instantiate connections.
	if err = amqpExchangeDeclare(ch, topic); err != nil {
		return nil, err
	}

	var queue amqp.Queue
	var name = topic + ":" + key

	// Args: name, durable, autoDelete, exclusive, noWait, args
	if queue, err = amqpQueueDeclare(ch, name); err != nil {
		return nil, err
	}

	// Args: name, key, exchange, noWait, args
	if err = ch.QueueBind(name, key, topic, false, nil); err != nil {
		return nil, err
	}

	return &queue, nil
}

// ---

func amqpQueueDeclare(ch *amqp.Channel, queue string) (amqp.Queue, error) {
	// Args: name, durable, autoDelete, exclusive, noWait, args
	return ch.QueueDeclare(queue, true, false, false, false, nil)
}
<|endoftext|>


<commit_before>package lshensemble

import (
	"bytes"
	"testing"
)

func Test_Serialization(t *testing.T) {
	key := "abcd"
	size := 123
	sig := randomSignature(8, int64(1))
	record := &DomainRecord{
		Key:       key,
		Size:      size,
		Signature: sig,
	}
	var data bytes.Buffer
	n, err := record.Write(&data, func(s string) ([]byte, error) {
		return []byte(s), nil
	})
	if err != nil {
		t.Error(err)
	}
	t.Log("Number of bytes written to buffer", n)

	var record2 DomainRecord
	n, err = record2.Read(&data, 4, func(b []byte) (string, error) {
		return string(b), nil
	})
	if err != nil {
		t.Error(err)
	}
	t.Log("Number of bytes read", n)
	t.Log(record2)
	if record2.Key != record.Key || record2.Size != record.Size {
		t.Error("Incorrect record after deserialization")
	}
}
<commit_msg>improv test<commit_after>package lshensemble

import (
	"bytes"
	"testing"
)

func Test_Serialization(t *testing.T) {
	key := "abcd"
	size := 123
	sig := randomSignature(8, int64(1))
	record := &DomainRecord{
		Key:       key,
		Size:      size,
		Signature: sig,
	}
	var data bytes.Buffer
	n, err := record.Write(&data, func(s string) ([]byte, error) {
		return []byte(s), nil
	})
	if err != nil {
		t.Error(err)
	}
	t.Log("Number of bytes written to buffer", n)
	t.Log(record)

	var record2 DomainRecord
	n, err = record2.Read(&data, 4, func(b []byte) (string, error) {
		return string(b), nil
	})
	if err != nil {
		t.Error(err)
	}
	t.Log("Number of bytes read", n)
	t.Log(record2)
	if record2.Key != record.Key || record2.Size != record.Size {
		t.Error("Incorrect record after deserialization")
	}
	for i := range record.Signature {
		if record.Signature[i] != record2.Signature[i] {
			t.Error("Signature does not match after serialization")
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"

	"github.com/crgimenes/goConfig"
)

// Declare config struct

type mongoDB struct {
	Host string `cfgDefault:"example.com"`
	Port int    `cfgDefault:"999"`
}

type systemUser struct {
	Name     string `json:"name" cfg:"name"`
	Password string `json:"passwd" cfg:"passwd"`
}

type configTest struct {
	Domain    string
	User      systemUser `json:"user" cfg:"user"`
	MongoDB   mongoDB    `json:"mongo" cfg:"mongo"`
}

func main() {

	// Instance config struct
	config := configTest{}

	// Adds a prefix to the environment variables.
	goConfig.PrefixEnv = "EXAMPLE"

	// Pass the struct instance pointer to the parser
	err := goConfig.Parse(&config)
	if err != nil {
		println(err)
		return
	}

	// it just print the config struct on the screen
	j, _ := json.MarshalIndent(config, "", "\t")
	println(string(j))
}
<commit_msg>gofmt and add error check<commit_after>package main

import (
	"encoding/json"

	"github.com/crgimenes/goConfig"
)

// Declare config struct

type mongoDB struct {
	Host string `cfgDefault:"example.com"`
	Port int    `cfgDefault:"999"`
}

type systemUser struct {
	Name     string `json:"name" cfg:"name"`
	Password string `json:"passwd" cfg:"passwd"`
}

type configTest struct {
	Domain  string
	User    systemUser `json:"user" cfg:"user"`
	MongoDB mongoDB    `json:"mongo" cfg:"mongo"`
}

func main() {

	// Instance config struct
	config := configTest{}

	// Adds a prefix to the environment variables.
	goConfig.PrefixEnv = "EXAMPLE"

	// Pass the struct instance pointer to the parser
	err := goConfig.Parse(&config)
	if err != nil {
		println(err)
		return
	}

	// it just print the config struct on the screen
	j, err := json.MarshalIndent(config, "", "\t")
	if err != nil {
		println(err)
		return
	}
	println(string(j))
}
<|endoftext|>


<commit_before>package ramsql

import (
	"database/sql/driver"
	"errors"
	"fmt"
	"io"

	"github.com/proullon/ramsql/engine/log"
	"github.com/proullon/ramsql/engine/parser"
)

// Rows implements the sql/driver Rows interface
type Rows struct {
	rowsChannel chan []string
	columns     []string
}

func newRows(channel chan []string) *Rows {
	r := &Rows{rowsChannel: channel}
	c, ok := <-channel
	if !ok {
		log.Critical("Cannot receive column names form channel")
		return nil
	}

	r.columns = c
	return r
}

// Columns returns the names of the columns. The number of
// columns of the result is inferred from the length of the
// slice.  If a particular column name isn't known, an empty
// string should be returned for that entry.
func (r *Rows) Columns() []string {
	return r.columns
}

// Close closes the rows iterator.
func (r *Rows) Close() error {

	if r.rowsChannel == nil {
		return nil
	}

	_, ok := <-r.rowsChannel
	if !ok {
		return nil
	}

	// Tels UnlimitedRowsChannel to close itself
	r.rowsChannel <- []string{}
	return nil
}

// Next is called to populate the next row of data into
// the provided slice. The provided slice will be the same
// size as the Columns() are wide.
//
// The dest slice may be populated only with
// a driver Value type, but excluding string.
// All string values must be converted to []byte.
//
// Next should return io.EOF when there are no more rows.
func (r *Rows) Next(dest []driver.Value) (err error) {

	value, ok := <-r.rowsChannel
	if !ok {
		r.rowsChannel = nil
		return io.EOF
	}

	if len(dest) < len(value) {
		return fmt.Errorf("slice too short (%d slots for %d values)", len(dest), len(value))
	}

	for i, v := range value {
		// TODO: make rowsChannel send virtualRows,
		// so we have the type and don't blindy try to parse date here
		if t, err := parser.ParseDate(string(v)); err == nil {
			dest[i] = *t
		} else {

			dest[i] = []byte(v)
		}
	}

	return nil
}

func (r *Rows) setColumns(columns []string) {
	r.columns = columns
}

func assignValue(s string, v driver.Value) error {
	dest, ok := v.(*string)
	if !ok {
		err := errors.New("cannot assign value")
		return err
	}

	*dest = s
	return nil
}
<commit_msg>fix (rows): add mutex to prevent race condition in rows.Close()<commit_after>package ramsql

import (
	"database/sql/driver"
	"errors"
	"fmt"
	"io"
	"sync"

	"github.com/proullon/ramsql/engine/log"
	"github.com/proullon/ramsql/engine/parser"
)

// Rows implements the sql/driver Rows interface
type Rows struct {
	rowsChannel chan []string
	columns     []string

	sync.Mutex
}

func newRows(channel chan []string) *Rows {
	r := &Rows{rowsChannel: channel}
	c, ok := <-channel
	if !ok {
		log.Critical("Cannot receive column names form channel")
		return nil
	}

	r.columns = c
	return r
}

// Columns returns the names of the columns. The number of
// columns of the result is inferred from the length of the
// slice.  If a particular column name isn't known, an empty
// string should be returned for that entry.
func (r *Rows) Columns() []string {
	return r.columns
}

// Close closes the rows iterator.
func (r *Rows) Close() error {
	r.Lock()
	defer r.Unlock()

	if r.rowsChannel == nil {
		return nil
	}

	_, ok := <-r.rowsChannel
	if !ok {
		return nil
	}

	// Tels UnlimitedRowsChannel to close itself
	//r.rowsChannel <- []string{}
	r.rowsChannel = nil
	return nil
}

// Next is called to populate the next row of data into
// the provided slice. The provided slice will be the same
// size as the Columns() are wide.
//
// The dest slice may be populated only with
// a driver Value type, but excluding string.
// All string values must be converted to []byte.
//
// Next should return io.EOF when there are no more rows.
func (r *Rows) Next(dest []driver.Value) (err error) {
	r.Lock()
	defer r.Unlock()

	if r.rowsChannel == nil {
		return io.EOF
	}

	value, ok := <-r.rowsChannel
	if !ok {
		r.rowsChannel = nil
		return io.EOF
	}

	if len(dest) < len(value) {
		return fmt.Errorf("slice too short (%d slots for %d values)", len(dest), len(value))
	}

	for i, v := range value {
		// TODO: make rowsChannel send virtualRows,
		// so we have the type and don't blindy try to parse date here
		if t, err := parser.ParseDate(string(v)); err == nil {
			dest[i] = *t
		} else {

			dest[i] = []byte(v)
		}
	}

	return nil
}

func (r *Rows) setColumns(columns []string) {
	r.columns = columns
}

func assignValue(s string, v driver.Value) error {
	dest, ok := v.(*string)
	if !ok {
		err := errors.New("cannot assign value")
		return err
	}

	*dest = s
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"log"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/dutchcoders/goftp"
	"github.com/garyburd/redigo/redis"
)

var (
	sshKey    = flag.String("key", "", "SSH key to use for cloning")
	ftpUrl    = flag.String("ftp", "", "FTP server to save backups to")
	redisUrl  = flag.String("redis", "", "Address of redis")
	frequency = flag.Duration("frequency", 24*time.Hour, "Frequency of backups")
	force     = flag.Bool("force", false, "Force download")
	help      = flag.Bool("help", false, "Show this help")
)

var (
	badCharacters = regexp.MustCompilePOSIX("[/@:!?*\\&]")
)

func main() {
	flag.Parse()
	if *help {
		flag.PrintDefaults()
		return
	}
	if *ftpUrl == "" || *redisUrl == "" {
		log.Fatalf("-ftp and -redis have to be set")
	}

	if *sshKey != "" {
		if err := addSshKey(*sshKey); err != nil {
			log.Fatalf("Could not add SSH key: %s", err)
		}
	}

	redisConn, err := connectRedis(*redisUrl)
	if err != nil {
		log.Fatalf("Could not connect to redis: %s", err)
	}
	defer redisConn.Close()

	ftpConn, ftpUrl, err := connectFtp(*ftpUrl)
	if err != nil {
		log.Fatalf("Could not connect to FTP server: %s", err)
	}
	defer ftpConn.Close()
	if err := ftpConn.Cwd(ftpUrl.Path); err != nil {
		log.Fatalf("Could not cd to target directory: %s", err)
	}

	for {
		if !*force {
			nextRun := lastRun(redisConn).Add(*frequency)
			if nextRun.After(time.Now()) {
				time.Sleep(nextRun.Sub(time.Now()))
				continue
			}
		}
		*force = false

		log.Printf("Downloading all the repos...")
		repos := repos(redisConn)
		for _, repo := range repos {
			log.Printf("Downloading %s...", repo)
			safeName := badCharacters.ReplaceAllString(repo, "_")
			buf, err := downloadRepository(repo)
			if err != nil {
				log.Printf("Error downloading repository: %s", err)
				continue
			}

			if err := ftpConn.Stor(safeName+".tar.gz", buf); err != nil {
				log.Printf("Error uploading: %s", err)
			}
		}
		log.Printf("Finished.")
		timestampLastRun(redisConn)
	}
}

func lastRun(conn redis.Conn) time.Time {
	ok, err := redis.Bool(conn.Do("EXISTS", "github-backup:lastrun"))
	if err != nil {
		log.Fatalf("Error querying database: %s", err)
	}
	if !ok {
		return time.Unix(0, 0)
	}

	ts, err := redis.String(conn.Do("GET", "github-backup:lastrun"))
	if err != nil {
		log.Fatalf("Error retrieving timestamp: %s", err)
	}
	t, err := time.Parse(time.RFC3339, ts)
	if err != nil {
		log.Fatalf("Error parsing timestamp: %s", err)
	}
	return t
}

func timestampLastRun(conn redis.Conn) {
	_, err := conn.Do("SET", "github-backup:lastrun", time.Now().Format(time.RFC3339))
	if err != nil {
		log.Fatalf("Error writing timestamp: %s", err)
	}
}

func repos(conn redis.Conn) []string {
	repos, err := redis.Values(conn.Do("LRANGE", "github-backup:repos", 0, 1000))
	if err == redis.ErrNil {
		return []string{}
	}
	if err != nil {
		log.Fatalf("Error retrieving repo list: %s", err)
	}
	r := make([]string, 0, len(repos))
	if err := redis.ScanSlice(repos, &r); err != nil {
		log.Fatalf("Error parsing repo list: %s", err)
	}
	return r
}

func connectRedis(s string) (redis.Conn, error) {
	redisUrl, err := url.Parse(s)
	if err != nil {
		return nil, fmt.Errorf("Could not parse redis url: %s", err)
	}
	if redisUrl.Scheme != "redis" {
		return nil, fmt.Errorf("Unsupported redis scheme %s", redisUrl.Scheme)
	}

	conn, err := redis.Dial("tcp", redisUrl.Host)
	if err != nil {
		return conn, err
	}
	if redisUrl.User != nil {
		pass, ok := redisUrl.User.Password()
		if !ok {
			pass = redisUrl.User.Username()
		}
		_, err := conn.Do("AUTH", pass)
		if err != nil {
			return conn, err
		}
	}
	_, err = conn.Do("EXISTS", "github-backup:lastrun")
	return conn, err
}

func connectFtp(s string) (*goftp.FTP, *url.URL, error) {
	ftpUrl, err := url.Parse(s)
	if err != nil {
		log.Fatalf("Invalid ftp url: %s", err)
	}
	if ftpUrl.Scheme != "ftp" {
		log.Fatalf("Unsupported target scheme %s", ftpUrl.Scheme)
	}
	if !strings.Contains(ftpUrl.Host, ":") {
		ftpUrl.Host += ":21"
	}

	ftp, err := goftp.Connect(ftpUrl.Host)
	if err != nil {
		return ftp, ftpUrl, err
	}
	if ftpUrl.User == nil {
		return ftp, ftpUrl, err
	}
	user := ftpUrl.User.Username()
	pass, _ := ftpUrl.User.Password()
	return ftp, ftpUrl, ftp.Login(user, pass)
}

func downloadRepository(path string) (*bytes.Buffer, error) {
	repo := os.TempDir() + "github-backup"

	if err := os.MkdirAll(repo, os.FileMode(0700)); err != nil {
		return nil, err
	}
	defer os.RemoveAll(repo)

	cmd := exec.Command("git", "clone", "--bare", path)
	cmd.Dir = repo
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return nil, err
	}

	return tarDir(repo)
}

func tarDir(root string) (*bytes.Buffer, error) {
	buf := &bytes.Buffer{}
	gzbuf := gzip.NewWriter(buf)
	defer gzbuf.Close()
	defer gzbuf.Flush()
	archive := tar.NewWriter(gzbuf)
	defer archive.Close()
	defer archive.Flush()
	err := filepath.Walk(root, filepath.WalkFunc(func(path string, info os.FileInfo, err error) error {
		if path == root {
			return nil
		}
		relPath := strings.TrimPrefix(path, root)
		hdr := &tar.Header{
			Name:     strings.TrimPrefix(relPath, "/"),
			Mode:     int64(info.Mode() & os.ModePerm),
			Uid:      1000,
			Gid:      1000,
			Size:     info.Size(),
			Typeflag: tar.TypeReg,
		}
		if info.IsDir() {
			hdr.Typeflag = tar.TypeDir
			hdr.Size = 0
		}

		if err := archive.WriteHeader(hdr); err != nil {
			return err
		}

		if info.IsDir() {
			return nil
		}

		f, err := os.Open(path)
		if err != nil {
			return err
		}
		defer f.Close()
		if _, err := io.Copy(archive, f); err != nil && err != io.EOF {
			return err
		}
		return nil
	}))
	if err != nil {
		return nil, err
	}
	return buf, nil
}

const (
	sshConfig = `
	IdentityFile ~/.ssh/id-rsa
	StrictHostKeyChecking no
	`
)

func writeFile(path string, content []byte) error {
	f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, os.FileMode(0700))
	if err != nil {
		return fmt.Errorf("Error creating file %s: %s", path, err)
	}
	defer f.Close()
	if _, err := f.Write(content); err != nil {
		return fmt.Errorf("Error writing file %s: %s", path, err)
	}
	return nil
}

func addSshKey(encKey string) error {
	key, err := base64.StdEncoding.DecodeString(encKey)
	if err != nil {
		return fmt.Errorf("Error decoding key: %s", err)
	}
	if err := os.MkdirAll("~/.ssh", os.FileMode(0700)); err != nil {
		return fmt.Errorf("Error creating .ssh folder: %s", err)
	}

	if err := writeFile("~/.ssh/id-rsa", key); err != nil {
		return err
	}
	if err := writeFile("~/.ssh/config", []byte(sshConfig)); err != nil {
		return err
	}
	return nil
}
<commit_msg>Use $HOME explicitly<commit_after>package main

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"log"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/dutchcoders/goftp"
	"github.com/garyburd/redigo/redis"
)

var (
	sshKey    = flag.String("key", "", "SSH key to use for cloning")
	ftpUrl    = flag.String("ftp", "", "FTP server to save backups to")
	redisUrl  = flag.String("redis", "", "Address of redis")
	frequency = flag.Duration("frequency", 24*time.Hour, "Frequency of backups")
	force     = flag.Bool("force", false, "Force download")
	help      = flag.Bool("help", false, "Show this help")
)

var (
	badCharacters = regexp.MustCompilePOSIX("[/@:!?*\\&]")
)

func main() {
	flag.Parse()
	if *help {
		flag.PrintDefaults()
		return
	}
	if *ftpUrl == "" || *redisUrl == "" {
		log.Fatalf("-ftp and -redis have to be set")
	}

	if *sshKey != "" {
		if err := addSshKey(*sshKey); err != nil {
			log.Fatalf("Could not add SSH key: %s", err)
		}
	}

	redisConn, err := connectRedis(*redisUrl)
	if err != nil {
		log.Fatalf("Could not connect to redis: %s", err)
	}
	defer redisConn.Close()

	ftpConn, ftpUrl, err := connectFtp(*ftpUrl)
	if err != nil {
		log.Fatalf("Could not connect to FTP server: %s", err)
	}
	defer ftpConn.Close()
	if err := ftpConn.Cwd(ftpUrl.Path); err != nil {
		log.Fatalf("Could not cd to target directory: %s", err)
	}

	for {
		if !*force {
			nextRun := lastRun(redisConn).Add(*frequency)
			if nextRun.After(time.Now()) {
				time.Sleep(nextRun.Sub(time.Now()))
				continue
			}
		}
		*force = false

		log.Printf("Downloading all the repos...")
		repos := repos(redisConn)
		for _, repo := range repos {
			log.Printf("Downloading %s...", repo)
			safeName := badCharacters.ReplaceAllString(repo, "_")
			buf, err := downloadRepository(repo)
			if err != nil {
				log.Printf("Error downloading repository: %s", err)
				continue
			}

			if err := ftpConn.Stor(safeName+".tar.gz", buf); err != nil {
				log.Printf("Error uploading: %s", err)
			}
		}
		log.Printf("Finished.")
		timestampLastRun(redisConn)
	}
}

func lastRun(conn redis.Conn) time.Time {
	ok, err := redis.Bool(conn.Do("EXISTS", "github-backup:lastrun"))
	if err != nil {
		log.Fatalf("Error querying database: %s", err)
	}
	if !ok {
		return time.Unix(0, 0)
	}

	ts, err := redis.String(conn.Do("GET", "github-backup:lastrun"))
	if err != nil {
		log.Fatalf("Error retrieving timestamp: %s", err)
	}
	t, err := time.Parse(time.RFC3339, ts)
	if err != nil {
		log.Fatalf("Error parsing timestamp: %s", err)
	}
	return t
}

func timestampLastRun(conn redis.Conn) {
	_, err := conn.Do("SET", "github-backup:lastrun", time.Now().Format(time.RFC3339))
	if err != nil {
		log.Fatalf("Error writing timestamp: %s", err)
	}
}

func repos(conn redis.Conn) []string {
	repos, err := redis.Values(conn.Do("LRANGE", "github-backup:repos", 0, 1000))
	if err == redis.ErrNil {
		return []string{}
	}
	if err != nil {
		log.Fatalf("Error retrieving repo list: %s", err)
	}
	r := make([]string, 0, len(repos))
	if err := redis.ScanSlice(repos, &r); err != nil {
		log.Fatalf("Error parsing repo list: %s", err)
	}
	return r
}

func connectRedis(s string) (redis.Conn, error) {
	redisUrl, err := url.Parse(s)
	if err != nil {
		return nil, fmt.Errorf("Could not parse redis url: %s", err)
	}
	if redisUrl.Scheme != "redis" {
		return nil, fmt.Errorf("Unsupported redis scheme %s", redisUrl.Scheme)
	}

	conn, err := redis.Dial("tcp", redisUrl.Host)
	if err != nil {
		return conn, err
	}
	if redisUrl.User != nil {
		pass, ok := redisUrl.User.Password()
		if !ok {
			pass = redisUrl.User.Username()
		}
		_, err := conn.Do("AUTH", pass)
		if err != nil {
			return conn, err
		}
	}
	_, err = conn.Do("EXISTS", "github-backup:lastrun")
	return conn, err
}

func connectFtp(s string) (*goftp.FTP, *url.URL, error) {
	ftpUrl, err := url.Parse(s)
	if err != nil {
		log.Fatalf("Invalid ftp url: %s", err)
	}
	if ftpUrl.Scheme != "ftp" {
		log.Fatalf("Unsupported target scheme %s", ftpUrl.Scheme)
	}
	if !strings.Contains(ftpUrl.Host, ":") {
		ftpUrl.Host += ":21"
	}

	ftp, err := goftp.Connect(ftpUrl.Host)
	if err != nil {
		return ftp, ftpUrl, err
	}
	if ftpUrl.User == nil {
		return ftp, ftpUrl, err
	}
	user := ftpUrl.User.Username()
	pass, _ := ftpUrl.User.Password()
	return ftp, ftpUrl, ftp.Login(user, pass)
}

func downloadRepository(path string) (*bytes.Buffer, error) {
	repo := os.TempDir() + "github-backup"

	if err := os.MkdirAll(repo, os.FileMode(0700)); err != nil {
		return nil, err
	}
	defer os.RemoveAll(repo)

	cmd := exec.Command("git", "clone", "--bare", path)
	cmd.Dir = repo
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return nil, err
	}

	return tarDir(repo)
}

func tarDir(root string) (*bytes.Buffer, error) {
	buf := &bytes.Buffer{}
	gzbuf := gzip.NewWriter(buf)
	defer gzbuf.Close()
	defer gzbuf.Flush()
	archive := tar.NewWriter(gzbuf)
	defer archive.Close()
	defer archive.Flush()
	err := filepath.Walk(root, filepath.WalkFunc(func(path string, info os.FileInfo, err error) error {
		if path == root {
			return nil
		}
		relPath := strings.TrimPrefix(path, root)
		hdr := &tar.Header{
			Name:     strings.TrimPrefix(relPath, "/"),
			Mode:     int64(info.Mode() & os.ModePerm),
			Uid:      1000,
			Gid:      1000,
			Size:     info.Size(),
			Typeflag: tar.TypeReg,
		}
		if info.IsDir() {
			hdr.Typeflag = tar.TypeDir
			hdr.Size = 0
		}

		if err := archive.WriteHeader(hdr); err != nil {
			return err
		}

		if info.IsDir() {
			return nil
		}

		f, err := os.Open(path)
		if err != nil {
			return err
		}
		defer f.Close()
		if _, err := io.Copy(archive, f); err != nil && err != io.EOF {
			return err
		}
		return nil
	}))
	if err != nil {
		return nil, err
	}
	return buf, nil
}

const (
	sshConfig = `
	IdentityFile ~/.ssh/id-rsa
	StrictHostKeyChecking no
	`
)

func writeFile(path string, content []byte) error {
	f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, os.FileMode(0700))
	if err != nil {
		return fmt.Errorf("Error creating file %s: %s", path, err)
	}
	defer f.Close()
	if _, err := f.Write(content); err != nil {
		return fmt.Errorf("Error writing file %s: %s", path, err)
	}
	return nil
}

func addSshKey(encKey string) error {
	key, err := base64.StdEncoding.DecodeString(encKey)
	if err != nil {
		return fmt.Errorf("Error decoding key: %s", err)
	}
	if err := os.MkdirAll(os.ExpandEnv("$HOME/.ssh"), os.FileMode(0700)); err != nil {
		return fmt.Errorf("Error creating .ssh folder: %s", err)
	}

	if err := writeFile(os.ExpandEnv("$HOME/.ssh/id-rsa"), key); err != nil {
		return err
	}
	if err := writeFile(os.ExpandEnv("$HOME/.ssh/config"), []byte(sshConfig)); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package dpn_test

import (
	"bufio"
	"fmt"
	"github.com/APTrust/bagman/bagman"
	"github.com/APTrust/bagman/dpn"
	"io"
	"net/url"
	"os"
	"os/user"
	"path/filepath"
	"strings"
	"testing"
)

var testConfig string = "test"

// These bags are on the DPN test server,
// and we have data for them in our local DPN REST
// service (they're in the fixture file TestServerData.json).
var TEST_BAGS = []string {
	"005e7793-8253-4585-7118-2da702e29aa0",
	"06edc30f-af04-4a4e-4ad6-3c679a195d46",
	"41e5376c-cc13-4c3e-6af3-297cc2e005aa",
	"472218b3-95ce-4b8e-6c21-6e514cfbe43f",
	"4d11736c-c0ab-44b0-66c6-a947f414d4f1",
	"8839f294-c7c0-404d-4467-7ab223ff461b",
}

var skipCopyMessagePrinted = false
var entryChecked = false
var hasConfigEntry = false

// Test to see if the current user has an ssh config entry for
// dpn-test. We need this to run the copier tests, because
// copier uses rsync over ssh.
func hasSSHConfigEntry() (bool) {
	if entryChecked == true {
		return hasConfigEntry
	}
	entryChecked = true
	usr, _ := user.Current()
	sshConfigFile := filepath.Join(usr.HomeDir, ".ssh", "config")
	f, err := os.Open(sshConfigFile)
	if err != nil {
		return false
	}
	defer f.Close()
	r := bufio.NewReader(f)
	for {
		line, err := r.ReadString('\n')
		if err == io.EOF {
			break
		} else if err != nil {
			break
		}
		if strings.HasPrefix(strings.TrimSpace(line), "Host dpn-test") {
			hasConfigEntry = true
		}
	}
	return hasConfigEntry
}

// Return true/false indicating whether we should try to run
// copy tests. Print a message if we can't run the tests.
func canRunCopyTests(t *testing.T) (bool) {
	ok := hasSSHConfigEntry() && runRestTests(t)
	if ok == false && skipCopyMessagePrinted == false {
		skipCopyMessagePrinted = true
		fmt.Println("**** Skipping DPN copy integration tests: "+
			"No dpn-test entry in ~/.ssh/config and/or local " +
			"DPN REST server is not running")
	}
	return ok
}

// For this to work, you need to have an account that
// can access the dpn test server, and you need a config
// entry in ~/.ssh/config with settings to connect to
// the dpn-test server.
func getTestLink(tarredBagName string) (string) {
	return fmt.Sprintf("dpn-test:/home/earthdiver/staging/%s", tarredBagName)
}

// This builds a DPNResult suitable for feeding to the Copier.RunTest()
// method. We get a real bag record and a real transfer record from the
// local DPN REST server, which should be loaded with the fixture data
// from TestServerData.json. We change the rsync link on the transfer
// request to point to the dpn-test server. The copy test will try to
// pull the tarred bag from that server.
func buildTestResult(bagIdentifier string, t *testing.T) (*dpn.DPNResult) {
	// Build a result object with a DPN Bag...
	result := dpn.NewDPNResult(bagIdentifier)

	localRestClient := getClient(t)
	dpnBag, err := localRestClient.DPNBagGet(bagIdentifier)
	if err != nil {
		t.Errorf("Can't get DPN bag '%s' from local DPN REST service: %v",
			bagIdentifier, err)
		return nil
	}
	result.DPNBag = dpnBag

	// Get a transfer request for this bag...
	params := &url.Values{}
	params.Set("uuid", bagIdentifier)
	xferRequests, err := localRestClient.DPNReplicationListGet(params)
	if err != nil {
		t.Errorf("Can't get transfer request for bag '%s' from local DPN REST service: %v",
			bagIdentifier, err)
		return nil
	}
	if len(xferRequests.Results) == 0 {
		t.Errorf("No transfer requests for bag '%s' in local DPN REST service",
			bagIdentifier)
		return nil
	}
	result.TransferRequest = xferRequests.Results[0]

	// Change the rsync link for the bag to point toward
	// our dpn test server.
	tarredBagName := fmt.Sprintf("%s.tar", bagIdentifier)
	result.TransferRequest.Link = getTestLink(tarredBagName)

	return result
}

// Delete rsynched files after testing
func copyTestCleanup() {
	procUtil := bagman.NewProcessUtil(&testConfig)
	for _, uuid := range TEST_BAGS {
		filePath := filepath.Join(procUtil.Config.DPNStagingDirectory, uuid + ".tar")
		os.Remove(filePath)
	}
}

func TestGetRsyncCommand(t *testing.T) {
	procUtil := bagman.NewProcessUtil(&testConfig)
	copyFrom := getTestLink(TEST_BAGS[0])
	copyTo := procUtil.Config.DPNStagingDirectory
	command := dpn.GetRsyncCommand(copyFrom, copyTo)
	if !strings.HasSuffix(command.Path, "rsync") {
		t.Errorf("Expected Path ending in 'rsync', got '%s'", command.Path)
	}
	if len(command.Args) < 6 {
		t.Errorf("rsync command has %d args, expected 5", len(command.Args))
		return
	}
	if command.Args[3] != "ssh" {
		t.Errorf("rsync command should be using ssh, but it's not")
	}
	if command.Args[4] != copyFrom {
		t.Errorf("rsync command is copying from '%s', expected '%s'",
			command.Args[4], copyFrom)
	}
	if command.Args[5] != copyTo {
		t.Errorf("rsync command is copying to '%s', expected '%s'",
			command.Args[5], copyTo)
	}
}

func TestCopier(t *testing.T) {
	if canRunCopyTests(t) == false {
		return
	}
	// runRestTests is defined in dpnrestclient_test.go
	if runRestTests(t) == false {
		return
	}

	procUtil := bagman.NewProcessUtil(&testConfig)
	dpnConfig := loadConfig(t, configFile)
	copier, err := dpn.NewCopier(procUtil, dpnConfig)
	if err != nil {
		t.Error(err)
		return
	}

	// Get ridda that shizzle
	defer copyTestCleanup()

	for _, uuid := range TEST_BAGS {
		dpnResult := buildTestResult(uuid, t)
		if dpnResult == nil {
			return
		}

		// RunTest will update DPNResult.CopyResult
		copier.RunTest(dpnResult)

		if dpnResult.CopyResult.BagWasCopied == false {
			if dpnResult.CopyResult.ErrorMessage != "" {
				t.Errorf("Error copying bag %s: %s",
					uuid, dpnResult.CopyResult.ErrorMessage)
			} else if dpnResult.CopyResult.InfoMessage != "" {
				t.Errorf("Bag %s was not copied: %s",
					uuid, dpnResult.CopyResult.InfoMessage)
			} else {
				// This shouldn't happen
				t.Errorf("Bag %s was not copied. Reason unknown.", uuid)
			}
		}
		if !bagman.FileExists(dpnResult.CopyResult.LocalPath) {
			t.Errorf("Bag %s was reported copied to %s, " +
				"but that file does not exist", uuid,
				dpnResult.CopyResult.LocalPath)
		}
		if dpnResult.DPNBag.Fixities.Sha256 != dpnResult.BagSha256Digest {
			t.Errorf("Fixity did not match for bag %s. Expected %s, " +
				"got %s", uuid, dpnResult.DPNBag.Fixities.Sha256,
				dpnResult.BagSha256Digest)
		}
		if len(dpnResult.BagMd5Digest) == 0 {
			t.Errorf("Bg MD5 digest is missing.")
		}
		if dpnResult.BagSize == 0 {
			t.Errorf("Bag size is missing")
		}
	}
}
<commit_msg>Replaced old test IDs with new ones<commit_after>package dpn_test

import (
	"bufio"
	"fmt"
	"github.com/APTrust/bagman/bagman"
	"github.com/APTrust/bagman/dpn"
	"io"
	"net/url"
	"os"
	"os/user"
	"path/filepath"
	"strings"
	"testing"
)

var testConfig string = "test"

// These bags are on the DPN test server,
// and we have data for them in our local DPN REST
// service (they're in the fixture file TestServerData.json).
var TEST_BAGS = []string {
	"10000000-0000-4000-a000-000000000001",
	"10000000-0000-4000-a000-000000000002",
	"10000000-0000-4000-a000-000000000003",
	"10000000-0000-4000-a000-000000000004",
	"10000000-0000-4000-a000-000000000005",
	"10000000-0000-4000-a000-000000000006",
}

var skipCopyMessagePrinted = false
var entryChecked = false
var hasConfigEntry = false

// Test to see if the current user has an ssh config entry for
// dpn-test. We need this to run the copier tests, because
// copier uses rsync over ssh.
func hasSSHConfigEntry() (bool) {
	if entryChecked == true {
		return hasConfigEntry
	}
	entryChecked = true
	usr, _ := user.Current()
	sshConfigFile := filepath.Join(usr.HomeDir, ".ssh", "config")
	f, err := os.Open(sshConfigFile)
	if err != nil {
		return false
	}
	defer f.Close()
	r := bufio.NewReader(f)
	for {
		line, err := r.ReadString('\n')
		if err == io.EOF {
			break
		} else if err != nil {
			break
		}
		if strings.HasPrefix(strings.TrimSpace(line), "Host dpn-test") {
			hasConfigEntry = true
		}
	}
	return hasConfigEntry
}

// Return true/false indicating whether we should try to run
// copy tests. Print a message if we can't run the tests.
func canRunCopyTests(t *testing.T) (bool) {
	ok := hasSSHConfigEntry() && runRestTests(t)
	if ok == false && skipCopyMessagePrinted == false {
		skipCopyMessagePrinted = true
		fmt.Println("**** Skipping DPN copy integration tests: "+
			"No dpn-test entry in ~/.ssh/config and/or local " +
			"DPN REST server is not running")
	}
	return ok
}

// For this to work, you need to have an account that
// can access the dpn test server, and you need a config
// entry in ~/.ssh/config with settings to connect to
// the dpn-test server.
func getTestLink(tarredBagName string) (string) {
	return fmt.Sprintf("dpn-test:/home/earthdiver/staging/%s", tarredBagName)
}

// This builds a DPNResult suitable for feeding to the Copier.RunTest()
// method. We get a real bag record and a real transfer record from the
// local DPN REST server, which should be loaded with the fixture data
// from TestServerData.json. We change the rsync link on the transfer
// request to point to the dpn-test server. The copy test will try to
// pull the tarred bag from that server.
func buildTestResult(bagIdentifier string, t *testing.T) (*dpn.DPNResult) {
	// Build a result object with a DPN Bag...
	result := dpn.NewDPNResult(bagIdentifier)

	localRestClient := getClient(t)
	dpnBag, err := localRestClient.DPNBagGet(bagIdentifier)
	if err != nil {
		t.Errorf("Can't get DPN bag '%s' from local DPN REST service: %v",
			bagIdentifier, err)
		return nil
	}
	result.DPNBag = dpnBag

	// Get a transfer request for this bag...
	params := &url.Values{}
	params.Set("uuid", bagIdentifier)
	xferRequests, err := localRestClient.DPNReplicationListGet(params)
	if err != nil {
		t.Errorf("Can't get transfer request for bag '%s' from local DPN REST service: %v",
			bagIdentifier, err)
		return nil
	}
	if len(xferRequests.Results) == 0 {
		t.Errorf("No transfer requests for bag '%s' in local DPN REST service",
			bagIdentifier)
		return nil
	}
	result.TransferRequest = xferRequests.Results[0]

	// Change the rsync link for the bag to point toward
	// our dpn test server.
	tarredBagName := fmt.Sprintf("%s.tar", bagIdentifier)
	result.TransferRequest.Link = getTestLink(tarredBagName)

	return result
}

// Delete rsynched files after testing
func copyTestCleanup() {
	procUtil := bagman.NewProcessUtil(&testConfig)
	for _, uuid := range TEST_BAGS {
		filePath := filepath.Join(procUtil.Config.DPNStagingDirectory, uuid + ".tar")
		os.Remove(filePath)
	}
}

func TestGetRsyncCommand(t *testing.T) {
	procUtil := bagman.NewProcessUtil(&testConfig)
	copyFrom := getTestLink(TEST_BAGS[0])
	copyTo := procUtil.Config.DPNStagingDirectory
	command := dpn.GetRsyncCommand(copyFrom, copyTo)
	if !strings.HasSuffix(command.Path, "rsync") {
		t.Errorf("Expected Path ending in 'rsync', got '%s'", command.Path)
	}
	if len(command.Args) < 6 {
		t.Errorf("rsync command has %d args, expected 5", len(command.Args))
		return
	}
	if command.Args[3] != "ssh" {
		t.Errorf("rsync command should be using ssh, but it's not")
	}
	if command.Args[4] != copyFrom {
		t.Errorf("rsync command is copying from '%s', expected '%s'",
			command.Args[4], copyFrom)
	}
	if command.Args[5] != copyTo {
		t.Errorf("rsync command is copying to '%s', expected '%s'",
			command.Args[5], copyTo)
	}
}

func TestCopier(t *testing.T) {
	if canRunCopyTests(t) == false {
		return
	}
	// runRestTests is defined in dpnrestclient_test.go
	if runRestTests(t) == false {
		return
	}

	procUtil := bagman.NewProcessUtil(&testConfig)
	dpnConfig := loadConfig(t, configFile)
	copier, err := dpn.NewCopier(procUtil, dpnConfig)
	if err != nil {
		t.Error(err)
		return
	}

	// Get ridda that shizzle
	defer copyTestCleanup()

	for _, uuid := range TEST_BAGS {
		dpnResult := buildTestResult(uuid, t)
		if dpnResult == nil {
			return
		}

		// RunTest will update DPNResult.CopyResult
		copier.RunTest(dpnResult)

		if dpnResult.CopyResult.BagWasCopied == false {
			if dpnResult.CopyResult.ErrorMessage != "" {
				t.Errorf("Error copying bag %s: %s",
					uuid, dpnResult.CopyResult.ErrorMessage)
			} else if dpnResult.CopyResult.InfoMessage != "" {
				t.Errorf("Bag %s was not copied: %s",
					uuid, dpnResult.CopyResult.InfoMessage)
			} else {
				// This shouldn't happen
				t.Errorf("Bag %s was not copied. Reason unknown.", uuid)
			}
		}
		if !bagman.FileExists(dpnResult.CopyResult.LocalPath) {
			t.Errorf("Bag %s was reported copied to %s, " +
				"but that file does not exist", uuid,
				dpnResult.CopyResult.LocalPath)
		}
		if dpnResult.DPNBag.Fixities.Sha256 != dpnResult.BagSha256Digest {
			t.Errorf("Fixity did not match for bag %s. Expected %s, " +
				"got %s", uuid, dpnResult.DPNBag.Fixities.Sha256,
				dpnResult.BagSha256Digest)
		}
		if len(dpnResult.BagMd5Digest) == 0 {
			t.Errorf("Bg MD5 digest is missing.")
		}
		if dpnResult.BagSize == 0 {
			t.Errorf("Bag size is missing")
		}
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Mirantis

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 flexvolume

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/Mirantis/virtlet/pkg/utils"
)

const (
	uuidOptionsKey     = "uuid"
	partOptionsKey     = "part"
	flexvolumeDataFile = "virtlet-flexvolume.json"
)

// flexVolumeDebug indicates whether flexvolume debugging should be enabled
var flexVolumeDebug = false

func init() {
	// XXX: invent a better way to decide whether debugging should
	// be used for flexvolume driver. For now we only enable it if
	// Docker-in-Docker env is used
	if fi, err := os.Stat("/dind/flexvolume_driver"); err == nil && !fi.IsDir() {
		flexVolumeDebug = true
	}
}

type Mounter interface {
	Mount(source string, target string, fstype string) error
	Unmount(target string) error
}

type nullMounter struct{}

func (m *nullMounter) Mount(source string, target string, fstype string) error {
	return nil
}

func (m *nullMounter) Unmount(target string) error {
	return nil
}

var NullMounter = &nullMounter{}

type UuidGen func() string

type FlexVolumeDriver struct {
	uuidGen UuidGen
	mounter Mounter
}

func NewFlexVolumeDriver(uuidGen UuidGen, mounter Mounter) *FlexVolumeDriver {
	return &FlexVolumeDriver{uuidGen: uuidGen, mounter: mounter}
}

func (d *FlexVolumeDriver) populateVolumeDir(targetDir string, opts map[string]interface{}) error {
	return utils.WriteJSON(filepath.Join(targetDir, flexvolumeDataFile), opts, 0700)
}

// The following functions are not currently needed, but still
// keeping them to make it easier to actually implement them

// Invocation: <driver executable> init
func (d *FlexVolumeDriver) init() (map[string]interface{}, error) {
	return nil, nil
}

// Invocation: <driver executable> attach <json options> <node name>
func (d *FlexVolumeDriver) attach(jsonOptions, nodeName string) (map[string]interface{}, error) {
	return nil, nil
}

// Invocation: <driver executable> detach <mount device> <node name>
func (d *FlexVolumeDriver) detach(mountDev, nodeName string) (map[string]interface{}, error) {
	return nil, nil
}

// Invocation: <driver executable> waitforattach <mount device> <json options>
func (d *FlexVolumeDriver) waitForAttach(mountDev, jsonOptions string) (map[string]interface{}, error) {
	return map[string]interface{}{"device": mountDev}, nil
}

// Invocation: <driver executable> isattached <json options> <node name>
func (d *FlexVolumeDriver) isAttached(jsonOptions, nodeName string) (map[string]interface{}, error) {
	return map[string]interface{}{"attached": true}, nil
}

//Invocation: <driver executable> mount <target mount dir> <mount device> <json options>
func (d *FlexVolumeDriver) mount(targetMountDir, jsonOptions string) (map[string]interface{}, error) {
	var opts map[string]interface{}
	if err := json.Unmarshal([]byte(jsonOptions), &opts); err != nil {
		return nil, fmt.Errorf("failed to unmarshal json options: %v", err)
	}
	opts[uuidOptionsKey] = d.uuidGen()
	if err := os.MkdirAll(targetMountDir, 0700); err != nil {
		return nil, fmt.Errorf("os.MkDirAll(): %v", err)
	}

	// Here we populate the volume directory twice.
	// That's because we need to do tmpfs mount to make kubelet happy -
	// it will not try to unmount the volume if it doesn't detect mount
	// point on the flexvolume dir, and using 'mount --bind' is not enough
	// because kubelet's mount point detection doesn't see bind mounts.
	// The problem is that hostPaths are mounted as private (no mount
	// propagation) and so tmpfs content is not visible inside Virtlet
	// pod. So, in order to avoid unneeded confusion down the road,
	// we place flexvolume contents to the volume dir twice,
	// both directly and onto the freshly mounted tmpfs.

	if err := d.populateVolumeDir(targetMountDir, opts); err != nil {
		return nil, err
	}

	if err := d.mounter.Mount("tmpfs", targetMountDir, "tmpfs"); err != nil {
		return nil, fmt.Errorf("error mounting tmpfs at %q: %v", targetMountDir, err)
	}

	done := false
	defer func() {
		// try to unmount upon error or panic
		if !done {
			d.mounter.Unmount(targetMountDir)
		}
	}()

	if err := d.populateVolumeDir(targetMountDir, opts); err != nil {
		return nil, err
	}

	done = true
	return nil, nil
}

// Invocation: <driver executable> unmount <mount dir>
func (d *FlexVolumeDriver) unmount(targetMountDir string) (map[string]interface{}, error) {
	if err := d.mounter.Unmount(targetMountDir); err != nil {
		return nil, fmt.Errorf("unmount %q: %v", targetMountDir, err.Error())
	}

	if err := os.RemoveAll(targetMountDir); err != nil {
		return nil, fmt.Errorf("os.RemoveAll(): %v", err.Error())
	}

	return nil, nil
}

type driverOp func(*FlexVolumeDriver, []string) (map[string]interface{}, error)

type cmdInfo struct {
	numArgs int
	run     driverOp
}

var commands = map[string]cmdInfo{
	"init": {
		0, func(d *FlexVolumeDriver, args []string) (map[string]interface{}, error) {
			return d.init()
		},
	},
	"attach": {
		2, func(d *FlexVolumeDriver, args []string) (map[string]interface{}, error) {
			return d.attach(args[0], args[1])
		},
	},
	"detach": {
		2, func(d *FlexVolumeDriver, args []string) (map[string]interface{}, error) {
			return d.detach(args[0], args[1])
		},
	},
	"waitforattach": {
		2, func(d *FlexVolumeDriver, args []string) (map[string]interface{}, error) {
			return d.waitForAttach(args[0], args[1])
		},
	},
	"isattached": {
		2, func(d *FlexVolumeDriver, args []string) (map[string]interface{}, error) {
			return d.isAttached(args[0], args[1])
		},
	},
	"mount": {
		2, func(d *FlexVolumeDriver, args []string) (map[string]interface{}, error) {
			return d.mount(args[0], args[1])
		},
	},
	"unmount": {
		1, func(d *FlexVolumeDriver, args []string) (map[string]interface{}, error) {
			return d.unmount(args[0])
		},
	},
}

func (d *FlexVolumeDriver) doRun(args []string) (map[string]interface{}, error) {
	if len(args) == 0 {
		return nil, errors.New("no arguments passed to flexvolume driver")
	}
	nArgs := len(args) - 1
	op := args[0]
	if cmdInfo, found := commands[op]; found {
		if cmdInfo.numArgs == nArgs {
			return cmdInfo.run(d, args[1:])
		} else {
			return nil, fmt.Errorf("unexpected number of args %d (expected %d) for operation %q", nArgs, cmdInfo.numArgs, op)
		}
	} else {
		return map[string]interface{}{
			"status": "Not supported",
		}, nil
	}
}

func (d *FlexVolumeDriver) Run(args []string) string {
	r := formatResult(d.doRun(args))

	if flexVolumeDebug {
		// This is for debugging purposes only.
		// The problem is that kubelet grabs CombinedOutput() from the process
		// and tries to parse it as JSON (need to recheck this,
		// maybe submit a PS to fix it)
		f, err := os.OpenFile("/tmp/flexvolume.log", os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0666)
		if err == nil {
			defer f.Close()
			fmt.Fprintf(f, "flexvolume %s -> %s\n", strings.Join(args, " "), r)
		}
	}

	return r
}

func formatResult(fields map[string]interface{}, err error) string {
	var data map[string]interface{}
	if err != nil {
		data = map[string]interface{}{
			"status":  "Failure",
			"message": err.Error(),
		}
	} else {
		data = map[string]interface{}{
			"status": "Success",
		}
		for k, v := range fields {
			data[k] = v
		}
	}
	s, err := json.Marshal(data)
	if err != nil {
		panic("error marshalling the data")
	}
	return string(s) + "\n"
}

// GetFlexvolumeInfo tries to extract flexvolume uuid and partition
// number from the specified directory. Negative partition number
// means that no partition number was specified.
func GetFlexvolumeInfo(dir string) (string, int, error) {
	dataFile := filepath.Join(dir, flexvolumeDataFile)
	var opts map[string]interface{}
	if err := utils.ReadJSON(dataFile, &opts); err != nil {
		return "", 0, fmt.Errorf("can't read flexvolume data file %q: %v", dataFile, err)
	}
	uuidRaw, found := opts[uuidOptionsKey]
	if !found {
		return "", 0, fmt.Errorf("%q: flexvolume doesn't have an uuid", dataFile)
	}
	uuid, ok := uuidRaw.(string)
	if !ok {
		return "", 0, fmt.Errorf("%q: flexvolume uuid is not a string", dataFile)
	}
	part := -1
	partRaw, found := opts[partOptionsKey]
	if found {
		partStr, ok := partRaw.(string)
		if !ok {
			return "", 0, fmt.Errorf("%q: 'part' value is not a string", dataFile)
		}
		var err error
		part, err = strconv.Atoi(partStr)
		if err != nil {
			return "", 0, fmt.Errorf("%q: malformed 'part' value (partition number): %q: %v", dataFile, partRaw, err)
		}
	}
	return uuid, part, nil
}
<commit_msg>incorrect comment for mount of flexvolume<commit_after>/*
Copyright 2017 Mirantis

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 flexvolume

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/Mirantis/virtlet/pkg/utils"
)

const (
	uuidOptionsKey     = "uuid"
	partOptionsKey     = "part"
	flexvolumeDataFile = "virtlet-flexvolume.json"
)

// flexVolumeDebug indicates whether flexvolume debugging should be enabled
var flexVolumeDebug = false

func init() {
	// XXX: invent a better way to decide whether debugging should
	// be used for flexvolume driver. For now we only enable it if
	// Docker-in-Docker env is used
	if fi, err := os.Stat("/dind/flexvolume_driver"); err == nil && !fi.IsDir() {
		flexVolumeDebug = true
	}
}

type Mounter interface {
	Mount(source string, target string, fstype string) error
	Unmount(target string) error
}

type nullMounter struct{}

func (m *nullMounter) Mount(source string, target string, fstype string) error {
	return nil
}

func (m *nullMounter) Unmount(target string) error {
	return nil
}

var NullMounter = &nullMounter{}

type UuidGen func() string

type FlexVolumeDriver struct {
	uuidGen UuidGen
	mounter Mounter
}

func NewFlexVolumeDriver(uuidGen UuidGen, mounter Mounter) *FlexVolumeDriver {
	return &FlexVolumeDriver{uuidGen: uuidGen, mounter: mounter}
}

func (d *FlexVolumeDriver) populateVolumeDir(targetDir string, opts map[string]interface{}) error {
	return utils.WriteJSON(filepath.Join(targetDir, flexvolumeDataFile), opts, 0700)
}

// The following functions are not currently needed, but still
// keeping them to make it easier to actually implement them

// Invocation: <driver executable> init
func (d *FlexVolumeDriver) init() (map[string]interface{}, error) {
	return nil, nil
}

// Invocation: <driver executable> attach <json options> <node name>
func (d *FlexVolumeDriver) attach(jsonOptions, nodeName string) (map[string]interface{}, error) {
	return nil, nil
}

// Invocation: <driver executable> detach <mount device> <node name>
func (d *FlexVolumeDriver) detach(mountDev, nodeName string) (map[string]interface{}, error) {
	return nil, nil
}

// Invocation: <driver executable> waitforattach <mount device> <json options>
func (d *FlexVolumeDriver) waitForAttach(mountDev, jsonOptions string) (map[string]interface{}, error) {
	return map[string]interface{}{"device": mountDev}, nil
}

// Invocation: <driver executable> isattached <json options> <node name>
func (d *FlexVolumeDriver) isAttached(jsonOptions, nodeName string) (map[string]interface{}, error) {
	return map[string]interface{}{"attached": true}, nil
}

//Invocation: <driver executable> mount <target mount dir> <json options>
func (d *FlexVolumeDriver) mount(targetMountDir, jsonOptions string) (map[string]interface{}, error) {
	var opts map[string]interface{}
	if err := json.Unmarshal([]byte(jsonOptions), &opts); err != nil {
		return nil, fmt.Errorf("failed to unmarshal json options: %v", err)
	}
	opts[uuidOptionsKey] = d.uuidGen()
	if err := os.MkdirAll(targetMountDir, 0700); err != nil {
		return nil, fmt.Errorf("os.MkDirAll(): %v", err)
	}

	// Here we populate the volume directory twice.
	// That's because we need to do tmpfs mount to make kubelet happy -
	// it will not try to unmount the volume if it doesn't detect mount
	// point on the flexvolume dir, and using 'mount --bind' is not enough
	// because kubelet's mount point detection doesn't see bind mounts.
	// The problem is that hostPaths are mounted as private (no mount
	// propagation) and so tmpfs content is not visible inside Virtlet
	// pod. So, in order to avoid unneeded confusion down the road,
	// we place flexvolume contents to the volume dir twice,
	// both directly and onto the freshly mounted tmpfs.

	if err := d.populateVolumeDir(targetMountDir, opts); err != nil {
		return nil, err
	}

	if err := d.mounter.Mount("tmpfs", targetMountDir, "tmpfs"); err != nil {
		return nil, fmt.Errorf("error mounting tmpfs at %q: %v", targetMountDir, err)
	}

	done := false
	defer func() {
		// try to unmount upon error or panic
		if !done {
			d.mounter.Unmount(targetMountDir)
		}
	}()

	if err := d.populateVolumeDir(targetMountDir, opts); err != nil {
		return nil, err
	}

	done = true
	return nil, nil
}

// Invocation: <driver executable> unmount <mount dir>
func (d *FlexVolumeDriver) unmount(targetMountDir string) (map[string]interface{}, error) {
	if err := d.mounter.Unmount(targetMountDir); err != nil {
		return nil, fmt.Errorf("unmount %q: %v", targetMountDir, err.Error())
	}

	if err := os.RemoveAll(targetMountDir); err != nil {
		return nil, fmt.Errorf("os.RemoveAll(): %v", err.Error())
	}

	return nil, nil
}

type driverOp func(*FlexVolumeDriver, []string) (map[string]interface{}, error)

type cmdInfo struct {
	numArgs int
	run     driverOp
}

var commands = map[string]cmdInfo{
	"init": {
		0, func(d *FlexVolumeDriver, args []string) (map[string]interface{}, error) {
			return d.init()
		},
	},
	"attach": {
		2, func(d *FlexVolumeDriver, args []string) (map[string]interface{}, error) {
			return d.attach(args[0], args[1])
		},
	},
	"detach": {
		2, func(d *FlexVolumeDriver, args []string) (map[string]interface{}, error) {
			return d.detach(args[0], args[1])
		},
	},
	"waitforattach": {
		2, func(d *FlexVolumeDriver, args []string) (map[string]interface{}, error) {
			return d.waitForAttach(args[0], args[1])
		},
	},
	"isattached": {
		2, func(d *FlexVolumeDriver, args []string) (map[string]interface{}, error) {
			return d.isAttached(args[0], args[1])
		},
	},
	"mount": {
		2, func(d *FlexVolumeDriver, args []string) (map[string]interface{}, error) {
			return d.mount(args[0], args[1])
		},
	},
	"unmount": {
		1, func(d *FlexVolumeDriver, args []string) (map[string]interface{}, error) {
			return d.unmount(args[0])
		},
	},
}

func (d *FlexVolumeDriver) doRun(args []string) (map[string]interface{}, error) {
	if len(args) == 0 {
		return nil, errors.New("no arguments passed to flexvolume driver")
	}
	nArgs := len(args) - 1
	op := args[0]
	if cmdInfo, found := commands[op]; found {
		if cmdInfo.numArgs == nArgs {
			return cmdInfo.run(d, args[1:])
		} else {
			return nil, fmt.Errorf("unexpected number of args %d (expected %d) for operation %q", nArgs, cmdInfo.numArgs, op)
		}
	} else {
		return map[string]interface{}{
			"status": "Not supported",
		}, nil
	}
}

func (d *FlexVolumeDriver) Run(args []string) string {
	r := formatResult(d.doRun(args))

	if flexVolumeDebug {
		// This is for debugging purposes only.
		// The problem is that kubelet grabs CombinedOutput() from the process
		// and tries to parse it as JSON (need to recheck this,
		// maybe submit a PS to fix it)
		f, err := os.OpenFile("/tmp/flexvolume.log", os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0666)
		if err == nil {
			defer f.Close()
			fmt.Fprintf(f, "flexvolume %s -> %s\n", strings.Join(args, " "), r)
		}
	}

	return r
}

func formatResult(fields map[string]interface{}, err error) string {
	var data map[string]interface{}
	if err != nil {
		data = map[string]interface{}{
			"status":  "Failure",
			"message": err.Error(),
		}
	} else {
		data = map[string]interface{}{
			"status": "Success",
		}
		for k, v := range fields {
			data[k] = v
		}
	}
	s, err := json.Marshal(data)
	if err != nil {
		panic("error marshalling the data")
	}
	return string(s) + "\n"
}

// GetFlexvolumeInfo tries to extract flexvolume uuid and partition
// number from the specified directory. Negative partition number
// means that no partition number was specified.
func GetFlexvolumeInfo(dir string) (string, int, error) {
	dataFile := filepath.Join(dir, flexvolumeDataFile)
	var opts map[string]interface{}
	if err := utils.ReadJSON(dataFile, &opts); err != nil {
		return "", 0, fmt.Errorf("can't read flexvolume data file %q: %v", dataFile, err)
	}
	uuidRaw, found := opts[uuidOptionsKey]
	if !found {
		return "", 0, fmt.Errorf("%q: flexvolume doesn't have an uuid", dataFile)
	}
	uuid, ok := uuidRaw.(string)
	if !ok {
		return "", 0, fmt.Errorf("%q: flexvolume uuid is not a string", dataFile)
	}
	part := -1
	partRaw, found := opts[partOptionsKey]
	if found {
		partStr, ok := partRaw.(string)
		if !ok {
			return "", 0, fmt.Errorf("%q: 'part' value is not a string", dataFile)
		}
		var err error
		part, err = strconv.Atoi(partStr)
		if err != nil {
			return "", 0, fmt.Errorf("%q: malformed 'part' value (partition number): %q: %v", dataFile, partRaw, err)
		}
	}
	return uuid, part, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package container

import (
	"fmt"

	"k8s.io/api/core/v1"
	ref "k8s.io/client-go/tools/reference"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
)

var ImplicitContainerPrefix string = "implicitly required container "

// GenerateContainerRef returns an *v1.ObjectReference which references the given container
// within the given pod. Returns an error if the reference can't be constructed or the
// container doesn't actually belong to the pod.
//
// This function will return an error if the provided Pod does not have a selfLink,
// but we expect selfLink to be populated at all call sites for the function.
func GenerateContainerRef(pod *v1.Pod, container *v1.Container) (*v1.ObjectReference, error) {
	fieldPath, err := fieldPath(pod, container)
	if err != nil {
		// TODO: figure out intelligent way to refer to containers that we implicitly
		// start (like the pod infra container). This is not a good way, ugh.
		fieldPath = ImplicitContainerPrefix + container.Name
	}
	ref, err := ref.GetPartialReference(legacyscheme.Scheme, pod, fieldPath)
	if err != nil {
		return nil, err
	}
	return ref, nil
}

// fieldPath returns a fieldPath locating container within pod.
// Returns an error if the container isn't part of the pod.
func fieldPath(pod *v1.Pod, container *v1.Container) (string, error) {
	for i := range pod.Spec.Containers {
		here := &pod.Spec.Containers[i]
		if here.Name == container.Name {
			if here.Name == "" {
				return fmt.Sprintf("spec.containers[%d]", i), nil
			} else {
				return fmt.Sprintf("spec.containers{%s}", here.Name), nil
			}
		}
	}
	for i := range pod.Spec.InitContainers {
		here := &pod.Spec.InitContainers[i]
		if here.Name == container.Name {
			if here.Name == "" {
				return fmt.Sprintf("spec.initContainers[%d]", i), nil
			} else {
				return fmt.Sprintf("spec.initContainers{%s}", here.Name), nil
			}
		}
	}
	return "", fmt.Errorf("container %#v not found in pod %#v", container, pod)
}
<commit_msg>Avoid using `%#v` for errors when using only a portion of the object<commit_after>/*
Copyright 2015 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package container

import (
	"fmt"

	"k8s.io/api/core/v1"
	ref "k8s.io/client-go/tools/reference"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
)

var ImplicitContainerPrefix string = "implicitly required container "

// GenerateContainerRef returns an *v1.ObjectReference which references the given container
// within the given pod. Returns an error if the reference can't be constructed or the
// container doesn't actually belong to the pod.
//
// This function will return an error if the provided Pod does not have a selfLink,
// but we expect selfLink to be populated at all call sites for the function.
func GenerateContainerRef(pod *v1.Pod, container *v1.Container) (*v1.ObjectReference, error) {
	fieldPath, err := fieldPath(pod, container)
	if err != nil {
		// TODO: figure out intelligent way to refer to containers that we implicitly
		// start (like the pod infra container). This is not a good way, ugh.
		fieldPath = ImplicitContainerPrefix + container.Name
	}
	ref, err := ref.GetPartialReference(legacyscheme.Scheme, pod, fieldPath)
	if err != nil {
		return nil, err
	}
	return ref, nil
}

// fieldPath returns a fieldPath locating container within pod.
// Returns an error if the container isn't part of the pod.
func fieldPath(pod *v1.Pod, container *v1.Container) (string, error) {
	for i := range pod.Spec.Containers {
		here := &pod.Spec.Containers[i]
		if here.Name == container.Name {
			if here.Name == "" {
				return fmt.Sprintf("spec.containers[%d]", i), nil
			}
			return fmt.Sprintf("spec.containers{%s}", here.Name), nil
		}
	}
	for i := range pod.Spec.InitContainers {
		here := &pod.Spec.InitContainers[i]
		if here.Name == container.Name {
			if here.Name == "" {
				return fmt.Sprintf("spec.initContainers[%d]", i), nil
			}
			return fmt.Sprintf("spec.initContainers{%s}", here.Name), nil
		}
	}
	return "", fmt.Errorf("container %q not found in pod %s/%s", container.Name, pod.Namespace, pod.Name)
}
<|endoftext|>