<commit_before>package mtree

import (
	"archive/tar"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"fmt"
	"hash"
	"io"
	"os"

	"github.com/vbatts/go-mtree/pkg/govis"
	"golang.org/x/crypto/ripemd160"
)

// KeywordFunc is the type of a function called on each file to be included in
// a DirectoryHierarchy, that will produce the string output of the keyword to
// be included for the file entry. Otherwise, empty string.
// io.Reader `r` is to the file stream for the file payload. While this
// function takes an io.Reader, the caller needs to reset it to the beginning
// for each new KeywordFunc
type KeywordFunc func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error)

var (
	// KeywordFuncs is the map of all keywords (and the functions to produce them)
	KeywordFuncs = map[Keyword]KeywordFunc{
		"size":            sizeKeywordFunc,                                      // The size, in bytes, of the file
		"type":            typeKeywordFunc,                                      // The type of the file
		"time":            timeKeywordFunc,                                      // The last modification time of the file
		"link":            linkKeywordFunc,                                      // The target of the symbolic link when type=link
		"uid":             uidKeywordFunc,                                       // The file owner as a numeric value
		"gid":             gidKeywordFunc,                                       // The file group as a numeric value
		"nlink":           nlinkKeywordFunc,                                     // The number of hard links the file is expected to have
		"uname":           unameKeywordFunc,                                     // The file owner as a symbolic name
		"gname":           gnameKeywordFunc,                                     // The file group as a symbolic name
		"mode":            modeKeywordFunc,                                      // The current file's permissions as a numeric (octal) or symbolic value
		"cksum":           cksumKeywordFunc,                                     // The checksum of the file using the default algorithm specified by the cksum(1) utility
		"md5":             hasherKeywordFunc("md5digest", md5.New),              // The MD5 message digest of the file
		"md5digest":       hasherKeywordFunc("md5digest", md5.New),              // A synonym for `md5`
		"rmd160":          hasherKeywordFunc("ripemd160digest", ripemd160.New),  // The RIPEMD160 message digest of the file
		"rmd160digest":    hasherKeywordFunc("ripemd160digest", ripemd160.New),  // A synonym for `rmd160`
		"ripemd160digest": hasherKeywordFunc("ripemd160digest", ripemd160.New),  // A synonym for `rmd160`
		"sha1":            hasherKeywordFunc("sha1digest", sha1.New),            // The SHA1 message digest of the file
		"sha1digest":      hasherKeywordFunc("sha1digest", sha1.New),            // A synonym for `sha1`
		"sha256":          hasherKeywordFunc("sha256digest", sha256.New),        // The SHA256 message digest of the file
		"sha256digest":    hasherKeywordFunc("sha256digest", sha256.New),        // A synonym for `sha256`
		"sha384":          hasherKeywordFunc("sha384digest", sha512.New384),     // The SHA384 message digest of the file
		"sha384digest":    hasherKeywordFunc("sha384digest", sha512.New384),     // A synonym for `sha384`
		"sha512":          hasherKeywordFunc("sha512digest", sha512.New),        // The SHA512 message digest of the file
		"sha512digest":    hasherKeywordFunc("sha512digest", sha512.New),        // A synonym for `sha512`
		"sha512256":       hasherKeywordFunc("sha512digest", sha512.New512_256), // The SHA512/256 message digest of the file
		"sha512256digest": hasherKeywordFunc("sha512digest", sha512.New512_256), // A synonym for `sha512256`

		"flags": flagsKeywordFunc, // NOTE: this is a noop, but here to support the presence of the "flags" keyword.

		// This is not an upstreamed keyword, but used to vary from "time", as tar
		// archives do not store nanosecond precision. So comparing on "time" will
		// be only seconds level accurate.
		"tar_time": tartimeKeywordFunc, // The last modification time of the file, from a tar archive mtime

		// This is not an upstreamed keyword, but a needed attribute for file validation.
		// The pattern for this keyword key is prefixed by "xattr." followed by the extended attribute "namespace.key".
		// The keyword value is the SHA1 digest of the extended attribute's value.
		// In this way, the order of the keys does not matter, and the contents of the value is not revealed.
		"xattr":  xattrKeywordFunc,
		"xattrs": xattrKeywordFunc,
	}
)
var (
	modeKeywordFunc = func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error) {
		permissions := info.Mode().Perm()
		if os.ModeSetuid&info.Mode() > 0 {
			permissions |= (1 << 11)
		}
		if os.ModeSetgid&info.Mode() > 0 {
			permissions |= (1 << 10)
		}
		if os.ModeSticky&info.Mode() > 0 {
			permissions |= (1 << 9)
		}
		return []KeyVal{KeyVal(fmt.Sprintf("mode=%#o", permissions))}, nil
	}
	sizeKeywordFunc = func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error) {
		if sys, ok := info.Sys().(*tar.Header); ok {
			if sys.Typeflag == tar.TypeSymlink {
				return []KeyVal{KeyVal(fmt.Sprintf("size=%d", len(sys.Linkname)))}, nil
			}
		}
		return []KeyVal{KeyVal(fmt.Sprintf("size=%d", info.Size()))}, nil
	}
	cksumKeywordFunc = func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error) {
		if !info.Mode().IsRegular() {
			return nil, nil
		}
		sum, _, err := cksum(r)
		if err != nil {
			return nil, err
		}
		return []KeyVal{KeyVal(fmt.Sprintf("cksum=%d", sum))}, nil
	}
	hasherKeywordFunc = func(name string, newHash func() hash.Hash) KeywordFunc {
		return func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error) {
			if !info.Mode().IsRegular() {
				return nil, nil
			}
			h := newHash()
			if _, err := io.Copy(h, r); err != nil {
				return nil, err
			}
			return []KeyVal{KeyVal(fmt.Sprintf("%s=%x", KeywordSynonym(name), h.Sum(nil)))}, nil
		}
	}
	tartimeKeywordFunc = func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error) {
		return []KeyVal{KeyVal(fmt.Sprintf("tar_time=%d.%9.9d", info.ModTime().Unix(), 0))}, nil
	}
	timeKeywordFunc = func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error) {
		tSec := info.ModTime().Unix()
		tNano := info.ModTime().Nanosecond()
		return []KeyVal{KeyVal(fmt.Sprintf("time=%d.%9.9d", tSec, tNano))}, nil
	}
	linkKeywordFunc = func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error) {
		if sys, ok := info.Sys().(*tar.Header); ok {
			if sys.Linkname != "" {
				linkname, err := govis.Vis(sys.Linkname, DefaultVisFlags)
				if err != nil {
					return nil, nil
				}
				return []KeyVal{KeyVal(fmt.Sprintf("link=%s", linkname))}, nil
			}
			return nil, nil
		}

		if info.Mode()&os.ModeSymlink != 0 {
			str, err := os.Readlink(path)
			if err != nil {
				return nil, nil
			}
			linkname, err := govis.Vis(str, DefaultVisFlags)
			if err != nil {
				return nil, nil
			}
			return []KeyVal{KeyVal(fmt.Sprintf("link=%s", linkname))}, nil
		}
		return nil, nil
	}
	typeKeywordFunc = func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error) {
		if info.Mode().IsDir() {
			return []KeyVal{"type=dir"}, nil
		}
		if info.Mode().IsRegular() {
			return []KeyVal{"type=file"}, nil
		}
		if info.Mode()&os.ModeSocket != 0 {
			return []KeyVal{"type=socket"}, nil
		}
		if info.Mode()&os.ModeSymlink != 0 {
			return []KeyVal{"type=link"}, nil
		}
		if info.Mode()&os.ModeNamedPipe != 0 {
			return []KeyVal{"type=fifo"}, nil
		}
		if info.Mode()&os.ModeDevice != 0 {
			if info.Mode()&os.ModeCharDevice != 0 {
				return []KeyVal{"type=char"}, nil
			}
			return []KeyVal{"type=block"}, nil
		}
		return nil, nil
	}
)
<commit_msg>keywordfunc: ignore ripemd160 for compat sake<commit_after>package mtree

import (
	"archive/tar"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"fmt"
	"hash"
	"io"
	"os"

	"github.com/vbatts/go-mtree/pkg/govis"

	//lint:ignore SA1019 yes ripemd160 is deprecated, but this is for mtree compatibility
	"golang.org/x/crypto/ripemd160"
)

// KeywordFunc is the type of a function called on each file to be included in
// a DirectoryHierarchy, that will produce the string output of the keyword to
// be included for the file entry. Otherwise, empty string.
// io.Reader `r` is to the file stream for the file payload. While this
// function takes an io.Reader, the caller needs to reset it to the beginning
// for each new KeywordFunc
type KeywordFunc func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error)

var (
	// KeywordFuncs is the map of all keywords (and the functions to produce them)
	KeywordFuncs = map[Keyword]KeywordFunc{
		"size":            sizeKeywordFunc,                                      // The size, in bytes, of the file
		"type":            typeKeywordFunc,                                      // The type of the file
		"time":            timeKeywordFunc,                                      // The last modification time of the file
		"link":            linkKeywordFunc,                                      // The target of the symbolic link when type=link
		"uid":             uidKeywordFunc,                                       // The file owner as a numeric value
		"gid":             gidKeywordFunc,                                       // The file group as a numeric value
		"nlink":           nlinkKeywordFunc,                                     // The number of hard links the file is expected to have
		"uname":           unameKeywordFunc,                                     // The file owner as a symbolic name
		"gname":           gnameKeywordFunc,                                     // The file group as a symbolic name
		"mode":            modeKeywordFunc,                                      // The current file's permissions as a numeric (octal) or symbolic value
		"cksum":           cksumKeywordFunc,                                     // The checksum of the file using the default algorithm specified by the cksum(1) utility
		"md5":             hasherKeywordFunc("md5digest", md5.New),              // The MD5 message digest of the file
		"md5digest":       hasherKeywordFunc("md5digest", md5.New),              // A synonym for `md5`
		"rmd160":          hasherKeywordFunc("ripemd160digest", ripemd160.New),  // The RIPEMD160 message digest of the file
		"rmd160digest":    hasherKeywordFunc("ripemd160digest", ripemd160.New),  // A synonym for `rmd160`
		"ripemd160digest": hasherKeywordFunc("ripemd160digest", ripemd160.New),  // A synonym for `rmd160`
		"sha1":            hasherKeywordFunc("sha1digest", sha1.New),            // The SHA1 message digest of the file
		"sha1digest":      hasherKeywordFunc("sha1digest", sha1.New),            // A synonym for `sha1`
		"sha256":          hasherKeywordFunc("sha256digest", sha256.New),        // The SHA256 message digest of the file
		"sha256digest":    hasherKeywordFunc("sha256digest", sha256.New),        // A synonym for `sha256`
		"sha384":          hasherKeywordFunc("sha384digest", sha512.New384),     // The SHA384 message digest of the file
		"sha384digest":    hasherKeywordFunc("sha384digest", sha512.New384),     // A synonym for `sha384`
		"sha512":          hasherKeywordFunc("sha512digest", sha512.New),        // The SHA512 message digest of the file
		"sha512digest":    hasherKeywordFunc("sha512digest", sha512.New),        // A synonym for `sha512`
		"sha512256":       hasherKeywordFunc("sha512digest", sha512.New512_256), // The SHA512/256 message digest of the file
		"sha512256digest": hasherKeywordFunc("sha512digest", sha512.New512_256), // A synonym for `sha512256`

		"flags": flagsKeywordFunc, // NOTE: this is a noop, but here to support the presence of the "flags" keyword.

		// This is not an upstreamed keyword, but used to vary from "time", as tar
		// archives do not store nanosecond precision. So comparing on "time" will
		// be only seconds level accurate.
		"tar_time": tartimeKeywordFunc, // The last modification time of the file, from a tar archive mtime

		// This is not an upstreamed keyword, but a needed attribute for file validation.
		// The pattern for this keyword key is prefixed by "xattr." followed by the extended attribute "namespace.key".
		// The keyword value is the SHA1 digest of the extended attribute's value.
		// In this way, the order of the keys does not matter, and the contents of the value is not revealed.
		"xattr":  xattrKeywordFunc,
		"xattrs": xattrKeywordFunc,
	}
)
var (
	modeKeywordFunc = func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error) {
		permissions := info.Mode().Perm()
		if os.ModeSetuid&info.Mode() > 0 {
			permissions |= (1 << 11)
		}
		if os.ModeSetgid&info.Mode() > 0 {
			permissions |= (1 << 10)
		}
		if os.ModeSticky&info.Mode() > 0 {
			permissions |= (1 << 9)
		}
		return []KeyVal{KeyVal(fmt.Sprintf("mode=%#o", permissions))}, nil
	}
	sizeKeywordFunc = func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error) {
		if sys, ok := info.Sys().(*tar.Header); ok {
			if sys.Typeflag == tar.TypeSymlink {
				return []KeyVal{KeyVal(fmt.Sprintf("size=%d", len(sys.Linkname)))}, nil
			}
		}
		return []KeyVal{KeyVal(fmt.Sprintf("size=%d", info.Size()))}, nil
	}
	cksumKeywordFunc = func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error) {
		if !info.Mode().IsRegular() {
			return nil, nil
		}
		sum, _, err := cksum(r)
		if err != nil {
			return nil, err
		}
		return []KeyVal{KeyVal(fmt.Sprintf("cksum=%d", sum))}, nil
	}
	hasherKeywordFunc = func(name string, newHash func() hash.Hash) KeywordFunc {
		return func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error) {
			if !info.Mode().IsRegular() {
				return nil, nil
			}
			h := newHash()
			if _, err := io.Copy(h, r); err != nil {
				return nil, err
			}
			return []KeyVal{KeyVal(fmt.Sprintf("%s=%x", KeywordSynonym(name), h.Sum(nil)))}, nil
		}
	}
	tartimeKeywordFunc = func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error) {
		return []KeyVal{KeyVal(fmt.Sprintf("tar_time=%d.%9.9d", info.ModTime().Unix(), 0))}, nil
	}
	timeKeywordFunc = func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error) {
		tSec := info.ModTime().Unix()
		tNano := info.ModTime().Nanosecond()
		return []KeyVal{KeyVal(fmt.Sprintf("time=%d.%9.9d", tSec, tNano))}, nil
	}
	linkKeywordFunc = func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error) {
		if sys, ok := info.Sys().(*tar.Header); ok {
			if sys.Linkname != "" {
				linkname, err := govis.Vis(sys.Linkname, DefaultVisFlags)
				if err != nil {
					return nil, nil
				}
				return []KeyVal{KeyVal(fmt.Sprintf("link=%s", linkname))}, nil
			}
			return nil, nil
		}

		if info.Mode()&os.ModeSymlink != 0 {
			str, err := os.Readlink(path)
			if err != nil {
				return nil, nil
			}
			linkname, err := govis.Vis(str, DefaultVisFlags)
			if err != nil {
				return nil, nil
			}
			return []KeyVal{KeyVal(fmt.Sprintf("link=%s", linkname))}, nil
		}
		return nil, nil
	}
	typeKeywordFunc = func(path string, info os.FileInfo, r io.Reader) ([]KeyVal, error) {
		if info.Mode().IsDir() {
			return []KeyVal{"type=dir"}, nil
		}
		if info.Mode().IsRegular() {
			return []KeyVal{"type=file"}, nil
		}
		if info.Mode()&os.ModeSocket != 0 {
			return []KeyVal{"type=socket"}, nil
		}
		if info.Mode()&os.ModeSymlink != 0 {
			return []KeyVal{"type=link"}, nil
		}
		if info.Mode()&os.ModeNamedPipe != 0 {
			return []KeyVal{"type=fifo"}, nil
		}
		if info.Mode()&os.ModeDevice != 0 {
			if info.Mode()&os.ModeCharDevice != 0 {
				return []KeyVal{"type=char"}, nil
			}
			return []KeyVal{"type=block"}, nil
		}
		return nil, nil
	}
)
<|endoftext|>


<commit_before>package negroni

import (
	"bufio"
	"net"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

type closeNotifyingRecorder struct {
	*httptest.ResponseRecorder
	closed chan bool
}

func newCloseNotifyingRecorder() *closeNotifyingRecorder {
	return &closeNotifyingRecorder{
		httptest.NewRecorder(),
		make(chan bool, 1),
	}
}

func (c *closeNotifyingRecorder) close() {
	c.closed <- true
}

func (c *closeNotifyingRecorder) CloseNotify() <-chan bool {
	return c.closed
}

type hijackableResponse struct {
	Hijacked bool
}

func newHijackableResponse() *hijackableResponse {
	return &hijackableResponse{}
}

func (h *hijackableResponse) Header() http.Header           { return nil }
func (h *hijackableResponse) Write(buf []byte) (int, error) { return 0, nil }
func (h *hijackableResponse) WriteHeader(code int)          {}
func (h *hijackableResponse) Flush()                        {}
func (h *hijackableResponse) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	h.Hijacked = true
	return nil, nil, nil
}

func TestResponseWriterBeforeWrite(t *testing.T) {
	rec := httptest.NewRecorder()
	rw := NewResponseWriter(rec)

	expect(t, rw.Status(), 0)
	expect(t, rw.Written(), false)
}

func TestResponseWriterWritingString(t *testing.T) {
	rec := httptest.NewRecorder()
	rw := NewResponseWriter(rec)

	rw.Write([]byte("Hello world"))

	expect(t, rec.Code, rw.Status())
	expect(t, rec.Body.String(), "Hello world")
	expect(t, rw.Status(), http.StatusOK)
	expect(t, rw.Size(), 11)
	expect(t, rw.Written(), true)
}

func TestResponseWriterWritingStrings(t *testing.T) {
	rec := httptest.NewRecorder()
	rw := NewResponseWriter(rec)

	rw.Write([]byte("Hello world"))
	rw.Write([]byte("foo bar bat baz"))

	expect(t, rec.Code, rw.Status())
	expect(t, rec.Body.String(), "Hello worldfoo bar bat baz")
	expect(t, rw.Status(), http.StatusOK)
	expect(t, rw.Size(), 26)
}

func TestResponseWriterWritingHeader(t *testing.T) {
	rec := httptest.NewRecorder()
	rw := NewResponseWriter(rec)

	rw.WriteHeader(http.StatusNotFound)

	expect(t, rec.Code, rw.Status())
	expect(t, rec.Body.String(), "")
	expect(t, rw.Status(), http.StatusNotFound)
	expect(t, rw.Size(), 0)
}

func TestResponseWriterBefore(t *testing.T) {
	rec := httptest.NewRecorder()
	rw := NewResponseWriter(rec)
	result := ""

	rw.Before(func(ResponseWriter) {
		result += "foo"
	})
	rw.Before(func(ResponseWriter) {
		result += "bar"
	})

	rw.WriteHeader(http.StatusNotFound)

	expect(t, rec.Code, rw.Status())
	expect(t, rec.Body.String(), "")
	expect(t, rw.Status(), http.StatusNotFound)
	expect(t, rw.Size(), 0)
	expect(t, result, "barfoo")
}

func TestResponseWriterHijack(t *testing.T) {
	hijackable := newHijackableResponse()
	rw := NewResponseWriter(hijackable)
	hijacker, ok := rw.(http.Hijacker)
	expect(t, ok, true)
	_, _, err := hijacker.Hijack()
	if err != nil {
		t.Error(err)
	}
	expect(t, hijackable.Hijacked, true)
}

func TestResponseWriteHijackNotOK(t *testing.T) {
	hijackable := new(http.ResponseWriter)
	rw := NewResponseWriter(*hijackable)
	hijacker, ok := rw.(http.Hijacker)
	expect(t, ok, true)
	_, _, err := hijacker.Hijack()

	refute(t, err, nil)
}

func TestResponseWriterCloseNotify(t *testing.T) {
	rec := newCloseNotifyingRecorder()
	rw := NewResponseWriter(rec)
	closed := false
	notifier := rw.(http.CloseNotifier).CloseNotify()
	rec.close()
	select {
	case <-notifier:
		closed = true
	case <-time.After(time.Second):
	}
	expect(t, closed, true)
}

func TestResponseWriterFlusher(t *testing.T) {
	rec := httptest.NewRecorder()
	rw := NewResponseWriter(rec)

	_, ok := rw.(http.Flusher)
	expect(t, ok, true)
}
<commit_msg>Add test asserting `beforeFunc`s have access to status<commit_after>package negroni

import (
	"bufio"
	"net"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

type closeNotifyingRecorder struct {
	*httptest.ResponseRecorder
	closed chan bool
}

func newCloseNotifyingRecorder() *closeNotifyingRecorder {
	return &closeNotifyingRecorder{
		httptest.NewRecorder(),
		make(chan bool, 1),
	}
}

func (c *closeNotifyingRecorder) close() {
	c.closed <- true
}

func (c *closeNotifyingRecorder) CloseNotify() <-chan bool {
	return c.closed
}

type hijackableResponse struct {
	Hijacked bool
}

func newHijackableResponse() *hijackableResponse {
	return &hijackableResponse{}
}

func (h *hijackableResponse) Header() http.Header           { return nil }
func (h *hijackableResponse) Write(buf []byte) (int, error) { return 0, nil }
func (h *hijackableResponse) WriteHeader(code int)          {}
func (h *hijackableResponse) Flush()                        {}
func (h *hijackableResponse) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	h.Hijacked = true
	return nil, nil, nil
}

func TestResponseWriterBeforeWrite(t *testing.T) {
	rec := httptest.NewRecorder()
	rw := NewResponseWriter(rec)

	expect(t, rw.Status(), 0)
	expect(t, rw.Written(), false)
}

func TestResponseWriterBeforeFuncHasAccessToStatus(t *testing.T) {
	var status int

	rec := httptest.NewRecorder()
	rw := NewResponseWriter(rec)

	rw.Before(func(w ResponseWriter) {
		status = w.Status()
	})
	rw.WriteHeader(http.StatusCreated)

	expect(t, status, http.StatusCreated)
}

func TestResponseWriterWritingString(t *testing.T) {
	rec := httptest.NewRecorder()
	rw := NewResponseWriter(rec)

	rw.Write([]byte("Hello world"))

	expect(t, rec.Code, rw.Status())
	expect(t, rec.Body.String(), "Hello world")
	expect(t, rw.Status(), http.StatusOK)
	expect(t, rw.Size(), 11)
	expect(t, rw.Written(), true)
}

func TestResponseWriterWritingStrings(t *testing.T) {
	rec := httptest.NewRecorder()
	rw := NewResponseWriter(rec)

	rw.Write([]byte("Hello world"))
	rw.Write([]byte("foo bar bat baz"))

	expect(t, rec.Code, rw.Status())
	expect(t, rec.Body.String(), "Hello worldfoo bar bat baz")
	expect(t, rw.Status(), http.StatusOK)
	expect(t, rw.Size(), 26)
}

func TestResponseWriterWritingHeader(t *testing.T) {
	rec := httptest.NewRecorder()
	rw := NewResponseWriter(rec)

	rw.WriteHeader(http.StatusNotFound)

	expect(t, rec.Code, rw.Status())
	expect(t, rec.Body.String(), "")
	expect(t, rw.Status(), http.StatusNotFound)
	expect(t, rw.Size(), 0)
}

func TestResponseWriterBefore(t *testing.T) {
	rec := httptest.NewRecorder()
	rw := NewResponseWriter(rec)
	result := ""

	rw.Before(func(ResponseWriter) {
		result += "foo"
	})
	rw.Before(func(ResponseWriter) {
		result += "bar"
	})

	rw.WriteHeader(http.StatusNotFound)

	expect(t, rec.Code, rw.Status())
	expect(t, rec.Body.String(), "")
	expect(t, rw.Status(), http.StatusNotFound)
	expect(t, rw.Size(), 0)
	expect(t, result, "barfoo")
}

func TestResponseWriterHijack(t *testing.T) {
	hijackable := newHijackableResponse()
	rw := NewResponseWriter(hijackable)
	hijacker, ok := rw.(http.Hijacker)
	expect(t, ok, true)
	_, _, err := hijacker.Hijack()
	if err != nil {
		t.Error(err)
	}
	expect(t, hijackable.Hijacked, true)
}

func TestResponseWriteHijackNotOK(t *testing.T) {
	hijackable := new(http.ResponseWriter)
	rw := NewResponseWriter(*hijackable)
	hijacker, ok := rw.(http.Hijacker)
	expect(t, ok, true)
	_, _, err := hijacker.Hijack()

	refute(t, err, nil)
}

func TestResponseWriterCloseNotify(t *testing.T) {
	rec := newCloseNotifyingRecorder()
	rw := NewResponseWriter(rec)
	closed := false
	notifier := rw.(http.CloseNotifier).CloseNotify()
	rec.close()
	select {
	case <-notifier:
		closed = true
	case <-time.After(time.Second):
	}
	expect(t, closed, true)
}

func TestResponseWriterFlusher(t *testing.T) {
	rec := httptest.NewRecorder()
	rw := NewResponseWriter(rec)

	_, ok := rw.(http.Flusher)
	expect(t, ok, true)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"sort"
	"strings"
	"text/tabwriter"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"

	"github.com/docker/machine/drivers"
	_ "github.com/docker/machine/drivers/amazonec2"
	_ "github.com/docker/machine/drivers/azure"
	_ "github.com/docker/machine/drivers/digitalocean"
	_ "github.com/docker/machine/drivers/google"
	_ "github.com/docker/machine/drivers/none"
	_ "github.com/docker/machine/drivers/virtualbox"
	_ "github.com/docker/machine/drivers/vmwarefusion"
	_ "github.com/docker/machine/drivers/vmwarevcloudair"
	_ "github.com/docker/machine/drivers/vmwarevsphere"
	"github.com/docker/machine/state"
)

type hostListItem struct {
	Name       string
	Active     bool
	DriverName string
	State      state.State
	URL        string
}

type hostListItemByName []hostListItem

func (h hostListItemByName) Len() int {
	return len(h)
}

func (h hostListItemByName) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

func (h hostListItemByName) Less(i, j int) bool {
	return strings.ToLower(h[i].Name) < strings.ToLower(h[j].Name)
}

var Commands = []cli.Command{
	{
		Name:   "active",
		Usage:  "Get or set the active machine",
		Action: cmdActive,
	},
	{
		Flags: append(
			drivers.GetCreateFlags(),
			cli.StringFlag{
				Name: "driver, d",
				Usage: fmt.Sprintf(
					"Driver to create machine with. Available drivers: %s",
					strings.Join(drivers.GetDriverNames(), ", "),
				),
				Value: "none",
			},
		),
		Name:   "create",
		Usage:  "Create a machine",
		Action: cmdCreate,
	},
	{
		Name:   "inspect",
		Usage:  "Inspect information about a machine",
		Action: cmdInspect,
	},
	{
		Name:   "ip",
		Usage:  "Get the IP address of a machine",
		Action: cmdIp,
	},
	{
		Name:   "kill",
		Usage:  "Kill a machine",
		Action: cmdKill,
	},
	{
		Flags: []cli.Flag{
			cli.BoolFlag{
				Name:  "quiet, q",
				Usage: "Enable quiet mode",
			},
		},
		Name:   "ls",
		Usage:  "List machines",
		Action: cmdLs,
	},
	{
		Name:   "restart",
		Usage:  "Restart a machine",
		Action: cmdRestart,
	},
	{
		Flags: []cli.Flag{
			cli.BoolFlag{
				Name:  "force, f",
				Usage: "Remove local configuration even if machine cannot be removed",
			},
		},
		Name:   "rm",
		Usage:  "Remove a machine",
		Action: cmdRm,
	},
	{
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "command, c",
				Usage: "SSH Command",
				Value: "",
			},
		},
		Name:   "ssh",
		Usage:  "Log into or run a command on a machine with SSH",
		Action: cmdSsh,
	},
	{
		Name:   "start",
		Usage:  "Start a machine",
		Action: cmdStart,
	},
	{
		Name:   "stop",
		Usage:  "Stop a machine",
		Action: cmdStop,
	},
	{
		Name:   "upgrade",
		Usage:  "Upgrade a machine to the latest version of Docker",
		Action: cmdUpgrade,
	},
	{
		Name:   "url",
		Usage:  "Get the URL of a machine",
		Action: cmdUrl,
	},
}

func cmdActive(c *cli.Context) {
	name := c.Args().First()
	store := NewStore(c.GlobalString("storage-path"))

	if name == "" {
		host, err := store.GetActive()
		if err != nil {
			log.Fatalf("error getting active host: %v", err)
		}
		if host != nil {
			fmt.Println(host.Name)
		}
	} else if name != "" {
		host, err := store.Load(name)
		if err != nil {
			log.Fatalf("error loading host: %v", err)
		}

		if err := store.SetActive(host); err != nil {
			log.Fatalf("error setting active host: %v", err)
		}
	} else {
		cli.ShowCommandHelp(c, "active")
	}
}

func cmdCreate(c *cli.Context) {
	driver := c.String("driver")
	name := c.Args().First()

	if name == "" {
		cli.ShowCommandHelp(c, "create")
		log.Fatal("You must specify a machine name")
	}

	keyExists, err := drivers.PublicKeyExists()
	if err != nil {
		log.Fatal(err)
	}

	if !keyExists {
		log.Fatalf("Identity authentication public key doesn't exist at %q. Create your public key by running the \"docker\" command.", drivers.PublicKeyPath())
	}

	store := NewStore(c.GlobalString("storage-path"))

	host, err := store.Create(name, driver, c)
	if err != nil {
		log.Errorf("Error creating host: %s", err)
		if c.GlobalBool("debug") {
			log.Fatal(err)
		} else {
			log.Warnf("Removing created machine. You can run machine with the --debug flag to avoid this.")
			// we know there was an error so do not check for the error on removal
			// instead we will Fatal with a message to prevent spamming with error messages
			_ = store.Remove(name, true)
			log.Warn("You will want to check the provider to make sure the machine and associated resources were properly removed.")
			log.Fatal("Error creating machine")
		}
	}
	if err := store.SetActive(host); err != nil {
		log.Fatalf("error setting active host: %v", err)
	}

	log.Infof("%q has been created and is now the active machine. To point Docker at this machine, run: export DOCKER_HOST=$(machine url) DOCKER_AUTH=identity", name)
}

func cmdInspect(c *cli.Context) {
	prettyJSON, err := json.MarshalIndent(getHost(c), "", "    ")
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println(string(prettyJSON))
}

func cmdIp(c *cli.Context) {
	ip, err := getHost(c).Driver.GetIP()
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println(ip)
}

func cmdKill(c *cli.Context) {
	if err := getHost(c).Driver.Kill(); err != nil {
		log.Fatal(err)
	}
}

func cmdLs(c *cli.Context) {
	quiet := c.Bool("quiet")
	store := NewStore(c.GlobalString("storage-path"))

	hostList, err := store.List()
	if err != nil {
		log.Fatal(err)
	}

	w := tabwriter.NewWriter(os.Stdout, 5, 1, 3, ' ', 0)

	if !quiet {
		fmt.Fprintln(w, "NAME\tACTIVE\tDRIVER\tSTATE\tURL")
	}

	items := []hostListItem{}
	hostListItems := make(chan hostListItem)

	for _, host := range hostList {
		if !quiet {
			go getHostState(host, *store, hostListItems)
		} else {
			fmt.Fprintf(w, "%s\n", host.Name)
		}
	}

	if !quiet {
		for i := 0; i < len(hostList); i++ {
			items = append(items, <-hostListItems)
		}
	}

	close(hostListItems)

	sort.Sort(hostListItemByName(items))

	for _, item := range items {
		activeString := ""
		if item.Active {
			activeString = "*"
		}
		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n",
			item.Name, activeString, item.DriverName, item.State, item.URL)
	}

	w.Flush()
}

func cmdRestart(c *cli.Context) {
	if err := getHost(c).Driver.Restart(); err != nil {
		log.Fatal(err)
	}
}

func cmdRm(c *cli.Context) {
	if len(c.Args()) == 0 {
		cli.ShowCommandHelp(c, "rm")
		log.Fatal("You must specify a machine name")
	}

	force := c.Bool("force")

	isError := false

	store := NewStore(c.GlobalString("storage-path"))
	for _, host := range c.Args() {
		if err := store.Remove(host, force); err != nil {
			log.Errorf("Error removing machine %s: %s", host, err)
			isError = true
		}
	}
	if isError {
		log.Fatal("There was an error removing a machine. To force remove it, pass the -f option. Warning: this might leave it running on the provider.")
	}
}

func cmdSsh(c *cli.Context) {
	name := c.Args().First()
	store := NewStore(c.GlobalString("storage-path"))

	if name == "" {
		host, err := store.GetActive()
		if err != nil {
			log.Fatalf("unable to get active host: %v", err)
		}

		name = host.Name
	}

	i := 1
	for i < len(os.Args) && os.Args[i-1] != name {
		i++
	}

	host, err := store.Load(name)
	if err != nil {
		log.Fatal(err)
	}

	var sshCmd *exec.Cmd
	if c.String("command") == "" {
		sshCmd, err = host.Driver.GetSSHCommand()
	} else {
		sshCmd, err = host.Driver.GetSSHCommand(c.String("command"))
	}
	if err != nil {
		log.Fatal(err)
	}

	sshCmd.Stdin = os.Stdin
	sshCmd.Stdout = os.Stdout
	sshCmd.Stderr = os.Stderr
	if err := sshCmd.Run(); err != nil {
		log.Fatal(err)
	}
}

func cmdStart(c *cli.Context) {
	if err := getHost(c).Start(); err != nil {
		log.Fatal(err)
	}
}

func cmdStop(c *cli.Context) {
	if err := getHost(c).Stop(); err != nil {
		log.Fatal(err)
	}
}

func cmdUpgrade(c *cli.Context) {
	if err := getHost(c).Upgrade(); err != nil {
		log.Fatal(err)
	}
}

func cmdUrl(c *cli.Context) {
	url, err := getHost(c).GetURL()
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println(url)
}

func cmdNotFound(c *cli.Context, command string) {
	log.Fatalf(
		"%s: '%s' is not a %s command. See '%s --help'.",
		c.App.Name,
		command,
		c.App.Name,
		c.App.Name,
	)
}

func getHost(c *cli.Context) *Host {
	name := c.Args().First()
	store := NewStore(c.GlobalString("storage-path"))

	if name == "" {
		host, err := store.GetActive()
		if err != nil {
			log.Fatalf("unable to get active host: %v", err)
		}
		return host
	}

	host, err := store.Load(name)
	if err != nil {
		log.Fatalf("unable to load host: %v", err)
	}
	return host
}

func getHostState(host Host, store Store, hostListItems chan<- hostListItem) {
	currentState, err := host.Driver.GetState()
	if err != nil {
		log.Errorf("error getting state for host %s: %s", host.Name, err)
	}

	url, err := host.GetURL()
	if err != nil {
		if err == drivers.ErrHostIsNotRunning {
			url = ""
		} else {
			log.Errorf("error getting URL for host %s: %s", host.Name, err)
		}
	}

	isActive, err := store.IsActive(&host)
	if err != nil {
		log.Errorf("error determining whether host %q is active: %s",
			host.Name, err)
	}

	hostListItems <- hostListItem{
		Name:       host.Name,
		Active:     isActive,
		DriverName: host.Driver.DriverName(),
		State:      currentState,
		URL:        url,
	}
}
<commit_msg>Stop machine from overriding an existing machine<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"sort"
	"strings"
	"text/tabwriter"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"

	"github.com/docker/machine/drivers"
	_ "github.com/docker/machine/drivers/amazonec2"
	_ "github.com/docker/machine/drivers/azure"
	_ "github.com/docker/machine/drivers/digitalocean"
	_ "github.com/docker/machine/drivers/google"
	_ "github.com/docker/machine/drivers/none"
	_ "github.com/docker/machine/drivers/virtualbox"
	_ "github.com/docker/machine/drivers/vmwarefusion"
	_ "github.com/docker/machine/drivers/vmwarevcloudair"
	_ "github.com/docker/machine/drivers/vmwarevsphere"
	"github.com/docker/machine/state"
)

type hostListItem struct {
	Name       string
	Active     bool
	DriverName string
	State      state.State
	URL        string
}

type hostListItemByName []hostListItem

func (h hostListItemByName) Len() int {
	return len(h)
}

func (h hostListItemByName) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

func (h hostListItemByName) Less(i, j int) bool {
	return strings.ToLower(h[i].Name) < strings.ToLower(h[j].Name)
}

var Commands = []cli.Command{
	{
		Name:   "active",
		Usage:  "Get or set the active machine",
		Action: cmdActive,
	},
	{
		Flags: append(
			drivers.GetCreateFlags(),
			cli.StringFlag{
				Name: "driver, d",
				Usage: fmt.Sprintf(
					"Driver to create machine with. Available drivers: %s",
					strings.Join(drivers.GetDriverNames(), ", "),
				),
				Value: "none",
			},
		),
		Name:   "create",
		Usage:  "Create a machine",
		Action: cmdCreate,
	},
	{
		Name:   "inspect",
		Usage:  "Inspect information about a machine",
		Action: cmdInspect,
	},
	{
		Name:   "ip",
		Usage:  "Get the IP address of a machine",
		Action: cmdIp,
	},
	{
		Name:   "kill",
		Usage:  "Kill a machine",
		Action: cmdKill,
	},
	{
		Flags: []cli.Flag{
			cli.BoolFlag{
				Name:  "quiet, q",
				Usage: "Enable quiet mode",
			},
		},
		Name:   "ls",
		Usage:  "List machines",
		Action: cmdLs,
	},
	{
		Name:   "restart",
		Usage:  "Restart a machine",
		Action: cmdRestart,
	},
	{
		Flags: []cli.Flag{
			cli.BoolFlag{
				Name:  "force, f",
				Usage: "Remove local configuration even if machine cannot be removed",
			},
		},
		Name:   "rm",
		Usage:  "Remove a machine",
		Action: cmdRm,
	},
	{
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "command, c",
				Usage: "SSH Command",
				Value: "",
			},
		},
		Name:   "ssh",
		Usage:  "Log into or run a command on a machine with SSH",
		Action: cmdSsh,
	},
	{
		Name:   "start",
		Usage:  "Start a machine",
		Action: cmdStart,
	},
	{
		Name:   "stop",
		Usage:  "Stop a machine",
		Action: cmdStop,
	},
	{
		Name:   "upgrade",
		Usage:  "Upgrade a machine to the latest version of Docker",
		Action: cmdUpgrade,
	},
	{
		Name:   "url",
		Usage:  "Get the URL of a machine",
		Action: cmdUrl,
	},
}

func cmdActive(c *cli.Context) {
	name := c.Args().First()
	store := NewStore(c.GlobalString("storage-path"))

	if name == "" {
		host, err := store.GetActive()
		if err != nil {
			log.Fatalf("error getting active host: %v", err)
		}
		if host != nil {
			fmt.Println(host.Name)
		}
	} else if name != "" {
		host, err := store.Load(name)
		if err != nil {
			log.Fatalf("error loading host: %v", err)
		}

		if err := store.SetActive(host); err != nil {
			log.Fatalf("error setting active host: %v", err)
		}
	} else {
		cli.ShowCommandHelp(c, "active")
	}
}

func cmdCreate(c *cli.Context) {
	driver := c.String("driver")
	name := c.Args().First()

	if name == "" {
		cli.ShowCommandHelp(c, "create")
		log.Fatal("You must specify a machine name")
	}

	store := NewStore(c.GlobalString("storage-path"))

	exists, err := store.Exists(name)
	if err != nil {
		log.Fatal(err)
	}

	if exists {
		log.Fatal("There's already a machine with the same name")
	}

	keyExists, err := drivers.PublicKeyExists()
	if err != nil {
		log.Fatal(err)
	}

	if !keyExists {
		log.Fatalf("Identity authentication public key doesn't exist at %q. Create your public key by running the \"docker\" command.", drivers.PublicKeyPath())
	}

	host, err := store.Create(name, driver, c)
	if err != nil {
		log.Errorf("Error creating host: %s", err)
		if c.GlobalBool("debug") {
			log.Fatal(err)
		} else {
			log.Warnf("Removing created machine. You can run machine with the --debug flag to avoid this.")
			// we know there was an error so do not check for the error on removal
			// instead we will Fatal with a message to prevent spamming with error messages
			_ = store.Remove(name, true)
			log.Warn("You will want to check the provider to make sure the machine and associated resources were properly removed.")
			log.Fatal("Error creating machine")
		}
	}
	if err := store.SetActive(host); err != nil {
		log.Fatalf("error setting active host: %v", err)
	}

	log.Infof("%q has been created and is now the active machine. To point Docker at this machine, run: export DOCKER_HOST=$(machine url) DOCKER_AUTH=identity", name)
}

func cmdInspect(c *cli.Context) {
	prettyJSON, err := json.MarshalIndent(getHost(c), "", "    ")
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println(string(prettyJSON))
}

func cmdIp(c *cli.Context) {
	ip, err := getHost(c).Driver.GetIP()
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println(ip)
}

func cmdKill(c *cli.Context) {
	if err := getHost(c).Driver.Kill(); err != nil {
		log.Fatal(err)
	}
}

func cmdLs(c *cli.Context) {
	quiet := c.Bool("quiet")
	store := NewStore(c.GlobalString("storage-path"))

	hostList, err := store.List()
	if err != nil {
		log.Fatal(err)
	}

	w := tabwriter.NewWriter(os.Stdout, 5, 1, 3, ' ', 0)

	if !quiet {
		fmt.Fprintln(w, "NAME\tACTIVE\tDRIVER\tSTATE\tURL")
	}

	items := []hostListItem{}
	hostListItems := make(chan hostListItem)

	for _, host := range hostList {
		if !quiet {
			go getHostState(host, *store, hostListItems)
		} else {
			fmt.Fprintf(w, "%s\n", host.Name)
		}
	}

	if !quiet {
		for i := 0; i < len(hostList); i++ {
			items = append(items, <-hostListItems)
		}
	}

	close(hostListItems)

	sort.Sort(hostListItemByName(items))

	for _, item := range items {
		activeString := ""
		if item.Active {
			activeString = "*"
		}
		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n",
			item.Name, activeString, item.DriverName, item.State, item.URL)
	}

	w.Flush()
}

func cmdRestart(c *cli.Context) {
	if err := getHost(c).Driver.Restart(); err != nil {
		log.Fatal(err)
	}
}

func cmdRm(c *cli.Context) {
	if len(c.Args()) == 0 {
		cli.ShowCommandHelp(c, "rm")
		log.Fatal("You must specify a machine name")
	}

	force := c.Bool("force")

	isError := false

	store := NewStore(c.GlobalString("storage-path"))
	for _, host := range c.Args() {
		if err := store.Remove(host, force); err != nil {
			log.Errorf("Error removing machine %s: %s", host, err)
			isError = true
		}
	}
	if isError {
		log.Fatal("There was an error removing a machine. To force remove it, pass the -f option. Warning: this might leave it running on the provider.")
	}
}

func cmdSsh(c *cli.Context) {
	name := c.Args().First()
	store := NewStore(c.GlobalString("storage-path"))

	if name == "" {
		host, err := store.GetActive()
		if err != nil {
			log.Fatalf("unable to get active host: %v", err)
		}

		name = host.Name
	}

	i := 1
	for i < len(os.Args) && os.Args[i-1] != name {
		i++
	}

	host, err := store.Load(name)
	if err != nil {
		log.Fatal(err)
	}

	var sshCmd *exec.Cmd
	if c.String("command") == "" {
		sshCmd, err = host.Driver.GetSSHCommand()
	} else {
		sshCmd, err = host.Driver.GetSSHCommand(c.String("command"))
	}
	if err != nil {
		log.Fatal(err)
	}

	sshCmd.Stdin = os.Stdin
	sshCmd.Stdout = os.Stdout
	sshCmd.Stderr = os.Stderr
	if err := sshCmd.Run(); err != nil {
		log.Fatal(err)
	}
}

func cmdStart(c *cli.Context) {
	if err := getHost(c).Start(); err != nil {
		log.Fatal(err)
	}
}

func cmdStop(c *cli.Context) {
	if err := getHost(c).Stop(); err != nil {
		log.Fatal(err)
	}
}

func cmdUpgrade(c *cli.Context) {
	if err := getHost(c).Upgrade(); err != nil {
		log.Fatal(err)
	}
}

func cmdUrl(c *cli.Context) {
	url, err := getHost(c).GetURL()
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println(url)
}

func cmdNotFound(c *cli.Context, command string) {
	log.Fatalf(
		"%s: '%s' is not a %s command. See '%s --help'.",
		c.App.Name,
		command,
		c.App.Name,
		c.App.Name,
	)
}

func getHost(c *cli.Context) *Host {
	name := c.Args().First()
	store := NewStore(c.GlobalString("storage-path"))

	if name == "" {
		host, err := store.GetActive()
		if err != nil {
			log.Fatalf("unable to get active host: %v", err)
		}
		return host
	}

	host, err := store.Load(name)
	if err != nil {
		log.Fatalf("unable to load host: %v", err)
	}
	return host
}

func getHostState(host Host, store Store, hostListItems chan<- hostListItem) {
	currentState, err := host.Driver.GetState()
	if err != nil {
		log.Errorf("error getting state for host %s: %s", host.Name, err)
	}

	url, err := host.GetURL()
	if err != nil {
		if err == drivers.ErrHostIsNotRunning {
			url = ""
		} else {
			log.Errorf("error getting URL for host %s: %s", host.Name, err)
		}
	}

	isActive, err := store.IsActive(&host)
	if err != nil {
		log.Errorf("error determining whether host %q is active: %s",
			host.Name, err)
	}

	hostListItems <- hostListItem{
		Name:       host.Name,
		Active:     isActive,
		DriverName: host.Driver.DriverName(),
		State:      currentState,
		URL:        url,
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"syscall"

	"github.com/codegangsta/cli"
	"github.com/motemen/ghq/utils"
)

var Commands = []cli.Command{
	commandGet,
	commandList,
	commandLook,
	commandImport,
	commandRoot,
}

var cloneFlags = []cli.Flag{
	cli.BoolFlag{Name: "update, u", Usage: "Update local repository if cloned already"},
	cli.BoolFlag{Name: "p", Usage: "Clone with SSH"},
	cli.BoolFlag{Name: "shallow", Usage: "Do a shallow clone"},
}

var commandGet = cli.Command{
	Name:  "get",
	Usage: "Clone/sync with a remote repository",
	Description: `
    Clone a GitHub repository under ghq root direcotry. If the repository is
    already cloned to local, nothing will happen unless '-u' ('--update')
    flag is supplied, in which case 'git remote update' is executed.
    When you use '-p' option, the repository is cloned via SSH.
`,
	Action: doGet,
	Flags:  cloneFlags,
}

var commandList = cli.Command{
	Name:  "list",
	Usage: "List local repositories",
	Description: `
    List locally cloned repositories. If a query argument is given, only
    repositories whose names contain that query text are listed. '-e'
    ('--exact') forces the match to be an exact one (i.e. the query equals to
    _project_ or _user_/_project_) If '-p' ('--full-path') is given, the full paths
    to the repository root are printed instead of relative ones.
`,
	Action: doList,
	Flags: []cli.Flag{
		cli.BoolFlag{Name: "exact, e", Usage: "Perform an exact match"},
		cli.BoolFlag{Name: "full-path, p", Usage: "Print full paths"},
		cli.BoolFlag{Name: "unique", Usage: "Print unique subpaths"},
	},
}

var commandLook = cli.Command{
	Name:  "look",
	Usage: "Look into a local repository",
	Description: `
    Look into a locally cloned repository with the shell.
`,
	Action: doLook,
}

var commandImport = cli.Command{
	Name:   "import",
	Usage:  "Bulk get repositories from stdin",
	Action: doImport,
	Flags:  cloneFlags,
}

var commandRoot = cli.Command{
	Name:   "root",
	Usage:  "Show repositories' root",
	Action: doRoot,
	Flags: []cli.Flag{
		cli.BoolFlag{Name: "all", Usage: "Show all roots"},
	},
}

type commandDoc struct {
	Parent    string
	Arguments string
}

var commandDocs = map[string]commandDoc{
	"get":    {"", "[-u] <repository URL> | [-u] [-p] <user>/<project>"},
	"list":   {"", "[-p] [-e] [<query>]"},
	"look":   {"", "<project> | <user>/<project> | <host>/<user>/<project>"},
	"import": {"", "< file"},
	"root":   {"", ""},
}

// Makes template conditionals to generate per-command documents.
func mkCommandsTemplate(genTemplate func(commandDoc) string) string {
	template := "{{if false}}"
	for _, command := range append(Commands) {
		template = template + fmt.Sprintf("{{else if (eq .Name %q)}}%s", command.Name, genTemplate(commandDocs[command.Name]))
	}
	return template + "{{end}}"
}

func init() {
	argsTemplate := mkCommandsTemplate(func(doc commandDoc) string { return doc.Arguments })
	parentTemplate := mkCommandsTemplate(func(doc commandDoc) string { return string(strings.TrimLeft(doc.Parent+" ", " ")) })

	cli.CommandHelpTemplate = `NAME:
    {{.Name}} - {{.Usage}}

USAGE:
    ghq ` + parentTemplate + `{{.Name}} ` + argsTemplate + `
{{if (len .Description)}}
DESCRIPTION: {{.Description}}
{{end}}{{if (len .Flags)}}
OPTIONS:
    {{range .Flags}}{{.}}
    {{end}}
{{end}}`
}

func doGet(c *cli.Context) {
	argURL := c.Args().Get(0)
	doUpdate := c.Bool("update")
	isShallow := c.Bool("shallow")

	if argURL == "" {
		cli.ShowCommandHelp(c, "get")
		os.Exit(1)
	}

	url, err := NewURL(argURL)
	utils.DieIf(err)

	isSSH := c.Bool("p")
	if isSSH {
		// Assume Git repository if `-p` is given.
		url, err = ConvertGitURLHTTPToSSH(url)
		utils.DieIf(err)
	}

	remote, err := NewRemoteRepository(url)
	utils.DieIf(err)

	if remote.IsValid() == false {
		utils.Log("error", fmt.Sprintf("Not a valid repository: %s", url))
		os.Exit(1)
	}

	getRemoteRepository(remote, doUpdate, isShallow)
}

// getRemoteRepository clones or updates a remote repository remote.
// If doUpdate is true, updates the locally cloned repository. Otherwise does nothing.
// If isShallow is true, does shallow cloning. (no effect if already cloned or the VCS is Mercurial and git-svn)
func getRemoteRepository(remote RemoteRepository, doUpdate bool, isShallow bool) {
	remoteURL := remote.URL()
	local := LocalRepositoryFromURL(remoteURL)

	path := local.FullPath
	newPath := false

	_, err := os.Stat(path)
	if err != nil {
		if os.IsNotExist(err) {
			newPath = true
			err = nil
		}
		utils.PanicIf(err)
	}

	if newPath {
		utils.Log("clone", fmt.Sprintf("%s -> %s", remoteURL, path))

		vcs := remote.VCS()
		if vcs == nil {
			utils.Log("error", fmt.Sprintf("Could not find version control system: %s", remoteURL))
			os.Exit(1)
		}

		err := vcs.Clone(remoteURL, path, isShallow)
		if err != nil {
			utils.Log("error", err.Error())
			os.Exit(1)
		}
	} else {
		if doUpdate {
			utils.Log("update", path)
			local.VCS().Update(path)
		} else {
			utils.Log("exists", path)
		}
	}
}

func doList(c *cli.Context) {
	query := c.Args().First()
	exact := c.Bool("exact")
	printFullPaths := c.Bool("full-path")
	printUniquePaths := c.Bool("unique")

	var filterFn func(*LocalRepository) bool
	if query == "" {
		filterFn = func(_ *LocalRepository) bool {
			return true
		}
	} else if exact {
		filterFn = func(repo *LocalRepository) bool {
			return repo.Matches(query)
		}
	} else {
		filterFn = func(repo *LocalRepository) bool {
			return strings.Contains(repo.NonHostPath(), query)
		}
	}

	repos := []*LocalRepository{}

	walkLocalRepositories(func(repo *LocalRepository) {
		if filterFn(repo) == false {
			return
		}

		repos = append(repos, repo)
	})

	if printUniquePaths {
		subpathCount := map[string]int{} // Count duplicated subpaths (ex. foo/dotfiles and bar/dotfiles)
		reposCount := map[string]int{}   // Check duplicated repositories among roots

		// Primary first
		for _, repo := range repos {
			if reposCount[repo.RelPath] == 0 {
				for _, p := range repo.Subpaths() {
					subpathCount[p] = subpathCount[p] + 1
				}
			}

			reposCount[repo.RelPath] = reposCount[repo.RelPath] + 1
		}

		for _, repo := range repos {
			if reposCount[repo.RelPath] > 1 && repo.IsUnderPrimaryRoot() == false {
				continue
			}

			for _, p := range repo.Subpaths() {
				if subpathCount[p] == 1 {
					fmt.Println(p)
					break
				}
			}
		}
	} else {
		for _, repo := range repos {
			if printFullPaths {
				fmt.Println(repo.FullPath)
			} else {
				fmt.Println(repo.RelPath)
			}
		}
	}
}

func doLook(c *cli.Context) {
	name := c.Args().First()

	if name == "" {
		cli.ShowCommandHelp(c, "look")
		os.Exit(1)
	}

	reposFound := []*LocalRepository{}
	walkLocalRepositories(func(repo *LocalRepository) {
		if repo.Matches(name) {
			reposFound = append(reposFound, repo)
		}
	})

	switch len(reposFound) {
	case 0:
		utils.Log("error", "No repository found")
		os.Exit(1)

	case 1:
		if runtime.GOOS == "windows" {
			cmd := exec.Command(os.Getenv("COMSPEC"))
			cmd.Stdin = os.Stdin
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
			cmd.Dir = reposFound[0].FullPath
			err := cmd.Start()
			if err == nil {
				cmd.Wait()
				os.Exit(0)
			}
		} else {
			shell := os.Getenv("SHELL")
			if shell == "" {
				shell = "/bin/sh"
			}

			utils.Log("cd", reposFound[0].FullPath)
			err := os.Chdir(reposFound[0].FullPath)
			utils.PanicIf(err)

			syscall.Exec(shell, []string{shell}, syscall.Environ())
		}

	default:
		utils.Log("error", "More than one repositories are found; Try more precise name")
		for _, repo := range reposFound {
			utils.Log("error", "- "+strings.Join(repo.PathParts, "/"))
		}
	}
}

func doImport(c *cli.Context) {
	var (
		doUpdate  = c.Bool("update")
		isSSH     = c.Bool("p")
		isShallow = c.Bool("shallow")
	)

	var (
		in       io.Reader
		finalize func() error
	)

	if len(c.Args()) == 0 {
		// `ghq import` reads URLs from stdin
		in = os.Stdin
		finalize = func() error { return nil }
	} else {
		// Handle `ghq import starred motemen` case
		// with `git config --global ghq.import.starred "!github-list-starred"`
		subCommand := c.Args().First()
		command, err := GitConfigSingle("ghq.import." + subCommand)
		if err == nil && command == "" {
			err = fmt.Errorf("ghq.import.%s configuration not found", subCommand)
		}
		utils.DieIf(err)

		// execute `sh -c 'COMMAND "$@"' -- ARG...`
		// TODO: Windows
		command = strings.TrimLeft(command, "!")
		shellCommand := append([]string{"sh", "-c", command + ` "$@"`, "--"}, c.Args().Tail()...)

		utils.Log("run", strings.Join(append([]string{command}, c.Args().Tail()...), " "))

		cmd := exec.Command(shellCommand[0], shellCommand[1:]...)
		cmd.Stderr = os.Stderr

		in, err = cmd.StdoutPipe()
		utils.DieIf(err)

		err = cmd.Start()
		utils.DieIf(err)

		finalize = cmd.Wait
	}

	scanner := bufio.NewScanner(in)
	for scanner.Scan() {
		line := scanner.Text()
		url, err := NewURL(line)
		if err != nil {
			utils.Log("error", fmt.Sprintf("Could not parse URL <%s>: %s", line, err))
			continue
		}
		if isSSH {
			url, err = ConvertGitURLHTTPToSSH(url)
			if err != nil {
				utils.Log("error", fmt.Sprintf("Could not convert URL <%s>: %s", url, err))
				continue
			}
		}

		remote, err := NewRemoteRepository(url)
		if utils.ErrorIf(err) {
			continue
		}
		if remote.IsValid() == false {
			utils.Log("error", fmt.Sprintf("Not a valid repository: %s", url))
			continue
		}

		getRemoteRepository(remote, doUpdate, isShallow)
	}
	if err := scanner.Err(); err != nil {
		utils.Log("error", fmt.Sprintf("While reading input: %s", err))
		os.Exit(1)
	}

	utils.DieIf(finalize())
}

func doRoot(c *cli.Context) {
	all := c.Bool("all")
	if all {
		for _, root := range localRepositoryRoots() {
			fmt.Println(root)
		}
	} else {
		fmt.Println(primaryLocalRepositoryRoot())
	}
}
<commit_msg>`look` command accepts remote repository url too.<commit_after>package main

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"syscall"

	"github.com/codegangsta/cli"
	"github.com/motemen/ghq/utils"
)

var Commands = []cli.Command{
	commandGet,
	commandList,
	commandLook,
	commandImport,
	commandRoot,
}

var cloneFlags = []cli.Flag{
	cli.BoolFlag{Name: "update, u", Usage: "Update local repository if cloned already"},
	cli.BoolFlag{Name: "p", Usage: "Clone with SSH"},
	cli.BoolFlag{Name: "shallow", Usage: "Do a shallow clone"},
}

var commandGet = cli.Command{
	Name:  "get",
	Usage: "Clone/sync with a remote repository",
	Description: `
    Clone a GitHub repository under ghq root direcotry. If the repository is
    already cloned to local, nothing will happen unless '-u' ('--update')
    flag is supplied, in which case 'git remote update' is executed.
    When you use '-p' option, the repository is cloned via SSH.
`,
	Action: doGet,
	Flags:  cloneFlags,
}

var commandList = cli.Command{
	Name:  "list",
	Usage: "List local repositories",
	Description: `
    List locally cloned repositories. If a query argument is given, only
    repositories whose names contain that query text are listed. '-e'
    ('--exact') forces the match to be an exact one (i.e. the query equals to
    _project_ or _user_/_project_) If '-p' ('--full-path') is given, the full paths
    to the repository root are printed instead of relative ones.
`,
	Action: doList,
	Flags: []cli.Flag{
		cli.BoolFlag{Name: "exact, e", Usage: "Perform an exact match"},
		cli.BoolFlag{Name: "full-path, p", Usage: "Print full paths"},
		cli.BoolFlag{Name: "unique", Usage: "Print unique subpaths"},
	},
}

var commandLook = cli.Command{
	Name:  "look",
	Usage: "Look into a local repository",
	Description: `
    Look into a locally cloned repository with the shell.
`,
	Action: doLook,
}

var commandImport = cli.Command{
	Name:   "import",
	Usage:  "Bulk get repositories from stdin",
	Action: doImport,
	Flags:  cloneFlags,
}

var commandRoot = cli.Command{
	Name:   "root",
	Usage:  "Show repositories' root",
	Action: doRoot,
	Flags: []cli.Flag{
		cli.BoolFlag{Name: "all", Usage: "Show all roots"},
	},
}

type commandDoc struct {
	Parent    string
	Arguments string
}

var commandDocs = map[string]commandDoc{
	"get":    {"", "[-u] <repository URL> | [-u] [-p] <user>/<project>"},
	"list":   {"", "[-p] [-e] [<query>]"},
	"look":   {"", "<project> | <user>/<project> | <host>/<user>/<project>"},
	"import": {"", "< file"},
	"root":   {"", ""},
}

// Makes template conditionals to generate per-command documents.
func mkCommandsTemplate(genTemplate func(commandDoc) string) string {
	template := "{{if false}}"
	for _, command := range append(Commands) {
		template = template + fmt.Sprintf("{{else if (eq .Name %q)}}%s", command.Name, genTemplate(commandDocs[command.Name]))
	}
	return template + "{{end}}"
}

func init() {
	argsTemplate := mkCommandsTemplate(func(doc commandDoc) string { return doc.Arguments })
	parentTemplate := mkCommandsTemplate(func(doc commandDoc) string { return string(strings.TrimLeft(doc.Parent+" ", " ")) })

	cli.CommandHelpTemplate = `NAME:
    {{.Name}} - {{.Usage}}

USAGE:
    ghq ` + parentTemplate + `{{.Name}} ` + argsTemplate + `
{{if (len .Description)}}
DESCRIPTION: {{.Description}}
{{end}}{{if (len .Flags)}}
OPTIONS:
    {{range .Flags}}{{.}}
    {{end}}
{{end}}`
}

func doGet(c *cli.Context) {
	argURL := c.Args().Get(0)
	doUpdate := c.Bool("update")
	isShallow := c.Bool("shallow")

	if argURL == "" {
		cli.ShowCommandHelp(c, "get")
		os.Exit(1)
	}

	url, err := NewURL(argURL)
	utils.DieIf(err)

	isSSH := c.Bool("p")
	if isSSH {
		// Assume Git repository if `-p` is given.
		url, err = ConvertGitURLHTTPToSSH(url)
		utils.DieIf(err)
	}

	remote, err := NewRemoteRepository(url)
	utils.DieIf(err)

	if remote.IsValid() == false {
		utils.Log("error", fmt.Sprintf("Not a valid repository: %s", url))
		os.Exit(1)
	}

	getRemoteRepository(remote, doUpdate, isShallow)
}

// getRemoteRepository clones or updates a remote repository remote.
// If doUpdate is true, updates the locally cloned repository. Otherwise does nothing.
// If isShallow is true, does shallow cloning. (no effect if already cloned or the VCS is Mercurial and git-svn)
func getRemoteRepository(remote RemoteRepository, doUpdate bool, isShallow bool) {
	remoteURL := remote.URL()
	local := LocalRepositoryFromURL(remoteURL)

	path := local.FullPath
	newPath := false

	_, err := os.Stat(path)
	if err != nil {
		if os.IsNotExist(err) {
			newPath = true
			err = nil
		}
		utils.PanicIf(err)
	}

	if newPath {
		utils.Log("clone", fmt.Sprintf("%s -> %s", remoteURL, path))

		vcs := remote.VCS()
		if vcs == nil {
			utils.Log("error", fmt.Sprintf("Could not find version control system: %s", remoteURL))
			os.Exit(1)
		}

		err := vcs.Clone(remoteURL, path, isShallow)
		if err != nil {
			utils.Log("error", err.Error())
			os.Exit(1)
		}
	} else {
		if doUpdate {
			utils.Log("update", path)
			local.VCS().Update(path)
		} else {
			utils.Log("exists", path)
		}
	}
}

func doList(c *cli.Context) {
	query := c.Args().First()
	exact := c.Bool("exact")
	printFullPaths := c.Bool("full-path")
	printUniquePaths := c.Bool("unique")

	var filterFn func(*LocalRepository) bool
	if query == "" {
		filterFn = func(_ *LocalRepository) bool {
			return true
		}
	} else if exact {
		filterFn = func(repo *LocalRepository) bool {
			return repo.Matches(query)
		}
	} else {
		filterFn = func(repo *LocalRepository) bool {
			return strings.Contains(repo.NonHostPath(), query)
		}
	}

	repos := []*LocalRepository{}

	walkLocalRepositories(func(repo *LocalRepository) {
		if filterFn(repo) == false {
			return
		}

		repos = append(repos, repo)
	})

	if printUniquePaths {
		subpathCount := map[string]int{} // Count duplicated subpaths (ex. foo/dotfiles and bar/dotfiles)
		reposCount := map[string]int{}   // Check duplicated repositories among roots

		// Primary first
		for _, repo := range repos {
			if reposCount[repo.RelPath] == 0 {
				for _, p := range repo.Subpaths() {
					subpathCount[p] = subpathCount[p] + 1
				}
			}

			reposCount[repo.RelPath] = reposCount[repo.RelPath] + 1
		}

		for _, repo := range repos {
			if reposCount[repo.RelPath] > 1 && repo.IsUnderPrimaryRoot() == false {
				continue
			}

			for _, p := range repo.Subpaths() {
				if subpathCount[p] == 1 {
					fmt.Println(p)
					break
				}
			}
		}
	} else {
		for _, repo := range repos {
			if printFullPaths {
				fmt.Println(repo.FullPath)
			} else {
				fmt.Println(repo.RelPath)
			}
		}
	}
}

func doLook(c *cli.Context) {
	name := c.Args().First()

	if name == "" {
		cli.ShowCommandHelp(c, "look")
		os.Exit(1)
	}

	reposFound := []*LocalRepository{}
	walkLocalRepositories(func(repo *LocalRepository) {
		if repo.Matches(name) {
			reposFound = append(reposFound, repo)
		}
	})

	if len(reposFound) == 0 {
		url, err := NewURL(name)

		if err == nil {
			repo := LocalRepositoryFromURL(url)
			_, err := os.Stat(repo.FullPath)

			// if the directory exists
			if err == nil {
				reposFound = append(reposFound, repo)
			}
		}
	}

	switch len(reposFound) {
	case 0:
		utils.Log("error", "No repository found")
		os.Exit(1)

	case 1:
		if runtime.GOOS == "windows" {
			cmd := exec.Command(os.Getenv("COMSPEC"))
			cmd.Stdin = os.Stdin
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
			cmd.Dir = reposFound[0].FullPath
			err := cmd.Start()
			if err == nil {
				cmd.Wait()
				os.Exit(0)
			}
		} else {
			shell := os.Getenv("SHELL")
			if shell == "" {
				shell = "/bin/sh"
			}

			utils.Log("cd", reposFound[0].FullPath)
			err := os.Chdir(reposFound[0].FullPath)
			utils.PanicIf(err)

			syscall.Exec(shell, []string{shell}, syscall.Environ())
		}

	default:
		utils.Log("error", "More than one repositories are found; Try more precise name")
		for _, repo := range reposFound {
			utils.Log("error", "- "+strings.Join(repo.PathParts, "/"))
		}
	}
}

func doImport(c *cli.Context) {
	var (
		doUpdate  = c.Bool("update")
		isSSH     = c.Bool("p")
		isShallow = c.Bool("shallow")
	)

	var (
		in       io.Reader
		finalize func() error
	)

	if len(c.Args()) == 0 {
		// `ghq import` reads URLs from stdin
		in = os.Stdin
		finalize = func() error { return nil }
	} else {
		// Handle `ghq import starred motemen` case
		// with `git config --global ghq.import.starred "!github-list-starred"`
		subCommand := c.Args().First()
		command, err := GitConfigSingle("ghq.import." + subCommand)
		if err == nil && command == "" {
			err = fmt.Errorf("ghq.import.%s configuration not found", subCommand)
		}
		utils.DieIf(err)

		// execute `sh -c 'COMMAND "$@"' -- ARG...`
		// TODO: Windows
		command = strings.TrimLeft(command, "!")
		shellCommand := append([]string{"sh", "-c", command + ` "$@"`, "--"}, c.Args().Tail()...)

		utils.Log("run", strings.Join(append([]string{command}, c.Args().Tail()...), " "))

		cmd := exec.Command(shellCommand[0], shellCommand[1:]...)
		cmd.Stderr = os.Stderr

		in, err = cmd.StdoutPipe()
		utils.DieIf(err)

		err = cmd.Start()
		utils.DieIf(err)

		finalize = cmd.Wait
	}

	scanner := bufio.NewScanner(in)
	for scanner.Scan() {
		line := scanner.Text()
		url, err := NewURL(line)
		if err != nil {
			utils.Log("error", fmt.Sprintf("Could not parse URL <%s>: %s", line, err))
			continue
		}
		if isSSH {
			url, err = ConvertGitURLHTTPToSSH(url)
			if err != nil {
				utils.Log("error", fmt.Sprintf("Could not convert URL <%s>: %s", url, err))
				continue
			}
		}

		remote, err := NewRemoteRepository(url)
		if utils.ErrorIf(err) {
			continue
		}
		if remote.IsValid() == false {
			utils.Log("error", fmt.Sprintf("Not a valid repository: %s", url))
			continue
		}

		getRemoteRepository(remote, doUpdate, isShallow)
	}
	if err := scanner.Err(); err != nil {
		utils.Log("error", fmt.Sprintf("While reading input: %s", err))
		os.Exit(1)
	}

	utils.DieIf(finalize())
}

func doRoot(c *cli.Context) {
	all := c.Bool("all")
	if all {
		for _, root := range localRepositoryRoots() {
			fmt.Println(root)
		}
	} else {
		fmt.Println(primaryLocalRepositoryRoot())
	}
}
<|endoftext|>


<commit_before>package isolated

import (
	"code.cloudfoundry.org/cli/integration/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("proxy", func() {
	var proxyURL string

	BeforeEach(func() {
		proxyURL = "http://127.0.0.1:9999"
	})

	Context("V2", func() {
		It("errors when proxy is not setup properly", func() {
			session := helpers.CFWithEnv(map[string]string{"https_proxy": proxyURL}, "api", apiURL)
			Eventually(session.Err).Should(Say("Get %s/v2/info: http: error connecting to proxy %s", apiURL, proxyURL))
			Eventually(session.Err).Should(Say("TIP: If you are behind a firewall and require an HTTP proxy, verify the https_proxy environment variable is correctly set. Else, check your network connection."))
			Eventually(session).Should(Exit(1))
		})
	})

	Context("V3", func() {
		It("errors when proxy is not setup properly", func() {
			session := helpers.CFWithEnv(map[string]string{"https_proxy": proxyURL}, "run-task", "app", "echo")
			Eventually(session.Err).Should(Say("Get %s: http: error connecting to proxy %s", apiURL, proxyURL))
			Eventually(session.Err).Should(Say("TIP: If you are behind a firewall and require an HTTP proxy, verify the https_proxy environment variable is correctly set. Else, check your network connection."))
			Eventually(session).Should(Exit(1))
		})
	})
})
<commit_msg>adjust error message due to 1.8 upgrade<commit_after>package isolated

import (
	"code.cloudfoundry.org/cli/integration/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("proxy", func() {
	var proxyURL string

	BeforeEach(func() {
		proxyURL = "http://127.0.0.1:9999"
	})

	Context("V2", func() {
		It("errors when proxy is not setup properly", func() {
			session := helpers.CFWithEnv(map[string]string{"https_proxy": proxyURL}, "api", apiURL)
			Eventually(session.Err).Should(Say("%s/v2/info.*proxy.*%s", apiURL, proxyURL))
			Eventually(session.Err).Should(Say("TIP: If you are behind a firewall and require an HTTP proxy, verify the https_proxy environment variable is correctly set. Else, check your network connection."))
			Eventually(session).Should(Exit(1))
		})
	})

	Context("V3", func() {
		It("errors when proxy is not setup properly", func() {
			session := helpers.CFWithEnv(map[string]string{"https_proxy": proxyURL}, "run-task", "app", "echo")
			Eventually(session.Err).Should(Say("Get %s: http: error connecting to proxy %s", apiURL, proxyURL))
			Eventually(session.Err).Should(Say("TIP: If you are behind a firewall and require an HTTP proxy, verify the https_proxy environment variable is correctly set. Else, check your network connection."))
			Eventually(session).Should(Exit(1))
		})
	})
})
<|endoftext|>


<commit_before>package main

import "fmt"

func (r *Rrd) addTemp(name, title string, colorid, which int) {
	r.creator.DS(name, "GAUGE", "30", "-273", "1000")
	vname := fmt.Sprintf("t%d", which)
	cname := fmt.Sprintf("f%d", which)
	r.grapher.Def(vname, r.path, name, "AVERAGE")
	if name == "solar" {
		r.grapher.CDef(cname, vname+",10,/")
	} else {
		r.grapher.CDef(cname, "9,5,/,"+vname+",*,32,+")
	}
	r.grapher.Line(2.0, cname, colorStr(colorid), title)
}

func (ppc *PoolPumpController) createRrds() {
	ppc.tempRrd.addTemp("pump", "Pump", 8, 1)
	ppc.tempRrd.addTemp("weather", "Weather", 1, 2)
	ppc.tempRrd.addTemp("roof", "Roof", 2, 3)
	ppc.tempRrd.addTemp("solar", "SolRad w/sqm", 4, 4)
	ppc.tempRrd.addTemp("pool", "Pool", 0, 5)
	ppc.tempRrd.addTemp("target", "Target", 6, 6)
	ppc.tempRrd.AddStandardRRAs()
	ppc.tempRrd.Creator().Create(*ppc.config.forceRrd)

	tg := ppc.tempRrd.grapher
	tg.SetTitle("Temperatures and Solar Radiation")
	tg.SetVLabel("Degrees Farenheit")
	tg.SetRightAxis(1, 0.0)
	tg.SetRightAxisLabel("dekawatts/sqm")
	tg.SetSize(640, 300) // Config?
	tg.SetImageFormat("PNG")

	pc := ppc.pumpRrd.Creator()
	pc.DS("status", "GAUGE", "30", "-1", "10")
	pc.DS("solar", "GAUGE", "30", "-1", "10")
	pc.DS("manual", "GAUGE", "30", "-1", "10")
	ppc.pumpRrd.AddStandardRRAs()
	pc.Create(*ppc.config.forceRrd) // fails if already exists

	pg := ppc.pumpRrd.grapher
	pg.SetTitle("Pump Activity")
	pg.SetVLabel("Status Code")
	pg.SetUpperLimit(5.0)
	pg.SetRightAxis(1, 0.0)
	pg.SetRightAxisLabel("Status Code")
	pg.SetSize(640, 200) // Config?
	pg.SetImageFormat("PNG")

	pg.Def("t1", ppc.pumpRrd.path, "status", "AVERAGE")
	pg.Line(2.0, "t1", colorStr(0), "Pump Status")
	pg.Def("t2", ppc.pumpRrd.path, "solar", "AVERAGE")
	pg.Line(2.0, "t2", colorStr(2), "Solar Status")
	pg.Def("t3", ppc.pumpRrd.path, "manual", "AVERAGE")
	pg.Line(2.0, "t3", colorStr(6), "Manual Operation")
}

// Writes updates to RRD files and generates cached graphs
func (ppc *PoolPumpController) UpdateRrd() {
	update := fmt.Sprintf("N:%f:%f:%f:%f:%f:%f",
		ppc.pumpTemp.Temperature(), ppc.WeatherC(), ppc.roofTemp.Temperature(),
		ppc.weather.GetSolarRadiation(ppc.zipcode),
		ppc.runningTemp.Temperature(), *ppc.solar.target)
	Debug("Updating TempRrd: %s", update)
	err := ppc.tempRrd.Updater().Update(update)
	if err != nil {
		Error("Update failed for TempRrd {%s}: %s", update, err.Error())
	}

	solar := 0.01
	if ppc.switches.solar.isOn() {
		solar = 1.03
	}
	manual := 0.02
	if ppc.switches.ManualState() {
		manual = 1.06
	}
	update = fmt.Sprintf("N:%d.001:%0.3f:%0.3f", ppc.switches.State(), solar, manual)
	Debug("Updating PumpRrd: %s", update)
	err = ppc.pumpRrd.Updater().Update(update)
	if err != nil {
		Error("Could not create PumpRrd: %s", err.Error())
	}
}
<commit_msg>Cleaning up comments<commit_after>package main

import "fmt"

func (r *Rrd) addTemp(name, title string, colorid, which int) {
	r.creator.DS(name, "GAUGE", "30", "-273", "1000")
	vname := fmt.Sprintf("t%d", which)
	cname := fmt.Sprintf("f%d", which)
	r.grapher.Def(vname, r.path, name, "AVERAGE")
	if name == "solar" {
		r.grapher.CDef(cname, vname+",10,/")
	} else {
		r.grapher.CDef(cname, "9,5,/,"+vname+",*,32,+")
	}
	r.grapher.Line(2.0, cname, colorStr(colorid), title)
}

func (ppc *PoolPumpController) createRrds() {
	ppc.tempRrd.addTemp("pump", "Pump", 8, 1)
	ppc.tempRrd.addTemp("weather", "Weather", 1, 2)
	ppc.tempRrd.addTemp("roof", "Roof", 2, 3)
	ppc.tempRrd.addTemp("solar", "SolRad w/sqm", 4, 4)
	ppc.tempRrd.addTemp("pool", "Pool", 0, 5)
	ppc.tempRrd.addTemp("target", "Target", 6, 6)
	ppc.tempRrd.AddStandardRRAs()
	ppc.tempRrd.Creator().Create(*ppc.config.forceRrd)

	tg := ppc.tempRrd.grapher
	tg.SetTitle("Temperatures and Solar Radiation")
	tg.SetVLabel("Degrees Farenheit")
	tg.SetRightAxis(1, 0.0)
	tg.SetRightAxisLabel("dekawatts/sqm")
	tg.SetSize(640, 300) // Config?
	tg.SetImageFormat("PNG")

	pc := ppc.pumpRrd.Creator()
	pc.DS("status", "GAUGE", "30", "-1", "10")
	pc.DS("solar", "GAUGE", "30", "-1", "10")
	pc.DS("manual", "GAUGE", "30", "-1", "10")
	ppc.pumpRrd.AddStandardRRAs()
	pc.Create(*ppc.config.forceRrd) // fails if already exists

	pg := ppc.pumpRrd.grapher
	pg.SetTitle("Pump Activity")
	pg.SetVLabel("Status Code")
	pg.SetUpperLimit(5.0)
	pg.SetRightAxis(1, 0.0)
	pg.SetRightAxisLabel("Status Code")
	pg.SetSize(640, 200) // Config?
	pg.SetImageFormat("PNG")

	pg.Def("t1", ppc.pumpRrd.path, "status", "AVERAGE")
	pg.Line(2.0, "t1", colorStr(0), "Pump Status")
	pg.Def("t2", ppc.pumpRrd.path, "solar", "AVERAGE")
	pg.Line(2.0, "t2", colorStr(2), "Solar Status")
	pg.Def("t3", ppc.pumpRrd.path, "manual", "AVERAGE")
	pg.Line(2.0, "t3", colorStr(6), "Manual Operation")
}

// UpdateRrd writes updates to RRD files and generates cached graphs
func (ppc *PoolPumpController) UpdateRrd() {
	update := fmt.Sprintf("N:%f:%f:%f:%f:%f:%f",
		ppc.pumpTemp.Temperature(), ppc.WeatherC(), ppc.roofTemp.Temperature(),
		ppc.weather.GetSolarRadiation(ppc.zipcode),
		ppc.runningTemp.Temperature(), *ppc.solar.target)
	Debug("Updating TempRrd: %s", update)
	err := ppc.tempRrd.Updater().Update(update)
	if err != nil {
		Error("Update failed for TempRrd {%s}: %s", update, err.Error())
	}

	solar := 0.01
	if ppc.switches.solar.isOn() {
		solar = 1.03
	}
	manual := 0.02
	if ppc.switches.ManualState() {
		manual = 1.06
	}
	update = fmt.Sprintf("N:%d.001:%0.3f:%0.3f", ppc.switches.State(), solar, manual)
	Debug("Updating PumpRrd: %s", update)
	err = ppc.pumpRrd.Updater().Update(update)
	if err != nil {
		Error("Could not create PumpRrd: %s", err.Error())
	}
}
<|endoftext|>


<commit_before>package sitemap

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func TestGet(t *testing.T) {
	server := server()
	defer server.Close()

	data, err := Get(server.URL + "/sitemap.xml")

	if len(data.URL) != 13 {
		t.Error("Get() should return Sitemap.Url(13 length)")
	}

	if err != nil {
		t.Error("Get() should not has error")
	}
}

func TestGetRecivedInvalidSitemapURL(t *testing.T) {
	server := server()
	defer server.Close()

	_, err := Get(server.URL + "/emptymap.xml")

	if err == nil {
		t.Error("Get() should return error")
	}
}

func TestGetRecivedSitemapIndexURL(t *testing.T) {
	server := server()
	defer server.Close()

	SetInterval(time.Nanosecond)
	data, err := Get(server.URL + "/sitemapindex.xml")

	if len(data.URL) != 39 {
		t.Error("Get() should return Sitemap.Url(39 length)")
	}

	if err != nil {
		t.Error("Get() should not has error")
	}
}

func TestParse(t *testing.T) {
	data, _ := ioutil.ReadFile("./testdata/sitemap.xml")
	sitemap, _ := Parse(data)

	if len(sitemap.URL) != 13 {
		t.Error("Parse() should return Sitemap.URL(13 length)")
	}
}

func TestParseIndex(t *testing.T) {
	data, _ := ioutil.ReadFile("./testdata/sitemapindex.xml")
	index, _ := ParseIndex(data)

	if len(index.Sitemap) != 3 {
		t.Error("ParseIndex() should return Index.Sitemap(3 length)")
	}
}

func TestSetInterval(t *testing.T) {
	newInterval := 3 * time.Second
	SetInterval(newInterval)

	if interval != newInterval {
		t.Error("interval should be time.Minute")
	}

	if interval == time.Second {
		t.Error("interval should not be Default(time.Second)")
	}
}

func TestSetFetch(t *testing.T) {
	f := func(url string) ([]byte, error) {
		var err error
		return []byte(url), err
	}

	SetFetch(f)

	url := "http://example.com"
	data, _ := fetch(url)

	if string(data) != url {
		t.Error("fetch(url) should return " + url)
	}
}

func BenchmarkReadSitemapXML(b *testing.B) {
	server := server()
	defer server.Close()
	Get(server.URL + "/sitemap.xml")
}

func BenchmarkReadSitemapIndex(b *testing.B) {
	server := server()
	defer server.Close()
	Get(server.URL + "/sitemapindex.xml")
}

func server() *httptest.Server {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.RequestURI == "" {
			// index page is always not found
			http.NotFound(w, r)
		}

		res, err := ioutil.ReadFile("./testdata" + r.RequestURI)
		if err != nil {
			http.NotFound(w, r)
		}
		str := strings.Replace(string(res), "HOST", r.Host, -1)
		w.WriteHeader(http.StatusOK)
		fmt.Fprintf(w, str)
	}))

	return server
}
<commit_msg>Add parse* benchmark<commit_after>package sitemap

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func TestGet(t *testing.T) {
	server := server()
	defer server.Close()

	data, err := Get(server.URL + "/sitemap.xml")

	if len(data.URL) != 13 {
		t.Error("Get() should return Sitemap.Url(13 length)")
	}

	if err != nil {
		t.Error("Get() should not has error")
	}
}

func TestGetRecivedInvalidSitemapURL(t *testing.T) {
	server := server()
	defer server.Close()

	_, err := Get(server.URL + "/emptymap.xml")

	if err == nil {
		t.Error("Get() should return error")
	}
}

func TestGetRecivedSitemapIndexURL(t *testing.T) {
	server := server()
	defer server.Close()

	SetInterval(time.Nanosecond)
	data, err := Get(server.URL + "/sitemapindex.xml")

	if len(data.URL) != 39 {
		t.Error("Get() should return Sitemap.Url(39 length)")
	}

	if err != nil {
		t.Error("Get() should not has error")
	}
}

func TestParse(t *testing.T) {
	data, _ := ioutil.ReadFile("./testdata/sitemap.xml")
	sitemap, _ := Parse(data)

	if len(sitemap.URL) != 13 {
		t.Error("Parse() should return Sitemap.URL(13 length)")
	}
}

func TestParseIndex(t *testing.T) {
	data, _ := ioutil.ReadFile("./testdata/sitemapindex.xml")
	index, _ := ParseIndex(data)

	if len(index.Sitemap) != 3 {
		t.Error("ParseIndex() should return Index.Sitemap(3 length)")
	}
}

func TestSetInterval(t *testing.T) {
	newInterval := 3 * time.Second
	SetInterval(newInterval)

	if interval != newInterval {
		t.Error("interval should be time.Minute")
	}

	if interval == time.Second {
		t.Error("interval should not be Default(time.Second)")
	}
}

func TestSetFetch(t *testing.T) {
	f := func(url string) ([]byte, error) {
		var err error
		return []byte(url), err
	}

	SetFetch(f)

	url := "http://example.com"
	data, _ := fetch(url)

	if string(data) != url {
		t.Error("fetch(url) should return " + url)
	}
}

func BenchmarkGetSitemap(b *testing.B) {
	server := server()
	defer server.Close()
	Get(server.URL + "/sitemap.xml")
}

func BenchmarkGetSitemapIndex(b *testing.B) {
	server := server()
	defer server.Close()
	Get(server.URL + "/sitemapindex.xml")
}

func BenchmarkParseSitemap(b *testing.B) {
	data, _ := ioutil.ReadFile("./testdata/sitemap.xml")
	Parse(data)
}

func BenchmarkParseSitemapIndex(b *testing.B) {
	data, _ := ioutil.ReadFile("./testdata/sitemapindex.xml")
	ParseIndex(data)
}

func server() *httptest.Server {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.RequestURI == "" {
			// index page is always not found
			http.NotFound(w, r)
		}

		res, err := ioutil.ReadFile("./testdata" + r.RequestURI)
		if err != nil {
			http.NotFound(w, r)
		}
		str := strings.Replace(string(res), "HOST", r.Host, -1)
		w.WriteHeader(http.StatusOK)
		fmt.Fprintf(w, str)
	}))

	return server
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

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

// kube2consul is a bridge between Kubernetes and Consul.  It watches the
// Kubernetes master for changes in Services and creates new DNS records on the
// consul agent.

package main // import "github.com/jmccarty3/kube2consul"

import (
	"flag"
	"fmt"
	"strings"
	"net/url"
	"os"
	"time"
	"reflect"

	kapi "k8s.io/kubernetes/pkg/api"
	kclient "k8s.io/kubernetes/pkg/client/unversioned"
	kcache "k8s.io/kubernetes/pkg/client/unversioned/cache"
	kclientcmd "k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	kframework "k8s.io/kubernetes/pkg/controller/framework"
	kcontrollerFramework "k8s.io/kubernetes/pkg/controller/framework"
	kSelector "k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/util"
	"github.com/golang/glog"
	consulapi "github.com/hashicorp/consul/api"

  "k8s.io/kubernetes/pkg/api"
  "k8s.io/kubernetes/pkg/fields"
  "k8s.io/kubernetes/pkg/labels"
)

var (
	argConsulAgent         = flag.String("consul-agent", "http://127.0.0.1:8500", "URL to consul agent")
	argKubecfgFile         = flag.String("kubecfg_file", "", "Location of kubecfg file for access to kubernetes service")
	argKubeMasterUrl       = flag.String("kube_master_url", "https://${KUBERNETES_SERVICE_HOST}:${KUBERNETES_SERVICE_PORT}", "Url to reach kubernetes master. Env variables in this flag will be expanded.")
)

const (
	// Maximum number of attempts to connect to consul server.
	maxConnectAttempts = 12
	// Resync period for the kube controller loop.
	resyncPeriod = 5 * time.Second
)

type nodeInformation struct {
	name string

	address string

	ready bool
	// map[service] DNS IDs
	ids map[string][]string
}

type kube2consul struct {
	// Consul client.
	consulClient *consulapi.Client

	// DNS domain name.
	domain string

	//Nodes Name / valid
	nodes map[string]nodeInformation

	//All Services.
	services map[string]*kapi.Service
}

func Newkube2consul() *kube2consul {
	var k kube2consul
	k.nodes = make(map[string]nodeInformation)
	k.services = make(map[string]*kapi.Service)

	return &k
}

func NewnodeInformation() *nodeInformation {
	var n nodeInformation
	n.ready = false
	n.ids = make(map[string][]string)

	return &n
}

func Contains(s []string, e string) bool {
	for _,i := range s {
		if i == e {
			return true
		}
	}
	return false
}

func (ks *kube2consul) removeDNS(recordID string) error {
	glog.Infof("Removing %s from DNS", recordID)
	ks.consulClient.Agent().ServiceDeregister(recordID)
	return nil
}

func (ks *kube2consul) createDNS(record string, service *kapi.Service, node *nodeInformation) error {
	if strings.Contains(record, ".") {
		glog.Infof("Service names containing '.' are not supported: %s\n", service.Name)
		return nil
	}

	// if ClusterIP is not set, do not create a DNS records
	if !kapi.IsServiceIPSet(service) {
		glog.Infof("Skipping dns record for headless service: %s\n", service.Name)
		return nil
	}

	for i := range service.Spec.Ports {
			newId := node.name+record + service.Spec.Ports[i].Name

			if len(service.Spec.Ports[i].Name) > 0 {
				asrName := record + "-" + service.Spec.Ports[i].Name
			}
			else {
				asrName := record
			}

			asr := &consulapi.AgentServiceRegistration{
				ID:			 newId,
				Name: 	 asrName,
				Address: node.address,
				Port:    service.Spec.Ports[i].NodePort,
				Tags: []string{"Kube"},
			}

			if Contains(node.ids[record], newId) == false {
				glog.Infof("Setting DNS record: %v -> %v:%d\n", asr.Name, asr.Address, asr.Port)
				if err := ks.consulClient.Agent().ServiceRegister(asr); err != nil {
					return err
				}

				node.ids[record] = append(node.ids[record], newId)
		}
	}
	return nil
}

func newConsulClient(consulAgent string) (*consulapi.Client, error) {
	var (
		client *consulapi.Client
		err    error
	)

	consulConfig := consulapi.DefaultConfig()
	consulAgentUrl, err := url.Parse(consulAgent)
	if err != nil {
			glog.Infof("Error parsing Consul url")
			return nil, err
	}

	if consulAgentUrl.Host != "" {
	  consulConfig.Address = consulAgentUrl.Host
	}

	if consulAgentUrl.Scheme != "" {
		consulConfig.Scheme = consulAgentUrl.Scheme
	}

	client, err = consulapi.NewClient(consulConfig)
	if err != nil {
			glog.Infof("Error creating Consul client")
			return nil, err
	}

	for attempt := 1; attempt <= maxConnectAttempts; attempt++ {
		if _, err = client.Agent().Self(); err == nil {
			break
		}

		if attempt == maxConnectAttempts {
			break
		}

		glog.Infof("[Attempt: %d] Attempting access to Consul after 5 second sleep", attempt)
		time.Sleep(5 * time.Second)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to connect to Consul agent: %v, error: %v", consulAgent, err)
	}
	glog.Infof("Consul agent found: %v", consulAgent)

	return client, nil
}

func getKubeMasterUrl() (string, error) {
	if *argKubeMasterUrl == "" {
		return "", fmt.Errorf("no --kube_master_url specified")
	}
	parsedUrl, err := url.Parse(os.ExpandEnv(*argKubeMasterUrl))
	if err != nil {
		return "", fmt.Errorf("failed to parse --kube_master_url %s - %v", *argKubeMasterUrl, err)
	}
	if parsedUrl.Scheme == "" || parsedUrl.Host == "" || parsedUrl.Host == ":" {
		return "", fmt.Errorf("invalid --kube_master_url specified %s", *argKubeMasterUrl)
	}
	return parsedUrl.String(), nil
}

// TODO: evaluate using pkg/client/clientcmd
func newKubeClient() (*kclient.Client, error) {
	var config *kclient.Config
	masterUrl, err := getKubeMasterUrl()
	if err != nil {
		return nil, err
	}
	if *argKubecfgFile == "" {
		config = &kclient.Config{
			Host:    masterUrl,
			Version: "v1",
		}
	} else {
		var err error
		overrides := &kclientcmd.ConfigOverrides{}
		overrides.ClusterInfo.Server = masterUrl                                     // might be "", but that is OK
		rules := &kclientcmd.ClientConfigLoadingRules{ExplicitPath: *argKubecfgFile} // might be "", but that is OK
		if config, err = kclientcmd.NewNonInteractiveDeferredLoadingClientConfig(rules, overrides).ClientConfig(); err != nil {
			return nil, err
		}
	}
	glog.Infof("Using %s for kubernetes master", config.Host)
	glog.Infof("Using kubernetes API %s", config.Version)
	return kclient.New(config)
}

func buildNameString(service, namespace string) string {
	//glog.Infof("Name String: %s  %s", service, namespace)
	//return fmt.Sprintf("%s.%s", service, namespace)
	return fmt.Sprintf("%s", service)
}

// Returns a cache.ListWatch that gets all changes to services.
func createServiceLW(kubeClient *kclient.Client) *kcache.ListWatch {
	return kcache.NewListWatchFromClient(kubeClient, "services", kapi.NamespaceAll, kSelector.Everything())
}

// Returns a cache.ListWatch that gets all changes to services.
func createNodeLW(kubeClient *kclient.Client) *kcache.ListWatch {
	return kcache.NewListWatchFromClient(kubeClient, "nodes", kapi.NamespaceAll, kSelector.Everything())
}

func (ks *kube2consul) newService(obj interface{}) {
	if s, ok := obj.(*kapi.Service); ok {
		name := buildNameString(s.Name, s.Namespace)
		ks.services[name] = s

		glog.V(2).Info("Creating Service: ", name)
		//Add to all existing nodes
		for _,node := range ks.nodes {
			if node.ready {
				ks.createDNS(name, s, &node)
			}
		}
	}
}

func (ks *kube2consul) removeService(obj interface{}) {
	if s, ok := obj.(*kapi.Service); ok {
		name := buildNameString(s.Name, s.Namespace)
		//Remove service from node
		for _,node := range ks.nodes {
			for _,id := range node.ids[name] {
				ks.removeDNS(id)
			}
			delete(node.ids, name)
		}
		delete(ks.services,name)
	}
}

func (ks *kube2consul) updateService(oldObj, newObj interface{}) {
	if old, ok := oldObj.(*kapi.Service); ok {
		if new, ok := newObj.(*kapi.Service); ok {
			if reflect.DeepEqual(*old, *new) == false {
				ks.removeService(old)
				ks.newService(new)
			}
		}
	}
}

func (ks *kube2consul) updateNode(oldObj, newObj interface{}) {
	if n, ok := newObj.(*kapi.Node); ok {
		ready := n.Status.Conditions[0].Status == kapi.ConditionTrue
		nodeInfo := ks.nodes[n.Name]

		if nodeInfo.ready != ready {
			nodeInfo.ready = ready
			if ready {
				for serviceName, service := range ks.services {
					ks.createDNS(serviceName, service, &nodeInfo)
				}
			}	else {
				for _,serviceIDs := range nodeInfo.ids {
					for _,serviceID := range serviceIDs {
						ks.removeDNS(serviceID)
					}
				}
				nodeInfo.ids = make(map[string][]string) //Clear tha map
			}
		}
	}
}

func (ks *kube2consul) newNode(newObj interface{}) {
	if node, ok := newObj.(*kapi.Node); ok {
		if _, ok := ks.nodes[node.Name]; !ok {
			glog.Info("Adding Node: ", node.Name)
			var newNode = *NewnodeInformation()
			newNode.name = node.Name
			newNode.address = node.Status.Addresses[0].Address

			ks.nodes[node.Name] = newNode
		}
	}
}

func (ks *kube2consul) removeNode(oldObj interface{}) {
	if node, ok := oldObj.(*kapi.Node); ok {
		if info, ok := ks.nodes[node.Name]; ok {
			glog.Info("Removing Node: ", node.Name)
			for _,idSet := range info.ids {
				for _,id := range idSet {
					ks.removeDNS(id)
				}
			}

			delete(ks.nodes, node.Name)
		}
	}
}

func watchForServices(kubeClient *kclient.Client, ks *kube2consul) {
	var serviceController *kcontrollerFramework.Controller
	_, serviceController = kframework.NewInformer(
		createServiceLW(kubeClient),
		&kapi.Service{},
		resyncPeriod,
		kframework.ResourceEventHandlerFuncs{
			AddFunc:    ks.newService,
			DeleteFunc: ks.removeService,
			UpdateFunc: ks.updateService,
		},
	)
	go serviceController.Run(util.NeverStop)
}

func watchForNodes(kubeClient *kclient.Client, ks *kube2consul) kcache.Store {
	store, serviceController := kframework.NewInformer(
		createNodeLW(kubeClient),
		&kapi.Node{},
		resyncPeriod,
		kframework.ResourceEventHandlerFuncs{
			AddFunc: ks.newNode,
			DeleteFunc: ks.removeNode,
			UpdateFunc: ks.updateNode,
		},
	)
	glog.Info("About to call run!")
	go serviceController.Run(util.NeverStop)
	return store
}

func main() {
	flag.Parse()
	var err error
	// TODO: Validate input flags.
	ks := Newkube2consul()

	if ks.consulClient, err = newConsulClient(*argConsulAgent); err != nil {
		glog.Fatalf("Failed to create Consul client - %v", err)
	}

	kubeClient, err := newKubeClient()
	if err != nil {
		glog.Fatalf("Failed to create a kubernetes client: %v", err)
	}

	glog.Info(kubeClient.ServerVersion())
	glog.Info(kubeClient.Services(kapi.NamespaceAll).Get("sensu-core"))

	pods, err := kubeClient.Pods(api.NamespaceDefault).List(labels.Everything(), fields.Everything())
	if err != nil {
	  for pod := range pods.Items {
			glog.Info(pod)
		}
	}

	watchForServices(kubeClient, ks)
	watchForNodes(kubeClient, ks)
	glog.Info("Watchers running")
	select {}
}
<commit_msg>Updating for kube API changes<commit_after>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

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

// kube2consul is a bridge between Kubernetes and Consul.  It watches the
// Kubernetes master for changes in Services and creates new DNS records on the
// consul agent.

package main // import "github.com/jmccarty3/kube2consul"

import (
	"flag"
	"fmt"
	"strings"
	"net/url"
	"os"
	"time"
	"reflect"

	kapi "k8s.io/kubernetes/pkg/api"
	kclient "k8s.io/kubernetes/pkg/client/unversioned"
	kcache "k8s.io/kubernetes/pkg/client/cache"
	kclientcmd "k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	kframework "k8s.io/kubernetes/pkg/controller/framework"
	kcontrollerFramework "k8s.io/kubernetes/pkg/controller/framework"
	kSelector "k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/util"
	"github.com/golang/glog"
	consulapi "github.com/hashicorp/consul/api"

  "k8s.io/kubernetes/pkg/api"
  "k8s.io/kubernetes/pkg/fields"
  "k8s.io/kubernetes/pkg/labels"
)

var (
	argConsulAgent         = flag.String("consul-agent", "http://127.0.0.1:8500", "URL to consul agent")
	argKubecfgFile         = flag.String("kubecfg_file", "", "Location of kubecfg file for access to kubernetes service")
	argKubeMasterUrl       = flag.String("kube_master_url", "https://${KUBERNETES_SERVICE_HOST}:${KUBERNETES_SERVICE_PORT}", "Url to reach kubernetes master. Env variables in this flag will be expanded.")
)

const (
	// Maximum number of attempts to connect to consul server.
	maxConnectAttempts = 12
	// Resync period for the kube controller loop.
	resyncPeriod = 5 * time.Second
)

type nodeInformation struct {
	name string

	address string

	ready bool
	// map[service] DNS IDs
	ids map[string][]string
}

type kube2consul struct {
	// Consul client.
	consulClient *consulapi.Client

	// DNS domain name.
	domain string

	//Nodes Name / valid
	nodes map[string]nodeInformation

	//All Services.
	services map[string]*kapi.Service
}

func Newkube2consul() *kube2consul {
	var k kube2consul
	k.nodes = make(map[string]nodeInformation)
	k.services = make(map[string]*kapi.Service)

	return &k
}

func NewnodeInformation() *nodeInformation {
	var n nodeInformation
	n.ready = false
	n.ids = make(map[string][]string)

	return &n
}

func Contains(s []string, e string) bool {
	for _,i := range s {
		if i == e {
			return true
		}
	}
	return false
}

func (ks *kube2consul) removeDNS(recordID string) error {
	glog.Infof("Removing %s from DNS", recordID)
	ks.consulClient.Agent().ServiceDeregister(recordID)
	return nil
}

func (ks *kube2consul) createDNS(record string, service *kapi.Service, node *nodeInformation) error {
	if strings.Contains(record, ".") {
		glog.Infof("Service names containing '.' are not supported: %s\n", service.Name)
		return nil
	}

	// if ClusterIP is not set, do not create a DNS records
	if !kapi.IsServiceIPSet(service) {
		glog.Infof("Skipping dns record for headless service: %s\n", service.Name)
		return nil
	}

	for i := range service.Spec.Ports {
			newId := node.name+record + service.Spec.Ports[i].Name
      var asrName string

			if len(service.Spec.Ports[i].Name) > 0 {
				asrName = record + "-" + service.Spec.Ports[i].Name
			} else {
				asrName = record
			}

			asr := &consulapi.AgentServiceRegistration{
				ID:			 newId,
				Name: 	 asrName,
				Address: node.address,
				Port:    service.Spec.Ports[i].NodePort,
				Tags: []string{"Kube"},
			}

			if Contains(node.ids[record], newId) == false {
				glog.Infof("Setting DNS record: %v -> %v:%d\n", asr.Name, asr.Address, asr.Port)
				if err := ks.consulClient.Agent().ServiceRegister(asr); err != nil {
					return err
				}

				node.ids[record] = append(node.ids[record], newId)
		}
	}
	return nil
}

func newConsulClient(consulAgent string) (*consulapi.Client, error) {
	var (
		client *consulapi.Client
		err    error
	)

	consulConfig := consulapi.DefaultConfig()
	consulAgentUrl, err := url.Parse(consulAgent)
	if err != nil {
			glog.Infof("Error parsing Consul url")
			return nil, err
	}

	if consulAgentUrl.Host != "" {
	  consulConfig.Address = consulAgentUrl.Host
	}

	if consulAgentUrl.Scheme != "" {
		consulConfig.Scheme = consulAgentUrl.Scheme
	}

	client, err = consulapi.NewClient(consulConfig)
	if err != nil {
			glog.Infof("Error creating Consul client")
			return nil, err
	}

	for attempt := 1; attempt <= maxConnectAttempts; attempt++ {
		if _, err = client.Agent().Self(); err == nil {
			break
		}

		if attempt == maxConnectAttempts {
			break
		}

		glog.Infof("[Attempt: %d] Attempting access to Consul after 5 second sleep", attempt)
		time.Sleep(5 * time.Second)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to connect to Consul agent: %v, error: %v", consulAgent, err)
	}
	glog.Infof("Consul agent found: %v", consulAgent)

	return client, nil
}

func getKubeMasterUrl() (string, error) {
	if *argKubeMasterUrl == "" {
		return "", fmt.Errorf("no --kube_master_url specified")
	}
	parsedUrl, err := url.Parse(os.ExpandEnv(*argKubeMasterUrl))
	if err != nil {
		return "", fmt.Errorf("failed to parse --kube_master_url %s - %v", *argKubeMasterUrl, err)
	}
	if parsedUrl.Scheme == "" || parsedUrl.Host == "" || parsedUrl.Host == ":" {
		return "", fmt.Errorf("invalid --kube_master_url specified %s", *argKubeMasterUrl)
	}
	return parsedUrl.String(), nil
}

// TODO: evaluate using pkg/client/clientcmd
func newKubeClient() (*kclient.Client, error) {
	var config *kclient.Config
	masterUrl, err := getKubeMasterUrl()
	if err != nil {
		return nil, err
	}
	if *argKubecfgFile == "" {
		config = &kclient.Config{
			Host:    masterUrl,
			Version: "v1",
		}
	} else {
		var err error
		overrides := &kclientcmd.ConfigOverrides{}
		overrides.ClusterInfo.Server = masterUrl                                     // might be "", but that is OK
		rules := &kclientcmd.ClientConfigLoadingRules{ExplicitPath: *argKubecfgFile} // might be "", but that is OK
		if config, err = kclientcmd.NewNonInteractiveDeferredLoadingClientConfig(rules, overrides).ClientConfig(); err != nil {
			return nil, err
		}
	}
	glog.Infof("Using %s for kubernetes master", config.Host)
	glog.Infof("Using kubernetes API %s", config.Version)
	return kclient.New(config)
}

func buildNameString(service, namespace string) string {
	//glog.Infof("Name String: %s  %s", service, namespace)
	//return fmt.Sprintf("%s.%s", service, namespace)
	return fmt.Sprintf("%s", service)
}

// Returns a cache.ListWatch that gets all changes to services.
func createServiceLW(kubeClient *kclient.Client) *kcache.ListWatch {
	return kcache.NewListWatchFromClient(kubeClient, "services", kapi.NamespaceAll, kSelector.Everything())
}

// Returns a cache.ListWatch that gets all changes to services.
func createNodeLW(kubeClient *kclient.Client) *kcache.ListWatch {
	return kcache.NewListWatchFromClient(kubeClient, "nodes", kapi.NamespaceAll, kSelector.Everything())
}

func (ks *kube2consul) newService(obj interface{}) {
	if s, ok := obj.(*kapi.Service); ok {
		name := buildNameString(s.Name, s.Namespace)
		ks.services[name] = s

		glog.V(2).Info("Creating Service: ", name)
		//Add to all existing nodes
		for _,node := range ks.nodes {
			if node.ready {
				ks.createDNS(name, s, &node)
			}
		}
	}
}

func (ks *kube2consul) removeService(obj interface{}) {
	if s, ok := obj.(*kapi.Service); ok {
		name := buildNameString(s.Name, s.Namespace)
		//Remove service from node
		for _,node := range ks.nodes {
			for _,id := range node.ids[name] {
				ks.removeDNS(id)
			}
			delete(node.ids, name)
		}
		delete(ks.services,name)
	}
}

func (ks *kube2consul) updateService(oldObj, newObj interface{}) {
	if old, ok := oldObj.(*kapi.Service); ok {
		if new, ok := newObj.(*kapi.Service); ok {
			if reflect.DeepEqual(*old, *new) == false {
				ks.removeService(old)
				ks.newService(new)
			}
		}
	}
}

func (ks *kube2consul) updateNode(oldObj, newObj interface{}) {
	if n, ok := newObj.(*kapi.Node); ok {
		ready := n.Status.Conditions[0].Status == kapi.ConditionTrue
		nodeInfo := ks.nodes[n.Name]

		if nodeInfo.ready != ready {
			nodeInfo.ready = ready
			if ready {
				for serviceName, service := range ks.services {
					ks.createDNS(serviceName, service, &nodeInfo)
				}
			}	else {
				for _,serviceIDs := range nodeInfo.ids {
					for _,serviceID := range serviceIDs {
						ks.removeDNS(serviceID)
					}
				}
				nodeInfo.ids = make(map[string][]string) //Clear tha map
			}
		}
	}
}

func (ks *kube2consul) newNode(newObj interface{}) {
	if node, ok := newObj.(*kapi.Node); ok {
		if _, ok := ks.nodes[node.Name]; !ok {
			glog.Info("Adding Node: ", node.Name)
			var newNode = *NewnodeInformation()
			newNode.name = node.Name
			newNode.address = node.Status.Addresses[0].Address

			ks.nodes[node.Name] = newNode
		}
	}
}

func (ks *kube2consul) removeNode(oldObj interface{}) {
	if node, ok := oldObj.(*kapi.Node); ok {
		if info, ok := ks.nodes[node.Name]; ok {
			glog.Info("Removing Node: ", node.Name)
			for _,idSet := range info.ids {
				for _,id := range idSet {
					ks.removeDNS(id)
				}
			}

			delete(ks.nodes, node.Name)
		}
	}
}

func watchForServices(kubeClient *kclient.Client, ks *kube2consul) {
	var serviceController *kcontrollerFramework.Controller
	_, serviceController = kframework.NewInformer(
		createServiceLW(kubeClient),
		&kapi.Service{},
		resyncPeriod,
		kframework.ResourceEventHandlerFuncs{
			AddFunc:    ks.newService,
			DeleteFunc: ks.removeService,
			UpdateFunc: ks.updateService,
		},
	)
	go serviceController.Run(util.NeverStop)
}

func watchForNodes(kubeClient *kclient.Client, ks *kube2consul) kcache.Store {
	store, serviceController := kframework.NewInformer(
		createNodeLW(kubeClient),
		&kapi.Node{},
		resyncPeriod,
		kframework.ResourceEventHandlerFuncs{
			AddFunc: ks.newNode,
			DeleteFunc: ks.removeNode,
			UpdateFunc: ks.updateNode,
		},
	)
	glog.Info("About to call run!")
	go serviceController.Run(util.NeverStop)
	return store
}

func main() {
	flag.Parse()
	var err error
	// TODO: Validate input flags.
	ks := Newkube2consul()

	if ks.consulClient, err = newConsulClient(*argConsulAgent); err != nil {
		glog.Fatalf("Failed to create Consul client - %v", err)
	}

	kubeClient, err := newKubeClient()
	if err != nil {
		glog.Fatalf("Failed to create a kubernetes client: %v", err)
	}

	glog.Info(kubeClient.ServerVersion())
	glog.Info(kubeClient.Services(kapi.NamespaceAll).Get("sensu-core"))

	pods, err := kubeClient.Pods(api.NamespaceDefault).List(labels.Everything(), fields.Everything())
	if err != nil {
	  for pod := range pods.Items {
			glog.Info(pod)
		}
	}

	watchForServices(kubeClient, ks)
	watchForNodes(kubeClient, ks)
	glog.Info("Watchers running")
	select {}
}
<|endoftext|>