<commit_before>package dbr

// UpdateStmt builds `UPDATE ...`
type UpdateStmt interface {
	Builder

	Where(query interface{}, value ...interface{}) UpdateStmt
	Set(column string, value interface{}) UpdateStmt
	SetMap(m map[string]interface{}) UpdateStmt
}

type updateStmt struct {
	raw

	Table     string
	Value     map[string]interface{}
	WhereCond []Builder
}

// Build builds `UPDATE ...` in dialect
func (b *updateStmt) Build(d Dialect, buf Buffer) error {
	if b.raw.Query != "" {
		return b.raw.Build(d, buf)
	}

	if b.Table == "" {
		return ErrTableNotSpecified
	}

	if len(b.Value) == 0 {
		return ErrColumnNotSpecified
	}

	buf.WriteString("UPDATE ")
	buf.WriteString(d.QuoteIdent(b.Table))
	buf.WriteString(" SET ")

	i := 0
	for col, v := range b.Value {
		if i > 0 {
			buf.WriteString(", ")
		}
		buf.WriteString(d.QuoteIdent(col))
		buf.WriteString(" = ")
		buf.WriteString(placeholder)

		buf.WriteValue(v)
		i++
	}

	if len(b.WhereCond) > 0 {
		buf.WriteString(" WHERE ")
		err := And(b.WhereCond...).Build(d, buf)
		if err != nil {
			return err
		}
	}
	return nil
}

// Update creates an UpdateStmt
func Update(table string) UpdateStmt {
	return createUpdateStmt(table)
}

func createUpdateStmt(table string) *updateStmt {
	return &updateStmt{
		Table: table,
		Value: make(map[string]interface{}),
	}
}

// UpdateBySql creates an UpdateStmt with raw query
func UpdateBySql(query string, value ...interface{}) UpdateStmt {
	return createUpdateStmtBySQL(query, value)
}

func createUpdateStmtBySQL(query string, value []interface{}) *updateStmt {
	return &updateStmt{
		raw: raw{
			Query: query,
			Value: value,
		},
		Value: make(map[string]interface{}),
	}
}

// Where adds a where condition
func (b *updateStmt) Where(query interface{}, value ...interface{}) UpdateStmt {
	switch query := query.(type) {
	case string:
		b.WhereCond = append(b.WhereCond, Expr(query, value...))
	case Builder:
		b.WhereCond = append(b.WhereCond, query)
	}
	return b
}

// Set specifies a key-value pair
func (b *updateStmt) Set(column string, value interface{}) UpdateStmt {
	b.Value[column] = value
	return b
}

// SetMap specifies a list of key-value pair
func (b *updateStmt) SetMap(m map[string]interface{}) UpdateStmt {
	for col, val := range m {
		b.Set(col, val)
	}
	return b
}
<commit_msg>update: add Record() for updating from struct values<commit_after>package dbr

import "reflect"

// UpdateStmt builds `UPDATE ...`
type UpdateStmt interface {
	Builder

	Where(query interface{}, value ...interface{}) UpdateStmt
	Set(column string, value interface{}) UpdateStmt
	SetMap(m map[string]interface{}) UpdateStmt
}

type updateStmt struct {
	raw

	Table     string
	Value     map[string]interface{}
	WhereCond []Builder
}

// Build builds `UPDATE ...` in dialect
func (b *updateStmt) Build(d Dialect, buf Buffer) error {
	if b.raw.Query != "" {
		return b.raw.Build(d, buf)
	}

	if b.Table == "" {
		return ErrTableNotSpecified
	}

	if len(b.Value) == 0 {
		return ErrColumnNotSpecified
	}

	buf.WriteString("UPDATE ")
	buf.WriteString(d.QuoteIdent(b.Table))
	buf.WriteString(" SET ")

	i := 0
	for col, v := range b.Value {
		if i > 0 {
			buf.WriteString(", ")
		}
		buf.WriteString(d.QuoteIdent(col))
		buf.WriteString(" = ")
		buf.WriteString(placeholder)

		buf.WriteValue(v)
		i++
	}

	if len(b.WhereCond) > 0 {
		buf.WriteString(" WHERE ")
		err := And(b.WhereCond...).Build(d, buf)
		if err != nil {
			return err
		}
	}
	return nil
}

// Update creates an UpdateStmt
func Update(table string) UpdateStmt {
	return createUpdateStmt(table)
}

func createUpdateStmt(table string) *updateStmt {
	return &updateStmt{
		Table: table,
		Value: make(map[string]interface{}),
	}
}

// UpdateBySql creates an UpdateStmt with raw query
func UpdateBySql(query string, value ...interface{}) UpdateStmt {
	return createUpdateStmtBySQL(query, value)
}

func createUpdateStmtBySQL(query string, value []interface{}) *updateStmt {
	return &updateStmt{
		raw: raw{
			Query: query,
			Value: value,
		},
		Value: make(map[string]interface{}),
	}
}

// Where adds a where condition
func (b *updateStmt) Where(query interface{}, value ...interface{}) UpdateStmt {
	switch query := query.(type) {
	case string:
		b.WhereCond = append(b.WhereCond, Expr(query, value...))
	case Builder:
		b.WhereCond = append(b.WhereCond, query)
	}
	return b
}

// Set specifies a key-value pair
func (b *updateStmt) Set(column string, value interface{}) UpdateStmt {
	b.Value[column] = value
	return b
}

// SetMap specifies a list of key-value pair
func (b *updateStmt) SetMap(m map[string]interface{}) UpdateStmt {
	for col, val := range m {
		b.Set(col, val)
	}
	return b
}

// SetRecord specifies a record with field and values to set
func (b *updateStmt) SetRecord(structValue interface{}) UpdateStmt {
	v := reflect.Indirect(reflect.ValueOf(structValue))
	if v.Kind() == reflect.Struct {
		sm := structMap(v.Type())
		// populate columns from available record fields
		// if no columns were specified up to this point

		for col, index := range sm {
			b.Set(col, v.FieldByIndex(index).Interface())
		}
	}

	return b
}
<|endoftext|>


<commit_before>package main

import (
	"compress/gzip"
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"
)

var updateTopic = &Topic{
	Name:      "update",
	ShortHelp: "update heroku-cli",
}

var updateCmd = &Command{
	Topic:     "update",
	ShortHelp: "updates heroku-cli",
	Args:      []Arg{{Name: "channel", Optional: true}},
	Run: func(ctx *Context) {
		channel := ctx.Args["channel"]
		if channel == "" {
			channel = "master"
		}
		Print("updating to ", channel)
	},
}

var binPath = filepath.Join(AppDir, "heroku-cli")

// UpdateIfNeeded checks for and performs an autoupdate if there is a new version out.
func UpdateIfNeeded() {
	if !updateNeeded() {
		return
	}
	// TODO: update plugins
	manifest := getUpdateManifest()
	if manifest.Version == Version {
		// Set timestamp of bin so we don't update again
		os.Chtimes(binPath, time.Now(), time.Now())
		return
	}
	if !updatable() {
		Errf("Out of date: You are running %s but %s is out.\n", Version, manifest.Version)
		return
	}
	// Leave out updating text until heroku-cli is used in place of ruby cli
	// So it doesn't confuse users with 2 different version numbers
	//Errf("Updating to %s... ", manifest.Version)
	build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
	update(build.URL, build.Sha1)
	//Errln("done")
	execBin()
	os.Exit(0)
}

func updateNeeded() bool {
	if Version == "dev" {
		return false
	}
	f, err := os.Stat(binPath)
	if err != nil {
		Errln("WARNING: Error autoupdating. The CLI will not be able to update itself.")
		return false
	}
	// TODO: Increase the autoupdate time later
	return f.ModTime().Add(2 * time.Minute).Before(time.Now())
}

type manifest struct {
	Channel, Version string
	Builds           map[string]map[string]struct {
		URL, Sha1 string
	}
}

func getUpdateManifest() manifest {
	res, err := http.Get("https://d1gvo455cekpjp.cloudfront.net/heroku-cli/" + Channel + "/manifest.json")
	if err != nil {
		panic(err)
	}
	var m manifest
	json.NewDecoder(res.Body).Decode(&m)
	return m
}

func updatable() bool {
	path, err := filepath.Abs(os.Args[0])
	if err != nil {
		Errln(err)
	}
	return path == binPath
}

func update(url, sha1 string) {
	tmp, err := downloadBin(url)
	if err != nil {
		panic(err)
	}
	if fileSha1(tmp) != sha1 {
		panic("SHA mismatch")
	}
	if err := os.Rename(tmp, binPath); err != nil {
		panic(err)
	}
}

func downloadBin(url string) (string, error) {
	out, err := os.OpenFile(binPath+"~", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return "", err
	}
	defer out.Close()
	client := &http.Client{}
	req, err := http.NewRequest("GET", url+".gz", nil)
	if err != nil {
		return "", err
	}
	req.Header.Add("Accept-Encoding", "gzip")
	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	uncompressed, err := gzip.NewReader(resp.Body)
	if err != nil {
		return "", err
	}
	_, err = io.Copy(out, uncompressed)
	return out.Name(), err
}

func fileSha1(path string) string {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", sha1.Sum(data))
}

func execBin() {
	cmd := exec.Command(binPath, os.Args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		panic(err)
	}
}
<commit_msg>implemented update command<commit_after>package main

import (
	"compress/gzip"
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"
)

var updateTopic = &Topic{
	Name:      "update",
	ShortHelp: "update heroku-cli",
}

var updateCmd = &Command{
	Topic:     "update",
	ShortHelp: "updates heroku-cli",
	Args:      []Arg{{Name: "channel", Optional: true}},
	Run: func(ctx *Context) {
		channel := ctx.Args["channel"]
		if channel == "" {
			channel = "master"
		}
		manifest := getUpdateManifest(channel)
		build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
		Errf("updating to %s (%s)... ", manifest.Version, manifest.Channel)
		update(build.URL, build.Sha1)
		Errln("done")
	},
}

var binPath = filepath.Join(AppDir, "heroku-cli")

// UpdateIfNeeded checks for and performs an autoupdate if there is a new version out.
func UpdateIfNeeded() {
	if !updateNeeded() {
		return
	}
	// TODO: update plugins
	manifest := getUpdateManifest(Channel)
	if manifest.Version == Version {
		// Set timestamp of bin so we don't update again
		os.Chtimes(binPath, time.Now(), time.Now())
		return
	}
	if !updatable() {
		Errf("Out of date: You are running %s but %s is out.\n", Version, manifest.Version)
		return
	}
	// Leave out updating text until heroku-cli is used in place of ruby cli
	// So it doesn't confuse users with 2 different version numbers
	//Errf("Updating to %s... ", manifest.Version)
	build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
	update(build.URL, build.Sha1)
	//Errln("done")
	execBin()
	os.Exit(0)
}

func updateNeeded() bool {
	if Version == "dev" {
		return false
	}
	f, err := os.Stat(binPath)
	if err != nil {
		Errln("WARNING: Error autoupdating. The CLI will not be able to update itself.")
		return false
	}
	// TODO: Increase the autoupdate time later
	return f.ModTime().Add(2 * time.Minute).Before(time.Now())
}

type manifest struct {
	Channel, Version string
	Builds           map[string]map[string]struct {
		URL, Sha1 string
	}
}

func getUpdateManifest(channel string) manifest {
	res, err := http.Get("https://d1gvo455cekpjp.cloudfront.net/heroku-cli/" + channel + "/manifest.json")
	if err != nil {
		panic(err)
	}
	var m manifest
	json.NewDecoder(res.Body).Decode(&m)
	return m
}

func updatable() bool {
	path, err := filepath.Abs(os.Args[0])
	if err != nil {
		Errln(err)
	}
	return path == binPath
}

func update(url, sha1 string) {
	tmp, err := downloadBin(url)
	if err != nil {
		panic(err)
	}
	if fileSha1(tmp) != sha1 {
		panic("SHA mismatch")
	}
	if err := os.Rename(tmp, binPath); err != nil {
		panic(err)
	}
}

func downloadBin(url string) (string, error) {
	out, err := os.OpenFile(binPath+"~", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return "", err
	}
	defer out.Close()
	client := &http.Client{}
	req, err := http.NewRequest("GET", url+".gz", nil)
	if err != nil {
		return "", err
	}
	req.Header.Add("Accept-Encoding", "gzip")
	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	uncompressed, err := gzip.NewReader(resp.Body)
	if err != nil {
		return "", err
	}
	_, err = io.Copy(out, uncompressed)
	return out.Name(), err
}

func fileSha1(path string) string {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", sha1.Sum(data))
}

func execBin() {
	cmd := exec.Command(binPath, os.Args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package wrecker

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
)

type Wrecker struct {
	BaseURL    string
	HttpClient *http.Client
}

const (
	GET    = "GET"
	POST   = "POST"
	PUT    = "PUT"
	DELETE = "DELETE"
)

func (w *Wrecker) Get(endpoint string, params url.Values, response interface{}) error {
	return w.SendRequest(GET, endpoint, params, response)
}

func (w *Wrecker) Post(endpoint string, params url.Values, response interface{}) error {
	return w.SendRequest(POST, endpoint, params, response)
}

func (w *Wrecker) Put(endpoint string, params url.Values, response interface{}) error {
	return w.SendRequest(PUT, endpoint, params, response)
}

func (w *Wrecker) Delete(endpoint string, response interface{}) error {
	return w.SendRequest(DELETE, endpoint, nil, response)
}

func (w *Wrecker) SendRequest(verb string, endpoint string, params url.Values, response interface{}) error {
	// Preparing Request
	requestURL := strings.Join([]string{w.BaseURL, endpoint}, "")
	paramsReader := strings.NewReader(params.Encode())
	clientReq, err := http.NewRequest(verb, requestURL, paramsReader)
	if err != nil {
		return err
	}
	clientReq.Header.Add("Content-Type", "application/x-www-form-urlencoded")

	// Executing request
	clientRes, err := w.HttpClient.Do(clientReq)
	if err != nil {
		return err
	}
	defer clientRes.Body.Close()

	// Packing into response
	body, err := ioutil.ReadAll(clientRes.Body)
	if err != nil {
		return err
	}
	err = json.Unmarshal(body, &response)
	return err
}
<commit_msg>Handing GET with url params<commit_after>package wrecker

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
)

type Wrecker struct {
	BaseURL    string
	HttpClient *http.Client
}

const (
	GET    = "GET"
	POST   = "POST"
	PUT    = "PUT"
	DELETE = "DELETE"
)

func (w *Wrecker) Get(endpoint string, params url.Values, response interface{}) error {
	if params != nil {
		queryString := "?" + params.Encode()
		endpoint = strings.Join([]string{endpoint, queryString}, "")
	}
	return w.SendRequest(GET, endpoint, nil, response)
}

func (w *Wrecker) Post(endpoint string, params url.Values, response interface{}) error {
	return w.SendRequest(POST, endpoint, params, response)
}

func (w *Wrecker) Put(endpoint string, params url.Values, response interface{}) error {
	return w.SendRequest(PUT, endpoint, params, response)
}

func (w *Wrecker) Delete(endpoint string, response interface{}) error {
	return w.SendRequest(DELETE, endpoint, nil, response)
}

func (w *Wrecker) SendRequest(verb string, endpoint string, params url.Values, response interface{}) error {
	// Preparing Request
	requestURL := strings.Join([]string{w.BaseURL, endpoint}, "")
	paramsReader := strings.NewReader(params.Encode())
	clientReq, err := http.NewRequest(verb, requestURL, paramsReader)
	if err != nil {
		return err
	}
	clientReq.Header.Add("Content-Type", "application/x-www-form-urlencoded")

	// Executing request
	clientRes, err := w.HttpClient.Do(clientReq)
	if err != nil {
		return err
	}
	defer clientRes.Body.Close()

	// Packing into response
	body, err := ioutil.ReadAll(clientRes.Body)
	if err != nil {
		return err
	}
	err = json.Unmarshal(body, &response)
	return err
}
<|endoftext|>


<commit_before>package stun

import (
	"errors"
	"fmt"
	"net"
	"strconv"
)

const (
	familyIPv4 uint16 = 0x01
	familyIPv6 uint16 = 0x02
)

// XORMappedAddress implements XOR-MAPPED-ADDRESS attribute.
//
// https://tools.ietf.org/html/rfc5389#section-15.2
type XORMappedAddress struct {
	IP   net.IP
	Port int
}

func (a XORMappedAddress) String() string {
	return net.JoinHostPort(a.IP.String(), strconv.Itoa(a.Port))
}

// isIPv4 returns true if ip with len of net.IPv6Len seems to be ipv4.
func isIPv4(ip net.IP) bool {
	// Optimized for performance. Copied from net.IP.To4.
	return isZeros(ip[0:10]) && ip[10] == 0xff && ip[11] == 0xff
}

// Is p all zeros?
func isZeros(p net.IP) bool {
	for i := 0; i < len(p); i++ {
		if p[i] != 0 {
			return false
		}
	}
	return true
}

// ErrBadIPLength means that len(IP) is not net.{IPv6len,IPv4len}.
var ErrBadIPLength = errors.New("invalid length of IP value")

// AddToAs adds XOR-MAPPED-ADDRESS value to m as t attribute.
func (a *XORMappedAddress) AddToAs(m *Message, t AttrType) error {
	var (
		family = familyIPv4
		ip     = a.IP
	)
	if len(a.IP) == net.IPv6len {
		if isIPv4(ip) {
			ip = ip[12:16] // like in ip.To4()
		} else {
			family = familyIPv6
		}
	} else if len(ip) != net.IPv4len {
		return ErrBadIPLength
	}
	value := make([]byte, 32+128)
	value[0] = 0 // first 8 bits are zeroes
	xorValue := make([]byte, net.IPv6len)
	copy(xorValue[4:], m.TransactionID[:])
	bin.PutUint32(xorValue[0:4], magicCookie)
	bin.PutUint16(value[0:2], family)
	bin.PutUint16(value[2:4], uint16(a.Port^magicCookie>>16))
	xorBytes(value[4:4+len(ip)], ip, xorValue)
	m.Add(t, value[:4+len(ip)])
	return nil
}

// AddTo adds XOR-MAPPED-ADDRESS to m. Can return ErrBadIPLength
// if len(a.IP) is invalid.
func (a *XORMappedAddress) AddTo(m *Message) error {
	return a.AddToAs(m, AttrXORMappedAddress)
}

// GetFrom decodes XOR-MAPPED-ADDRESS attribute value in message
// getting it as for t type.
func (a *XORMappedAddress) GetFromAs(m *Message, t AttrType) error {
	v, err := m.Get(t)
	if err != nil {
		return err
	}
	family := bin.Uint16(v[0:2])
	if family != familyIPv6 && family != familyIPv4 {
		return newDecodeErr("xor-mapped address", "family",
			fmt.Sprintf("bad value %d", family),
		)
	}
	ipLen := net.IPv4len
	if family == familyIPv6 {
		ipLen = net.IPv6len
	}
	// Ensuring len(a.IP) == ipLen and reusing a.IP.
	if len(a.IP) < ipLen {
		a.IP = a.IP[:cap(a.IP)]
		for len(a.IP) < ipLen {
			a.IP = append(a.IP, 0)
		}
	}
	a.IP = a.IP[:ipLen]
	for i := range a.IP {
		a.IP[i] = 0
	}
	a.Port = int(bin.Uint16(v[2:4])) ^ (magicCookie >> 16)
	xorValue := make([]byte, 4+transactionIDSize)
	bin.PutUint32(xorValue[0:4], magicCookie)
	copy(xorValue[4:], m.TransactionID[:])
	xorBytes(a.IP, v[4:], xorValue)
	return nil
}

// GetFrom decodes XOR-MAPPED-ADDRESS attribute in message and returns
// error if any. While decoding, a.IP is reused if possible and can be
// rendered to invalid state (e.g. if a.IP was set to IPv6 and then
// IPv4 value were decoded into it), be careful.
//
// Example:
//
//  expectedIP := net.ParseIP("213.141.156.236")
//  expectedIP.String() // 213.141.156.236, 16 bytes, first 12 of them are zeroes
//  expectedPort := 21254
//  addr := &XORMappedAddress{
//    IP:   expectedIP,
//    Port: expectedPort,
//  }
//  // addr were added to message that is decoded as newMessage
//  // ...
//
//  addr.GetFrom(newMessage)
//  addr.IP.String()    // 213.141.156.236, net.IPv4Len
//  expectedIP.String() // d58d:9cec::ffff:d58d:9cec, 16 bytes, first 4 are IPv4
//  // now we have len(expectedIP) = 16 and len(addr.IP) = 4.
func (a *XORMappedAddress) GetFrom(m *Message) error {
	return a.GetFromAs(m, AttrXORMappedAddress)
}
<commit_msg>a:xoraddr: use non-pointer reciever<commit_after>package stun

import (
	"errors"
	"fmt"
	"net"
	"strconv"
)

const (
	familyIPv4 uint16 = 0x01
	familyIPv6 uint16 = 0x02
)

// XORMappedAddress implements XOR-MAPPED-ADDRESS attribute.
//
// https://tools.ietf.org/html/rfc5389#section-15.2
type XORMappedAddress struct {
	IP   net.IP
	Port int
}

func (a XORMappedAddress) String() string {
	return net.JoinHostPort(a.IP.String(), strconv.Itoa(a.Port))
}

// isIPv4 returns true if ip with len of net.IPv6Len seems to be ipv4.
func isIPv4(ip net.IP) bool {
	// Optimized for performance. Copied from net.IP.To4.
	return isZeros(ip[0:10]) && ip[10] == 0xff && ip[11] == 0xff
}

// Is p all zeros?
func isZeros(p net.IP) bool {
	for i := 0; i < len(p); i++ {
		if p[i] != 0 {
			return false
		}
	}
	return true
}

// ErrBadIPLength means that len(IP) is not net.{IPv6len,IPv4len}.
var ErrBadIPLength = errors.New("invalid length of IP value")

// AddToAs adds XOR-MAPPED-ADDRESS value to m as t attribute.
func (a XORMappedAddress) AddToAs(m *Message, t AttrType) error {
	var (
		family = familyIPv4
		ip     = a.IP
	)
	if len(a.IP) == net.IPv6len {
		if isIPv4(ip) {
			ip = ip[12:16] // like in ip.To4()
		} else {
			family = familyIPv6
		}
	} else if len(ip) != net.IPv4len {
		return ErrBadIPLength
	}
	value := make([]byte, 32+128)
	value[0] = 0 // first 8 bits are zeroes
	xorValue := make([]byte, net.IPv6len)
	copy(xorValue[4:], m.TransactionID[:])
	bin.PutUint32(xorValue[0:4], magicCookie)
	bin.PutUint16(value[0:2], family)
	bin.PutUint16(value[2:4], uint16(a.Port^magicCookie>>16))
	xorBytes(value[4:4+len(ip)], ip, xorValue)
	m.Add(t, value[:4+len(ip)])
	return nil
}

// AddTo adds XOR-MAPPED-ADDRESS to m. Can return ErrBadIPLength
// if len(a.IP) is invalid.
func (a XORMappedAddress) AddTo(m *Message) error {
	return a.AddToAs(m, AttrXORMappedAddress)
}

// GetFrom decodes XOR-MAPPED-ADDRESS attribute value in message
// getting it as for t type.
func (a *XORMappedAddress) GetFromAs(m *Message, t AttrType) error {
	v, err := m.Get(t)
	if err != nil {
		return err
	}
	family := bin.Uint16(v[0:2])
	if family != familyIPv6 && family != familyIPv4 {
		return newDecodeErr("xor-mapped address", "family",
			fmt.Sprintf("bad value %d", family),
		)
	}
	ipLen := net.IPv4len
	if family == familyIPv6 {
		ipLen = net.IPv6len
	}
	// Ensuring len(a.IP) == ipLen and reusing a.IP.
	if len(a.IP) < ipLen {
		a.IP = a.IP[:cap(a.IP)]
		for len(a.IP) < ipLen {
			a.IP = append(a.IP, 0)
		}
	}
	a.IP = a.IP[:ipLen]
	for i := range a.IP {
		a.IP[i] = 0
	}
	a.Port = int(bin.Uint16(v[2:4])) ^ (magicCookie >> 16)
	xorValue := make([]byte, 4+transactionIDSize)
	bin.PutUint32(xorValue[0:4], magicCookie)
	copy(xorValue[4:], m.TransactionID[:])
	xorBytes(a.IP, v[4:], xorValue)
	return nil
}

// GetFrom decodes XOR-MAPPED-ADDRESS attribute in message and returns
// error if any. While decoding, a.IP is reused if possible and can be
// rendered to invalid state (e.g. if a.IP was set to IPv6 and then
// IPv4 value were decoded into it), be careful.
//
// Example:
//
//  expectedIP := net.ParseIP("213.141.156.236")
//  expectedIP.String() // 213.141.156.236, 16 bytes, first 12 of them are zeroes
//  expectedPort := 21254
//  addr := &XORMappedAddress{
//    IP:   expectedIP,
//    Port: expectedPort,
//  }
//  // addr were added to message that is decoded as newMessage
//  // ...
//
//  addr.GetFrom(newMessage)
//  addr.IP.String()    // 213.141.156.236, net.IPv4Len
//  expectedIP.String() // d58d:9cec::ffff:d58d:9cec, 16 bytes, first 4 are IPv4
//  // now we have len(expectedIP) = 16 and len(addr.IP) = 4.
func (a *XORMappedAddress) GetFrom(m *Message) error {
	return a.GetFromAs(m, AttrXORMappedAddress)
}
<|endoftext|>


<commit_before>package kubernetes

import (
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/Masterminds/semver"
	"github.com/rusenask/cron"

	"k8s.io/api/core/v1"

	"github.com/prometheus/client_golang/prometheus"

	"github.com/keel-hq/keel/approvals"
	"github.com/keel-hq/keel/extension/notification"
	"github.com/keel-hq/keel/internal/k8s"
	"github.com/keel-hq/keel/internal/policy"
	"github.com/keel-hq/keel/types"
	"github.com/keel-hq/keel/util/image"
	"github.com/keel-hq/keel/util/policies"

	log "github.com/sirupsen/logrus"
)

var kubernetesVersionedUpdatesCounter = prometheus.NewCounterVec(
	prometheus.CounterOpts{
		Name: "kubernetes_versioned_updates_total",
		Help: "How many versioned deployments were updated, partitioned by deployment name.",
	},
	[]string{"kubernetes"},
)

var kubernetesUnversionedUpdatesCounter = prometheus.NewCounterVec(
	prometheus.CounterOpts{
		Name: "kubernetes_unversioned_updates_total",
		Help: "How many unversioned deployments were updated, partitioned by deployment name.",
	},
	[]string{"kubernetes"},
)

func init() {
	prometheus.MustRegister(kubernetesVersionedUpdatesCounter)
	prometheus.MustRegister(kubernetesUnversionedUpdatesCounter)
}

// ProviderName - provider name
const ProviderName = "kubernetes"

var versionreg = regexp.MustCompile(`:[^:]*$`)

// GenericResourceCache an interface for generic resource cache.
type GenericResourceCache interface {
	// Values returns a copy of the contents of the cache.
	// The slice and its contents should be treated as read-only.
	Values() []*k8s.GenericResource

	// Register registers ch to receive a value when Notify is called.
	Register(chan int, int)
}

// UpdatePlan - deployment update plan
type UpdatePlan struct {
	// Updated deployment version
	// Deployment v1beta1.Deployment
	Resource *k8s.GenericResource

	// Current (last seen cluster version)
	CurrentVersion string
	// New version that's already in the deployment
	NewVersion string
}

func (p *UpdatePlan) String() string {
	if p.Resource != nil {
		return fmt.Sprintf("%s %s->%s", p.Resource.Identifier, p.CurrentVersion, p.NewVersion)
	}
	return "empty plan"
}

// Provider - kubernetes provider for auto update
type Provider struct {
	implementer Implementer

	sender notification.Sender

	approvalManager approvals.Manager

	cache GenericResourceCache

	events chan *types.Event
	stop   chan struct{}
}

// NewProvider - create new kubernetes based provider
func NewProvider(implementer Implementer, sender notification.Sender, approvalManager approvals.Manager, cache GenericResourceCache) (*Provider, error) {
	return &Provider{
		implementer:     implementer,
		cache:           cache,
		approvalManager: approvalManager,
		events:          make(chan *types.Event, 100),
		stop:            make(chan struct{}),
		sender:          sender,
	}, nil
}

// Submit - submit event to provider
func (p *Provider) Submit(event types.Event) error {
	p.events <- &event
	return nil
}

// GetName - get provider name
func (p *Provider) GetName() string {
	return ProviderName
}

// Start - starts kubernetes provider, waits for events
func (p *Provider) Start() error {
	return p.startInternal()
}

// Stop - stops kubernetes provider
func (p *Provider) Stop() {
	close(p.stop)
}

// TrackedImages returns a list of tracked images.
func (p *Provider) TrackedImages() ([]*types.TrackedImage, error) {
	var trackedImages []*types.TrackedImage

	for _, gr := range p.cache.Values() {
		labels := gr.GetLabels()
		annotations := gr.GetAnnotations()

		// ignoring unlabelled deployments
		plc := policy.GetPolicyFromLabelsOrAnnotations(labels, annotations)
		if plc.Type() == policy.PolicyTypeNone {
			continue
		}

		schedule, ok := annotations[types.KeelPollScheduleAnnotation]
		if ok {
			_, err := cron.Parse(schedule)
			if err != nil {
				log.WithFields(log.Fields{
					"error":     err,
					"schedule":  schedule,
					"name":      gr.Name,
					"namespace": gr.Namespace,
				}).Error("provider.kubernetes: failed to parse poll schedule, setting default schedule")
				schedule = types.KeelPollDefaultSchedule
			}
		} else {
			schedule = types.KeelPollDefaultSchedule
		}

		// trigger type, we only care for "poll" type triggers
		trigger := policies.GetTriggerPolicy(labels, annotations)
		secrets := gr.GetImagePullSecrets()
		images := gr.GetImages()
		for _, img := range images {
			ref, err := image.Parse(img)
			if err != nil {
				log.WithFields(log.Fields{
					"error":     err,
					"image":     img,
					"namespace": gr.Namespace,
					"name":      gr.Name,
				}).Error("provider.kubernetes: failed to parse image")
				continue
			}
			svp := make(map[string]string)

			semverTag, err := semver.NewVersion(ref.Tag())
			if err == nil {
				if semverTag.Prerelease() != "" {
					svp[semverTag.Prerelease()] = ref.Tag()
				}
			}

			trackedImages = append(trackedImages, &types.TrackedImage{
				Image:                ref,
				PollSchedule:         schedule,
				Trigger:              trigger,
				Provider:             ProviderName,
				Namespace:            gr.Namespace,
				Secrets:              secrets,
				Meta:                 make(map[string]string),
				SemverPreReleaseTags: svp,
			})
		}
	}

	return trackedImages, nil
}

func (p *Provider) startInternal() error {
	for {
		select {
		case event := <-p.events:
			log.WithFields(log.Fields{
				"repository": event.Repository.Name,
				"tag":        event.Repository.Tag,
				"registry":   event.Repository.Host,
			}).Info("provider.kubernetes: processing event")
			_, err := p.processEvent(event)
			if err != nil {
				log.WithFields(log.Fields{
					"error": err,
					"image": event.Repository.Name,
					"tag":   event.Repository.Tag,
				}).Error("provider.kubernetes: failed to process event")
			}
		case <-p.stop:
			log.Info("provider.kubernetes: got shutdown signal, stopping...")
			return nil
		}
	}
}

func (p *Provider) processEvent(event *types.Event) (updated []*k8s.GenericResource, err error) {
	plans, err := p.createUpdatePlans(&event.Repository)
	if err != nil {
		return nil, err
	}

	if len(plans) == 0 {
		log.WithFields(log.Fields{
			"image": event.Repository.Name,
			"tag":   event.Repository.Tag,
		}).Debug("provider.kubernetes: no plans for deployment updates found for this event")
		return
	}

	approvedPlans := p.checkForApprovals(event, plans)

	return p.updateDeployments(approvedPlans)
}

func (p *Provider) updateDeployments(plans []*UpdatePlan) (updated []*k8s.GenericResource, err error) {
	for _, plan := range plans {
		resource := plan.Resource

		annotations := resource.GetAnnotations()

		notificationChannels := types.ParseEventNotificationChannels(annotations)

		p.sender.Send(types.EventNotification{
			Name:      "preparing to update resource",
			Message:   fmt.Sprintf("Preparing to update %s %s/%s %s->%s (%s)", resource.Kind(), resource.Namespace, resource.Name, plan.CurrentVersion, plan.NewVersion, strings.Join(resource.GetImages(), ", ")),
			CreatedAt: time.Now(),
			Type:      types.NotificationPreDeploymentUpdate,
			Level:     types.LevelDebug,
			Channels:  notificationChannels,
		})

		var err error

		timestamp := time.Now().Format(time.RFC3339)
		annotations["kubernetes.io/change-cause"] = fmt.Sprintf("keel automated update, version %s -> %s [%s]", plan.CurrentVersion, plan.NewVersion, timestamp)

		resource.SetAnnotations(annotations)

		err = p.implementer.Update(resource)
		kubernetesVersionedUpdatesCounter.With(prometheus.Labels{"kubernetes": fmt.Sprintf("%s/%s", resource.Namespace, resource.Name)}).Inc()
		if err != nil {
			log.WithFields(log.Fields{
				"error":      err,
				"namespace":  resource.Namespace,
				"deployment": resource.Name,
				"kind":       resource.Kind(),
				"update":     fmt.Sprintf("%s->%s", plan.CurrentVersion, plan.NewVersion),
			}).Error("provider.kubernetes: got error while updating resource")

			p.sender.Send(types.EventNotification{
				Name:      "update resource",
				Message:   fmt.Sprintf("%s %s/%s update %s->%s failed, error: %s", resource.Kind(), resource.Namespace, resource.Name, plan.CurrentVersion, plan.NewVersion, err),
				CreatedAt: time.Now(),
				Type:      types.NotificationDeploymentUpdate,
				Level:     types.LevelError,
				Channels:  notificationChannels,
			})

			continue
		}

		err = p.updateComplete(plan)
		if err != nil {
			log.WithFields(log.Fields{
				"error":     err,
				"name":      resource.Name,
				"kind":      resource.Kind(),
				"namespace": resource.Namespace,
			}).Warn("provider.kubernetes: got error while resetting approvals counter after successful update")
		}

		var msg string
		releaseNotes := types.ParseReleaseNotesURL(resource.GetAnnotations())
		if releaseNotes != "" {
			msg = fmt.Sprintf("Successfully updated %s %s/%s %s->%s (%s). Release notes: %s", resource.Kind(), resource.Namespace, resource.Name, plan.CurrentVersion, plan.NewVersion, strings.Join(resource.GetImages(), ", "), releaseNotes)
		} else {
			msg = fmt.Sprintf("Successfully updated %s %s/%s %s->%s (%s)", resource.Kind(), resource.Namespace, resource.Name, plan.CurrentVersion, plan.NewVersion, strings.Join(resource.GetImages(), ", "))
		}

		p.sender.Send(types.EventNotification{
			Name:      "update resource",
			Message:   msg,
			CreatedAt: time.Now(),
			Type:      types.NotificationDeploymentUpdate,
			Level:     types.LevelSuccess,
			Channels:  notificationChannels,
		})

		log.WithFields(log.Fields{
			"name":      resource.Name,
			"kind":      resource.Kind(),
			"previous":  plan.CurrentVersion,
			"new":       plan.NewVersion,
			"namespace": resource.Namespace,
		}).Info("provider.kubernetes: resource updated")
		updated = append(updated, resource)
	}

	return
}

func getDesiredImage(delta map[string]string, currentImage string) (string, error) {
	currentRef, err := image.Parse(currentImage)
	if err != nil {
		return "", err
	}
	for repository, tag := range delta {
		if repository == currentRef.Repository() {
			ref, err := image.Parse(repository)
			if err != nil {
				return "", err
			}

			// updating image
			if ref.Registry() == image.DefaultRegistryHostname {
				return fmt.Sprintf("%s:%s", ref.ShortName(), tag), nil
			}
			return fmt.Sprintf("%s:%s", ref.Repository(), tag), nil
		}
	}
	return "", fmt.Errorf("image %s not found in deltas", currentImage)
}

// createUpdatePlans - impacted deployments by changed repository
func (p *Provider) createUpdatePlans(repo *types.Repository) ([]*UpdatePlan, error) {
	impacted := []*UpdatePlan{}

	for _, resource := range p.cache.Values() {

		labels := resource.GetLabels()
		annotations := resource.GetAnnotations()

		plc := policy.GetPolicyFromLabelsOrAnnotations(labels, annotations)
		if plc.Type() == policy.PolicyTypeNone {
			log.Debugf("no policy defined, skipping: %s, labels: %s, annotations: %s", resource.Identifier, labels, annotations)
			continue
		}

		updated, shouldUpdateDeployment, err := checkForUpdate(plc, repo, resource)
		if err != nil {
			log.WithFields(log.Fields{
				"error":      err,
				"deployment": resource.Name,
				"kind":       resource.Kind(),
				"namespace":  resource.Namespace,
			}).Error("provider.kubernetes: got error while checking versioned resource")
			continue
		}

		if shouldUpdateDeployment {
			impacted = append(impacted, updated)
		}
	}

	return impacted, nil
}

func (p *Provider) namespaces() (*v1.NamespaceList, error) {
	return p.implementer.Namespaces()
}
<commit_msg>setting policies in the tracked images<commit_after>package kubernetes

import (
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/Masterminds/semver"
	"github.com/rusenask/cron"

	"k8s.io/api/core/v1"

	"github.com/prometheus/client_golang/prometheus"

	"github.com/keel-hq/keel/approvals"
	"github.com/keel-hq/keel/extension/notification"
	"github.com/keel-hq/keel/internal/k8s"
	"github.com/keel-hq/keel/internal/policy"
	"github.com/keel-hq/keel/types"
	"github.com/keel-hq/keel/util/image"
	"github.com/keel-hq/keel/util/policies"

	log "github.com/sirupsen/logrus"
)

var kubernetesVersionedUpdatesCounter = prometheus.NewCounterVec(
	prometheus.CounterOpts{
		Name: "kubernetes_versioned_updates_total",
		Help: "How many versioned deployments were updated, partitioned by deployment name.",
	},
	[]string{"kubernetes"},
)

var kubernetesUnversionedUpdatesCounter = prometheus.NewCounterVec(
	prometheus.CounterOpts{
		Name: "kubernetes_unversioned_updates_total",
		Help: "How many unversioned deployments were updated, partitioned by deployment name.",
	},
	[]string{"kubernetes"},
)

func init() {
	prometheus.MustRegister(kubernetesVersionedUpdatesCounter)
	prometheus.MustRegister(kubernetesUnversionedUpdatesCounter)
}

// ProviderName - provider name
const ProviderName = "kubernetes"

var versionreg = regexp.MustCompile(`:[^:]*$`)

// GenericResourceCache an interface for generic resource cache.
type GenericResourceCache interface {
	// Values returns a copy of the contents of the cache.
	// The slice and its contents should be treated as read-only.
	Values() []*k8s.GenericResource

	// Register registers ch to receive a value when Notify is called.
	Register(chan int, int)
}

// UpdatePlan - deployment update plan
type UpdatePlan struct {
	// Updated deployment version
	// Deployment v1beta1.Deployment
	Resource *k8s.GenericResource

	// Current (last seen cluster version)
	CurrentVersion string
	// New version that's already in the deployment
	NewVersion string
}

func (p *UpdatePlan) String() string {
	if p.Resource != nil {
		return fmt.Sprintf("%s %s->%s", p.Resource.Identifier, p.CurrentVersion, p.NewVersion)
	}
	return "empty plan"
}

// Provider - kubernetes provider for auto update
type Provider struct {
	implementer Implementer

	sender notification.Sender

	approvalManager approvals.Manager

	cache GenericResourceCache

	events chan *types.Event
	stop   chan struct{}
}

// NewProvider - create new kubernetes based provider
func NewProvider(implementer Implementer, sender notification.Sender, approvalManager approvals.Manager, cache GenericResourceCache) (*Provider, error) {
	return &Provider{
		implementer:     implementer,
		cache:           cache,
		approvalManager: approvalManager,
		events:          make(chan *types.Event, 100),
		stop:            make(chan struct{}),
		sender:          sender,
	}, nil
}

// Submit - submit event to provider
func (p *Provider) Submit(event types.Event) error {
	p.events <- &event
	return nil
}

// GetName - get provider name
func (p *Provider) GetName() string {
	return ProviderName
}

// Start - starts kubernetes provider, waits for events
func (p *Provider) Start() error {
	return p.startInternal()
}

// Stop - stops kubernetes provider
func (p *Provider) Stop() {
	close(p.stop)
}

// TrackedImages returns a list of tracked images.
func (p *Provider) TrackedImages() ([]*types.TrackedImage, error) {
	var trackedImages []*types.TrackedImage

	for _, gr := range p.cache.Values() {
		labels := gr.GetLabels()
		annotations := gr.GetAnnotations()

		// ignoring unlabelled deployments
		plc := policy.GetPolicyFromLabelsOrAnnotations(labels, annotations)
		if plc.Type() == policy.PolicyTypeNone {
			continue
		}

		schedule, ok := annotations[types.KeelPollScheduleAnnotation]
		if ok {
			_, err := cron.Parse(schedule)
			if err != nil {
				log.WithFields(log.Fields{
					"error":     err,
					"schedule":  schedule,
					"name":      gr.Name,
					"namespace": gr.Namespace,
				}).Error("provider.kubernetes: failed to parse poll schedule, setting default schedule")
				schedule = types.KeelPollDefaultSchedule
			}
		} else {
			schedule = types.KeelPollDefaultSchedule
		}

		// trigger type, we only care for "poll" type triggers
		trigger := policies.GetTriggerPolicy(labels, annotations)
		secrets := gr.GetImagePullSecrets()
		images := gr.GetImages()
		for _, img := range images {
			ref, err := image.Parse(img)
			if err != nil {
				log.WithFields(log.Fields{
					"error":     err,
					"image":     img,
					"namespace": gr.Namespace,
					"name":      gr.Name,
				}).Error("provider.kubernetes: failed to parse image")
				continue
			}
			svp := make(map[string]string)

			semverTag, err := semver.NewVersion(ref.Tag())
			if err == nil {
				if semverTag.Prerelease() != "" {
					svp[semverTag.Prerelease()] = ref.Tag()
				}
			}

			trackedImages = append(trackedImages, &types.TrackedImage{
				Image:        ref,
				PollSchedule: schedule,
				Trigger:      trigger,
				Provider:     ProviderName,
				Namespace:    gr.Namespace,
				Secrets:      secrets,
				Meta:         make(map[string]string),
				Policy:       plc,
			})
		}
	}

	return trackedImages, nil
}

func (p *Provider) startInternal() error {
	for {
		select {
		case event := <-p.events:
			log.WithFields(log.Fields{
				"repository": event.Repository.Name,
				"tag":        event.Repository.Tag,
				"registry":   event.Repository.Host,
			}).Info("provider.kubernetes: processing event")
			_, err := p.processEvent(event)
			if err != nil {
				log.WithFields(log.Fields{
					"error": err,
					"image": event.Repository.Name,
					"tag":   event.Repository.Tag,
				}).Error("provider.kubernetes: failed to process event")
			}
		case <-p.stop:
			log.Info("provider.kubernetes: got shutdown signal, stopping...")
			return nil
		}
	}
}

func (p *Provider) processEvent(event *types.Event) (updated []*k8s.GenericResource, err error) {
	plans, err := p.createUpdatePlans(&event.Repository)
	if err != nil {
		return nil, err
	}

	if len(plans) == 0 {
		log.WithFields(log.Fields{
			"image": event.Repository.Name,
			"tag":   event.Repository.Tag,
		}).Debug("provider.kubernetes: no plans for deployment updates found for this event")
		return
	}

	approvedPlans := p.checkForApprovals(event, plans)

	return p.updateDeployments(approvedPlans)
}

func (p *Provider) updateDeployments(plans []*UpdatePlan) (updated []*k8s.GenericResource, err error) {
	for _, plan := range plans {
		resource := plan.Resource

		annotations := resource.GetAnnotations()

		notificationChannels := types.ParseEventNotificationChannels(annotations)

		p.sender.Send(types.EventNotification{
			Name:      "preparing to update resource",
			Message:   fmt.Sprintf("Preparing to update %s %s/%s %s->%s (%s)", resource.Kind(), resource.Namespace, resource.Name, plan.CurrentVersion, plan.NewVersion, strings.Join(resource.GetImages(), ", ")),
			CreatedAt: time.Now(),
			Type:      types.NotificationPreDeploymentUpdate,
			Level:     types.LevelDebug,
			Channels:  notificationChannels,
		})

		var err error

		timestamp := time.Now().Format(time.RFC3339)
		annotations["kubernetes.io/change-cause"] = fmt.Sprintf("keel automated update, version %s -> %s [%s]", plan.CurrentVersion, plan.NewVersion, timestamp)

		resource.SetAnnotations(annotations)

		err = p.implementer.Update(resource)
		kubernetesVersionedUpdatesCounter.With(prometheus.Labels{"kubernetes": fmt.Sprintf("%s/%s", resource.Namespace, resource.Name)}).Inc()
		if err != nil {
			log.WithFields(log.Fields{
				"error":      err,
				"namespace":  resource.Namespace,
				"deployment": resource.Name,
				"kind":       resource.Kind(),
				"update":     fmt.Sprintf("%s->%s", plan.CurrentVersion, plan.NewVersion),
			}).Error("provider.kubernetes: got error while updating resource")

			p.sender.Send(types.EventNotification{
				Name:      "update resource",
				Message:   fmt.Sprintf("%s %s/%s update %s->%s failed, error: %s", resource.Kind(), resource.Namespace, resource.Name, plan.CurrentVersion, plan.NewVersion, err),
				CreatedAt: time.Now(),
				Type:      types.NotificationDeploymentUpdate,
				Level:     types.LevelError,
				Channels:  notificationChannels,
			})

			continue
		}

		err = p.updateComplete(plan)
		if err != nil {
			log.WithFields(log.Fields{
				"error":     err,
				"name":      resource.Name,
				"kind":      resource.Kind(),
				"namespace": resource.Namespace,
			}).Warn("provider.kubernetes: got error while resetting approvals counter after successful update")
		}

		var msg string
		releaseNotes := types.ParseReleaseNotesURL(resource.GetAnnotations())
		if releaseNotes != "" {
			msg = fmt.Sprintf("Successfully updated %s %s/%s %s->%s (%s). Release notes: %s", resource.Kind(), resource.Namespace, resource.Name, plan.CurrentVersion, plan.NewVersion, strings.Join(resource.GetImages(), ", "), releaseNotes)
		} else {
			msg = fmt.Sprintf("Successfully updated %s %s/%s %s->%s (%s)", resource.Kind(), resource.Namespace, resource.Name, plan.CurrentVersion, plan.NewVersion, strings.Join(resource.GetImages(), ", "))
		}

		p.sender.Send(types.EventNotification{
			Name:      "update resource",
			Message:   msg,
			CreatedAt: time.Now(),
			Type:      types.NotificationDeploymentUpdate,
			Level:     types.LevelSuccess,
			Channels:  notificationChannels,
		})

		log.WithFields(log.Fields{
			"name":      resource.Name,
			"kind":      resource.Kind(),
			"previous":  plan.CurrentVersion,
			"new":       plan.NewVersion,
			"namespace": resource.Namespace,
		}).Info("provider.kubernetes: resource updated")
		updated = append(updated, resource)
	}

	return
}

func getDesiredImage(delta map[string]string, currentImage string) (string, error) {
	currentRef, err := image.Parse(currentImage)
	if err != nil {
		return "", err
	}
	for repository, tag := range delta {
		if repository == currentRef.Repository() {
			ref, err := image.Parse(repository)
			if err != nil {
				return "", err
			}

			// updating image
			if ref.Registry() == image.DefaultRegistryHostname {
				return fmt.Sprintf("%s:%s", ref.ShortName(), tag), nil
			}
			return fmt.Sprintf("%s:%s", ref.Repository(), tag), nil
		}
	}
	return "", fmt.Errorf("image %s not found in deltas", currentImage)
}

// createUpdatePlans - impacted deployments by changed repository
func (p *Provider) createUpdatePlans(repo *types.Repository) ([]*UpdatePlan, error) {
	impacted := []*UpdatePlan{}

	for _, resource := range p.cache.Values() {

		labels := resource.GetLabels()
		annotations := resource.GetAnnotations()

		plc := policy.GetPolicyFromLabelsOrAnnotations(labels, annotations)
		if plc.Type() == policy.PolicyTypeNone {
			log.Debugf("no policy defined, skipping: %s, labels: %s, annotations: %s", resource.Identifier, labels, annotations)
			continue
		}

		updated, shouldUpdateDeployment, err := checkForUpdate(plc, repo, resource)
		if err != nil {
			log.WithFields(log.Fields{
				"error":      err,
				"deployment": resource.Name,
				"kind":       resource.Kind(),
				"namespace":  resource.Namespace,
			}).Error("provider.kubernetes: got error while checking versioned resource")
			continue
		}

		if shouldUpdateDeployment {
			impacted = append(impacted, updated)
		}
	}

	return impacted, nil
}

func (p *Provider) namespaces() (*v1.NamespaceList, error) {
	return p.implementer.Namespaces()
}
<|endoftext|>


<commit_before>package twitter

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"

	"github.com/gorilla/pat"
	"github.com/markbates/goth"
	"github.com/mrjones/oauth"
	"github.com/stretchr/testify/assert"
)

func Test_New(t *testing.T) {
	t.Parallel()
	a := assert.New(t)

	provider := twitterProvider()
	a.Equal(provider.ClientKey, os.Getenv("TWITTER_KEY"))
	a.Equal(provider.Secret, os.Getenv("TWITTER_SECRET"))
	a.Equal(provider.CallbackURL, "/foo")
}

func Test_Implements_Provider(t *testing.T) {
	t.Parallel()
	a := assert.New(t)

	a.Implements((*goth.Provider)(nil), twitterProvider())
}

func Test_BeginAuth(t *testing.T) {
	t.Parallel()
	a := assert.New(t)

	mockTwitter(func(ts *httptest.Server) {
		provider := twitterProvider()
		session, err := provider.BeginAuth("state")
		s := session.(*Session)
		a.NoError(err)
		a.Contains(s.AuthURL, "authorize?oauth_token=TOKEN")
		a.Equal("TOKEN", s.RequestToken.Token)
		a.Equal("SECRET", s.RequestToken.Secret)
	})
	mockTwitter(func(ts *httptest.Server) {
		provider := twitterProviderAuthenticate()
		session, err := provider.BeginAuth("state")
		s := session.(*Session)
		a.NoError(err)
		a.Contains(s.AuthURL, "authenticate?oauth_token=TOKEN")
		a.Equal("TOKEN", s.RequestToken.Token)
		a.Equal("SECRET", s.RequestToken.Secret)
	})
}

func Test_FetchUser(t *testing.T) {
	t.Parallel()
	a := assert.New(t)

	mockTwitter(func(ts *httptest.Server) {
		provider := twitterProvider()
		session := Session{AccessToken: &oauth.AccessToken{Token: "TOKEN", Secret: "SECRET"}}

		user, err := provider.FetchUser(&session)
		a.NoError(err)

		a.Equal("Homer", user.Name)
		a.Equal("duffman", user.NickName)
		a.Equal("Duff rules!!", user.Description)
		a.Equal("http://example.com/image.jpg", user.AvatarURL)
		a.Equal("1234", user.UserID)
		a.Equal("Springfield", user.Location)
		a.Equal("TOKEN", user.AccessToken)
		a.Equal("duffman@springfield.com", user.Email)
	})
}

func Test_SessionFromJSON(t *testing.T) {
	t.Parallel()
	a := assert.New(t)

	provider := twitterProvider()

	s, err := provider.UnmarshalSession(`{"AuthURL":"http://com/auth_url","AccessToken":{"Token":"1234567890","Secret":"secret!!","AdditionalData":{}},"RequestToken":{"Token":"0987654321","Secret":"!!secret"}}`)
	a.NoError(err)
	session := s.(*Session)
	a.Equal(session.AuthURL, "http://com/auth_url")
	a.Equal(session.AccessToken.Token, "1234567890")
	a.Equal(session.AccessToken.Secret, "secret!!")
	a.Equal(session.RequestToken.Token, "0987654321")
	a.Equal(session.RequestToken.Secret, "!!secret")
}

func twitterProvider() *Provider {
	return New(os.Getenv("TWITTER_KEY"), os.Getenv("TWITTER_SECRET"), "/foo")
}

func twitterProviderAuthenticate() *Provider {
	return NewAuthenticate(os.Getenv("TWITTER_KEY"), os.Getenv("TWITTER_SECRET"), "/foo")
}

func mockTwitter(f func(*httptest.Server)) {
	p := pat.New()
	p.Get("/oauth/request_token", func(res http.ResponseWriter, req *http.Request) {
		fmt.Fprint(res, "oauth_token=TOKEN&oauth_token_secret=SECRET")
	})
	p.Get("/1.1/account/verify_credentials.json", func(res http.ResponseWriter, req *http.Request) {
		data := map[string]string{
			"name":              "Homer",
			"screen_name":       "duffman",
			"description":       "Duff rules!!",
			"profile_image_url": "http://example.com/image.jpg",
			"id_str":            "1234",
			"location":          "Springfield",
			"email":			 "duffman@springfield.com",
		}
		json.NewEncoder(res).Encode(&data)
	})
	ts := httptest.NewServer(p)
	defer ts.Close()

	originalRequestURL := requestURL
	originalEndpointProfile := endpointProfile

	requestURL = ts.URL + "/oauth/request_token"
	endpointProfile = ts.URL + "/1.1/account/verify_credentials.json"

	f(ts)

	requestURL = originalRequestURL
	endpointProfile = originalEndpointProfile
}
<commit_msg>Ran go fmt to correct format.<commit_after>package twitter

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"

	"github.com/gorilla/pat"
	"github.com/markbates/goth"
	"github.com/mrjones/oauth"
	"github.com/stretchr/testify/assert"
)

func Test_New(t *testing.T) {
	t.Parallel()
	a := assert.New(t)

	provider := twitterProvider()
	a.Equal(provider.ClientKey, os.Getenv("TWITTER_KEY"))
	a.Equal(provider.Secret, os.Getenv("TWITTER_SECRET"))
	a.Equal(provider.CallbackURL, "/foo")
}

func Test_Implements_Provider(t *testing.T) {
	t.Parallel()
	a := assert.New(t)

	a.Implements((*goth.Provider)(nil), twitterProvider())
}

func Test_BeginAuth(t *testing.T) {
	t.Parallel()
	a := assert.New(t)

	mockTwitter(func(ts *httptest.Server) {
		provider := twitterProvider()
		session, err := provider.BeginAuth("state")
		s := session.(*Session)
		a.NoError(err)
		a.Contains(s.AuthURL, "authorize?oauth_token=TOKEN")
		a.Equal("TOKEN", s.RequestToken.Token)
		a.Equal("SECRET", s.RequestToken.Secret)
	})
	mockTwitter(func(ts *httptest.Server) {
		provider := twitterProviderAuthenticate()
		session, err := provider.BeginAuth("state")
		s := session.(*Session)
		a.NoError(err)
		a.Contains(s.AuthURL, "authenticate?oauth_token=TOKEN")
		a.Equal("TOKEN", s.RequestToken.Token)
		a.Equal("SECRET", s.RequestToken.Secret)
	})
}

func Test_FetchUser(t *testing.T) {
	t.Parallel()
	a := assert.New(t)

	mockTwitter(func(ts *httptest.Server) {
		provider := twitterProvider()
		session := Session{AccessToken: &oauth.AccessToken{Token: "TOKEN", Secret: "SECRET"}}

		user, err := provider.FetchUser(&session)
		a.NoError(err)

		a.Equal("Homer", user.Name)
		a.Equal("duffman", user.NickName)
		a.Equal("Duff rules!!", user.Description)
		a.Equal("http://example.com/image.jpg", user.AvatarURL)
		a.Equal("1234", user.UserID)
		a.Equal("Springfield", user.Location)
		a.Equal("TOKEN", user.AccessToken)
		a.Equal("duffman@springfield.com", user.Email)
	})
}

func Test_SessionFromJSON(t *testing.T) {
	t.Parallel()
	a := assert.New(t)

	provider := twitterProvider()

	s, err := provider.UnmarshalSession(`{"AuthURL":"http://com/auth_url","AccessToken":{"Token":"1234567890","Secret":"secret!!","AdditionalData":{}},"RequestToken":{"Token":"0987654321","Secret":"!!secret"}}`)
	a.NoError(err)
	session := s.(*Session)
	a.Equal(session.AuthURL, "http://com/auth_url")
	a.Equal(session.AccessToken.Token, "1234567890")
	a.Equal(session.AccessToken.Secret, "secret!!")
	a.Equal(session.RequestToken.Token, "0987654321")
	a.Equal(session.RequestToken.Secret, "!!secret")
}

func twitterProvider() *Provider {
	return New(os.Getenv("TWITTER_KEY"), os.Getenv("TWITTER_SECRET"), "/foo")
}

func twitterProviderAuthenticate() *Provider {
	return NewAuthenticate(os.Getenv("TWITTER_KEY"), os.Getenv("TWITTER_SECRET"), "/foo")
}

func mockTwitter(f func(*httptest.Server)) {
	p := pat.New()
	p.Get("/oauth/request_token", func(res http.ResponseWriter, req *http.Request) {
		fmt.Fprint(res, "oauth_token=TOKEN&oauth_token_secret=SECRET")
	})
	p.Get("/1.1/account/verify_credentials.json", func(res http.ResponseWriter, req *http.Request) {
		data := map[string]string{
			"name":              "Homer",
			"screen_name":       "duffman",
			"description":       "Duff rules!!",
			"profile_image_url": "http://example.com/image.jpg",
			"id_str":            "1234",
			"location":          "Springfield",
			"email":             "duffman@springfield.com",
		}
		json.NewEncoder(res).Encode(&data)
	})
	ts := httptest.NewServer(p)
	defer ts.Close()

	originalRequestURL := requestURL
	originalEndpointProfile := endpointProfile

	requestURL = ts.URL + "/oauth/request_token"
	endpointProfile = ts.URL + "/1.1/account/verify_credentials.json"

	f(ts)

	requestURL = originalRequestURL
	endpointProfile = originalEndpointProfile
}
<|endoftext|>


<commit_before>package vptree

import (
	"container/heap"
	"math"
	"math/rand"
)

type node struct {
	Item      interface{}
	Threshold float64
	Left      *node
	Right     *node
}

type heapItem struct {
	Item interface{}
	Dist float64
}

// A Metric is a function that measures the distance between two provided
// interface{}-values. The function *must* be a metric in the mathematical
// sense, that is, the metric d must fullfill the following requirements:
//
//	* d(x, y) >= 0
//	* d(x, y) = 0 if and only if x = y
//	* d(x, y) = d(y, x)
//	* d(x, z) <= d(x, y) + d(y, z) (triangle inequality)
type Metric func(a, b interface{}) float64

// A VPTree struct represents a Vantage-point tree. Vantage-point trees are
// useful for nearest-neighbour searches in high-dimensional metric spaces.
type VPTree struct {
	root           *node
	distanceMetric Metric
}

// New creates a new VP-tree using the metric and items provided. The metric
// measures the distance between two items, so that the VP-tree can find the
// nearest neighbour(s) of a target item.
func New(metric Metric, items []interface{}) (t *VPTree) {
	t = &VPTree{
		distanceMetric: metric,
	}
	t.root = t.buildFromPoints(items)
	return
}

// Search searches the VP-tree for the k nearest neighbours of target. It
// returns the up to k narest neighbours and the corresponding distances in
// order of least distance to largest distance.
func (vp *VPTree) Search(target interface{}, k int) (results []interface{}, distances []float64) {
	if k < 1 {
		return
	}

	h := make(priorityQueue, 0, k)

	tau := math.MaxFloat64
	vp.search(vp.root, &tau, target, k, &h)

	for h.Len() > 0 {
		hi := heap.Pop(&h)
		results = append(results, hi.(*heapItem).Item)
		distances = append(distances, hi.(*heapItem).Dist)
	}

	// Reverse results and distances, because we popped them from the heap
	// in large-to-small order
	for i, j := 0, len(results)-1; i < j; i, j = i+1, j-1 {
		results[i], results[j] = results[j], results[i]
		distances[i], distances[j] = distances[j], distances[i]
	}

	return
}

func (vp *VPTree) buildFromPoints(items []interface{}) (n *node) {
	if len(items) == 0 {
		return nil
	}

	n = &node{}

	// Take a random item out of the items slice and make it this node's item
	idx := rand.Intn(len(items))
	n.Item = items[idx]
	items[idx], items = items[len(items)-1], items[:len(items)-1]

	if len(items) > 0 {
		// Now partition the items into two equal-sized sets, one
		// closer to the node's item than the median, and one farther
		// away.
		median := len(items) / 2
		pivotDist := vp.distanceMetric(items[median], n.Item)
		items[median], items[len(items)-1] = items[len(items)-1], items[median]

		storeIndex := 0
		for i := 0; i < len(items)-1; i++ {
			if vp.distanceMetric(items[i], n.Item) <= pivotDist {
				items[storeIndex], items[i] = items[i], items[storeIndex]
				storeIndex++
			}
		}
		items[len(items)-1], items[storeIndex] = items[storeIndex], items[len(items)-1]
		median = storeIndex

		n.Threshold = vp.distanceMetric(items[median], n.Item)
		n.Left = vp.buildFromPoints(items[:median])
		n.Right = vp.buildFromPoints(items[median:])
	}
	return
}

func (vp *VPTree) search(n *node, tau *float64, target interface{}, k int, h *priorityQueue) {
	if n == nil {
		return
	}

	dist := vp.distanceMetric(n.Item, target)

	if dist < *tau {
		if h.Len() == k {
			heap.Pop(h)
		}
		heap.Push(h, &heapItem{n.Item, dist})
		if h.Len() == k {
			*tau = h.Top().(*heapItem).Dist
		}
	}

	if n.Left == nil && n.Right == nil {
		return
	}

	if dist < n.Threshold {
		if dist-*tau <= n.Threshold {
			vp.search(n.Left, tau, target, k, h)
		}

		if dist+*tau >= n.Threshold {
			vp.search(n.Right, tau, target, k, h)
		}
	} else {
		if dist+*tau >= n.Threshold {
			vp.search(n.Right, tau, target, k, h)
		}

		if dist-*tau <= n.Threshold {
			vp.search(n.Left, tau, target, k, h)
		}
	}
}
<commit_msg>Fix duplicate distanceMetric computation (thanks @vseledkin!)<commit_after>package vptree

import (
	"container/heap"
	"math"
	"math/rand"
)

type node struct {
	Item      interface{}
	Threshold float64
	Left      *node
	Right     *node
}

type heapItem struct {
	Item interface{}
	Dist float64
}

// A Metric is a function that measures the distance between two provided
// interface{}-values. The function *must* be a metric in the mathematical
// sense, that is, the metric d must fullfill the following requirements:
//
//	* d(x, y) >= 0
//	* d(x, y) = 0 if and only if x = y
//	* d(x, y) = d(y, x)
//	* d(x, z) <= d(x, y) + d(y, z) (triangle inequality)
type Metric func(a, b interface{}) float64

// A VPTree struct represents a Vantage-point tree. Vantage-point trees are
// useful for nearest-neighbour searches in high-dimensional metric spaces.
type VPTree struct {
	root           *node
	distanceMetric Metric
}

// New creates a new VP-tree using the metric and items provided. The metric
// measures the distance between two items, so that the VP-tree can find the
// nearest neighbour(s) of a target item.
func New(metric Metric, items []interface{}) (t *VPTree) {
	t = &VPTree{
		distanceMetric: metric,
	}
	t.root = t.buildFromPoints(items)
	return
}

// Search searches the VP-tree for the k nearest neighbours of target. It
// returns the up to k narest neighbours and the corresponding distances in
// order of least distance to largest distance.
func (vp *VPTree) Search(target interface{}, k int) (results []interface{}, distances []float64) {
	if k < 1 {
		return
	}

	h := make(priorityQueue, 0, k)

	tau := math.MaxFloat64
	vp.search(vp.root, &tau, target, k, &h)

	for h.Len() > 0 {
		hi := heap.Pop(&h)
		results = append(results, hi.(*heapItem).Item)
		distances = append(distances, hi.(*heapItem).Dist)
	}

	// Reverse results and distances, because we popped them from the heap
	// in large-to-small order
	for i, j := 0, len(results)-1; i < j; i, j = i+1, j-1 {
		results[i], results[j] = results[j], results[i]
		distances[i], distances[j] = distances[j], distances[i]
	}

	return
}

func (vp *VPTree) buildFromPoints(items []interface{}) (n *node) {
	if len(items) == 0 {
		return nil
	}

	n = &node{}

	// Take a random item out of the items slice and make it this node's item
	idx := rand.Intn(len(items))
	n.Item = items[idx]
	items[idx], items = items[len(items)-1], items[:len(items)-1]

	if len(items) > 0 {
		// Now partition the items into two equal-sized sets, one
		// closer to the node's item than the median, and one farther
		// away.
		median := len(items) / 2
		pivotDist := vp.distanceMetric(items[median], n.Item)
		items[median], items[len(items)-1] = items[len(items)-1], items[median]

		storeIndex := 0
		for i := 0; i < len(items)-1; i++ {
			if vp.distanceMetric(items[i], n.Item) <= pivotDist {
				items[storeIndex], items[i] = items[i], items[storeIndex]
				storeIndex++
			}
		}
		items[len(items)-1], items[storeIndex] = items[storeIndex], items[len(items)-1]
		median = storeIndex

		n.Threshold = pivotDist
		n.Left = vp.buildFromPoints(items[:median])
		n.Right = vp.buildFromPoints(items[median:])
	}
	return
}

func (vp *VPTree) search(n *node, tau *float64, target interface{}, k int, h *priorityQueue) {
	if n == nil {
		return
	}

	dist := vp.distanceMetric(n.Item, target)

	if dist < *tau {
		if h.Len() == k {
			heap.Pop(h)
		}
		heap.Push(h, &heapItem{n.Item, dist})
		if h.Len() == k {
			*tau = h.Top().(*heapItem).Dist
		}
	}

	if n.Left == nil && n.Right == nil {
		return
	}

	if dist < n.Threshold {
		if dist-*tau <= n.Threshold {
			vp.search(n.Left, tau, target, k, h)
		}

		if dist+*tau >= n.Threshold {
			vp.search(n.Right, tau, target, k, h)
		}
	} else {
		if dist+*tau >= n.Threshold {
			vp.search(n.Right, tau, target, k, h)
		}

		if dist-*tau <= n.Threshold {
			vp.search(n.Left, tau, target, k, h)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2011-2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package zip

import (
	"archive/zip"
	"bytes"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"
)

// FindAll returns the cleaned path of every file in the supplied zip reader.
func FindAll(reader *zip.Reader) ([]string, error) {
	return Find(reader, "*")
}

// Find returns the cleaned path of every file in the supplied zip reader whose
// base name matches the supplied pattern, which is interpreted as in path.Match.
func Find(reader *zip.Reader, pattern string) ([]string, error) {
	// path.Match will only return an error if the pattern is not
	// valid (*and* the supplied name is not empty, hence "check").
	if _, err := path.Match(pattern, "check"); err != nil {
		return nil, err
	}
	var matches []string
	for _, zipFile := range reader.File {
		cleanPath := path.Clean(zipFile.Name)
		baseName := path.Base(cleanPath)
		if match, _ := path.Match(pattern, baseName); match {
			matches = append(matches, cleanPath)
		}
	}
	return matches, nil
}

// ExtractAll extracts the supplied zip reader to the target path, overwriting
// existing files and directories only where necessary.
func ExtractAll(reader *zip.Reader, targetRoot string) error {
	return Extract(reader, targetRoot, "")
}

// Extract extracts files from the supplied zip reader, from the (internal, slash-
// separated) source path into the (external, OS-specific) target path. If the
// source path does not reference a directory, the referenced file will be written
// directly to the target path.
func Extract(reader *zip.Reader, targetRoot, sourceRoot string) error {
	sourceRoot = path.Clean(sourceRoot)
	if sourceRoot == "." {
		sourceRoot = ""
	}
	if !isSanePath(sourceRoot) {
		return fmt.Errorf("cannot extract files rooted at %q", sourceRoot)
	}
	extractor := extractor{targetRoot, sourceRoot}
	for _, zipFile := range reader.File {
		if err := extractor.extract(zipFile); err != nil {
			cleanName := path.Clean(zipFile.Name)
			return fmt.Errorf("cannot extract %q: %v", cleanName, err)
		}
	}
	return nil
}

type extractor struct {
	targetRoot string
	sourceRoot string
}

// targetPath returns the target path for a given zip file and whether
// it should be extracted.
func (x extractor) targetPath(zipFile *zip.File) (string, bool) {
	cleanPath := path.Clean(zipFile.Name)
	if cleanPath == x.sourceRoot {
		return x.targetRoot, true
	}
	if x.sourceRoot != "" {
		mustPrefix := x.sourceRoot + "/"
		if !strings.HasPrefix(cleanPath, mustPrefix) {
			return "", false
		}
		cleanPath = cleanPath[len(mustPrefix):]
	}
	return filepath.Join(x.targetRoot, filepath.FromSlash(cleanPath)), true
}

func (x extractor) extract(zipFile *zip.File) error {
	targetPath, ok := x.targetPath(zipFile)
	if !ok {
		return nil
	}
	parentPath := filepath.Dir(targetPath)
	if err := os.MkdirAll(parentPath, 0777); err != nil {
		return err
	}
	mode := zipFile.Mode()
	modePerm := mode & os.ModePerm
	modeType := mode & os.ModeType
	switch modeType {
	case os.ModeDir:
		return x.writeDir(targetPath, modePerm)
	case os.ModeSymlink:
		return x.writeSymlink(targetPath, zipFile)
	case 0:
		return x.writeFile(targetPath, zipFile, modePerm)
	}
	return fmt.Errorf("unknown file type %d", modeType)
}

func (x extractor) writeDir(targetPath string, modePerm os.FileMode) error {
	fileInfo, err := os.Lstat(targetPath)
	switch {
	case err == nil:
		mode := fileInfo.Mode()
		if mode.IsDir() {
			if mode&os.ModePerm != modePerm {
				return os.Chmod(targetPath, modePerm)
			}
			return nil
		}
		fallthrough
	case !os.IsNotExist(err):
		if err := os.RemoveAll(targetPath); err != nil {
			return err
		}
	}
	return os.MkdirAll(targetPath, modePerm)
}

func (x extractor) writeFile(targetPath string, zipFile *zip.File, modePerm os.FileMode) error {
	if _, err := os.Lstat(targetPath); !os.IsNotExist(err) {
		if err := os.RemoveAll(targetPath); err != nil {
			return err
		}
	}
	writer, err := os.OpenFile(targetPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, modePerm)
	if err != nil {
		return err
	}
	return readTo(writer, zipFile)
}

func (x extractor) writeSymlink(targetPath string, zipFile *zip.File) error {
	symlinkTarget, err := x.checkSymlink(targetPath, zipFile)
	if err != nil {
		return err
	}
	if _, err := os.Lstat(targetPath); !os.IsNotExist(err) {
		if err := os.RemoveAll(targetPath); err != nil {
			return err
		}
	}
	return os.Symlink(symlinkTarget, targetPath)
}

func (x extractor) checkSymlink(targetPath string, zipFile *zip.File) (string, error) {
	var buffer bytes.Buffer
	if err := readTo(&buffer, zipFile); err != nil {
		return "", err
	}
	symlinkTarget := buffer.String()
	if filepath.IsAbs(symlinkTarget) {
		return "", fmt.Errorf("symlink %q is absolute", symlinkTarget)
	}
	finalPath := filepath.Join(filepath.Dir(targetPath), symlinkTarget)
	relativePath, err := filepath.Rel(x.targetRoot, finalPath)
	if err != nil {
		// Not tested, because I don't know how to trigger this condition.
		return "", fmt.Errorf("symlink %q not comprehensible", symlinkTarget)
	}
	if !isSanePath(relativePath) {
		return "", fmt.Errorf("symlink %q leads out of scope", symlinkTarget)
	}
	return symlinkTarget, nil
}

func readTo(writer io.Writer, zipFile *zip.File) error {
	reader, err := zipFile.Open()
	if err != nil {
		return err
	}
	_, err = io.Copy(writer, reader)
	reader.Close()
	return err
}

func isSanePath(path string) bool {
	if path == ".." || strings.HasPrefix(path, "../") {
		return false
	}
	return true
}
<commit_msg>use Manifest when getting charm file list; *didn't* fix getting individual charm files<commit_after>// Copyright 2011-2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package zip

import (
	"archive/zip"
	"bytes"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"
)

// FindAll returns the cleaned path of every file in the supplied zip reader.
func FindAll(reader *zip.Reader) ([]string, error) {
	return Find(reader, "*")
}

// Find returns the cleaned path of every file in the supplied zip reader whose
// base name matches the supplied pattern, which is interpreted as in path.Match.
func Find(reader *zip.Reader, pattern string) ([]string, error) {
	// path.Match will only return an error if the pattern is not
	// valid (*and* the supplied name is not empty, hence "check").
	if _, err := path.Match(pattern, "check"); err != nil {
		return nil, err
	}
	var matches []string
	for _, zipFile := range reader.File {
		cleanPath := path.Clean(zipFile.Name)
		baseName := path.Base(cleanPath)
		if match, _ := path.Match(pattern, baseName); match {
			matches = append(matches, cleanPath)
		}
	}
	return matches, nil
}

// ExtractAll extracts the supplied zip reader to the target path, overwriting
// existing files and directories only where necessary.
func ExtractAll(reader *zip.Reader, targetRoot string) error {
	return Extract(reader, targetRoot, "")
}

// Extract extracts files from the supplied zip reader, from the (internal, slash-
// separated) source path into the (external, OS-specific) target path. If the
// source path does not reference a directory, the referenced file will be written
// directly to the target path.
func Extract(reader *zip.Reader, targetRoot, sourceRoot string) error {
	sourceRoot = path.Clean(sourceRoot)
	if sourceRoot == "." {
		sourceRoot = ""
	}
	if !isSanePath(sourceRoot) {
		return fmt.Errorf("cannot extract files rooted at %q", sourceRoot)
	}
	extractor := extractor{targetRoot, sourceRoot}
	for _, zipFile := range reader.File {
		if err := extractor.extract(zipFile); err != nil {
			cleanName := path.Clean(zipFile.Name)
			return fmt.Errorf("cannot extract %q: %v", cleanName, err)
		}
	}
	return nil
}

type extractor struct {
	targetRoot string
	sourceRoot string
}

// targetPath returns the target path for a given zip file and whether
// it should be extracted.
func (x extractor) targetPath(zipFile *zip.File) (string, bool) {
	cleanPath := path.Clean(zipFile.Name)
	if cleanPath == x.sourceRoot {
		return x.targetRoot, true
	}
	if x.sourceRoot != "" {
		mustPrefix := x.sourceRoot + "/"
		if !strings.HasPrefix(cleanPath, mustPrefix) {
			return "", false
		}
		cleanPath = cleanPath[len(mustPrefix):]
	}
	return filepath.Join(x.targetRoot, filepath.FromSlash(cleanPath)), true
}

func (x extractor) extract(zipFile *zip.File) error {
	targetPath, ok := x.targetPath(zipFile)
	if !ok {
		return nil
	}
	parentPath := filepath.Dir(targetPath)
	if err := os.MkdirAll(parentPath, 0777); err != nil {
		return err
	}
	mode := zipFile.Mode()
	modePerm := mode & os.ModePerm
	modeType := mode & os.ModeType
	switch modeType {
	case os.ModeDir:
		return x.writeDir(targetPath, modePerm)
	case os.ModeSymlink:
		return x.writeSymlink(targetPath, zipFile)
	case 0:
		return x.writeFile(targetPath, zipFile, modePerm)
	}
	return fmt.Errorf("unknown file type %d", modeType)
}

func (x extractor) writeDir(targetPath string, modePerm os.FileMode) error {
	fileInfo, err := os.Lstat(targetPath)
	switch {
	case err == nil:
		mode := fileInfo.Mode()
		if mode.IsDir() {
			if mode&os.ModePerm != modePerm {
				return os.Chmod(targetPath, modePerm)
			}
			return nil
		}
		fallthrough
	case !os.IsNotExist(err):
		if err := os.RemoveAll(targetPath); err != nil {
			return err
		}
	}
	return os.MkdirAll(targetPath, modePerm)
}

func (x extractor) writeFile(targetPath string, zipFile *zip.File, modePerm os.FileMode) error {
	if _, err := os.Lstat(targetPath); !os.IsNotExist(err) {
		if err := os.RemoveAll(targetPath); err != nil {
			return err
		}
	}
	writer, err := os.OpenFile(targetPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, modePerm)
	if err != nil {
		return err
	}
	return copyTo(writer, zipFile)
}

func (x extractor) writeSymlink(targetPath string, zipFile *zip.File) error {
	symlinkTarget, err := x.checkSymlink(targetPath, zipFile)
	if err != nil {
		return err
	}
	if _, err := os.Lstat(targetPath); !os.IsNotExist(err) {
		if err := os.RemoveAll(targetPath); err != nil {
			return err
		}
	}
	return os.Symlink(symlinkTarget, targetPath)
}

func (x extractor) checkSymlink(targetPath string, zipFile *zip.File) (string, error) {
	var buffer bytes.Buffer
	if err := copyTo(&buffer, zipFile); err != nil {
		return "", err
	}
	symlinkTarget := buffer.String()
	if filepath.IsAbs(symlinkTarget) {
		return "", fmt.Errorf("symlink %q is absolute", symlinkTarget)
	}
	finalPath := filepath.Join(filepath.Dir(targetPath), symlinkTarget)
	relativePath, err := filepath.Rel(x.targetRoot, finalPath)
	if err != nil {
		// Not tested, because I don't know how to trigger this condition.
		return "", fmt.Errorf("symlink %q not comprehensible", symlinkTarget)
	}
	if !isSanePath(relativePath) {
		return "", fmt.Errorf("symlink %q leads out of scope", symlinkTarget)
	}
	return symlinkTarget, nil
}

func copyTo(writer io.Writer, zipFile *zip.File) error {
	reader, err := zipFile.Open()
	if err != nil {
		return err
	}
	_, err = io.Copy(writer, reader)
	reader.Close()
	return err
}

func isSanePath(path string) bool {
	if path == ".." || strings.HasPrefix(path, "../") {
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package ezhttp

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"os"
	"strings"
)

// EzClient is a container for the http requests.
type EzClient struct {
	*http.Client
	body io.Reader
}

// NewInsecureClient creates a new SSL client but skips verification.
func NewInsecureClient() *EzClient {
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	return &EzClient{Client: &http.Client{Transport: tr}}
}

// NewClient creates a http connection.
func NewClient() *EzClient {
	return &EzClient{Client: &http.Client{}}
}

// JSONGet performs a get request and decodes the resulting JSON.
func (c *EzClient) JSONGet(url string, out interface{}) (int, error) {
	resp, err := c.Client.Get(url)
	if err != nil {
		return 0, err
	}
	return decodeJsonResponse(resp, out)
}

// FileGet downloads a file and writes it to path.
func (c *EzClient) FileGet(url, path string) (int, error) {
	resp, err := c.Client.Get(url)
	if err != nil {
		return 0, err
	}

	return writeToPath(resp, path)
}

// writeToPath writes the file in the Response to path.
func writeToPath(resp *http.Response, path string) (int, error) {
	defer resp.Body.Close()

	if resp.StatusCode > 299 {
		msg, _ := ioutil.ReadAll(resp.Body)
		return resp.StatusCode, errors.New(string(msg))
	}

	out, err := os.Create(path)
	if err != nil {
		return 0, err
	}
	defer out.Close()

	io.Copy(out, resp.Body)
	return resp.StatusCode, nil
}

// JSONStr takes a string formatted as JSON and makes it available for POST and PUT.
func (c *EzClient) JSONStr(j string) *EzClient {
	c.body = strings.NewReader(j)
	return c
}

// JSON takes an object, marshals it and make it avaiable for POST and PUT.
func (c *EzClient) JSON(j interface{}) *EzClient {
	b, err := json.Marshal(j)
	if err == nil {
		c.body = bytes.NewReader(b)
	}

	return c
}

// JSONPost performs a POST with JSON.
func (c *EzClient) JSONPost(url string, out interface{}) (int, error) {
	resp, err := c.Client.Post(url, "application/json", c.body)
	if err != nil {
		return 0, err
	}
	return decodeJsonResponse(resp, out)
}

// decodeJSONResponse decodes a response containing JSON.
func decodeJSONResponse(resp *http.Response, out interface{}) (int, error) {
	defer resp.Body.Close()
	body, _ := ioutil.ReadAll(resp.Body)

	if resp.StatusCode > 299 {
		return resp.StatusCode, errors.New(string(body))
	}

	err := json.Unmarshal(body, out)
	if err != nil {
		return resp.StatusCode, err
	}
	return resp.StatusCode, nil
}

// PostFile uploads a file.
func (c *EzClient) PostFile(url, filepath, formName string, params map[string]string) (int, error) {
	// Setup body
	body := bytes.NewBufferString("")
	writer := multipart.NewWriter(body)
	defer writer.Close()

	// Create file form
	part, err := writer.CreateFormFile(formName, filepath)
	if err != nil {
		return 0, err
	}

	// Read file into form
	file, err := os.Open(filepath)
	if err != nil {
		return 0, err
	}
	defer file.Close()

	fileContents, err := ioutil.ReadAll(file)
	part.Write(fileContents)

	// Add boundary
	boundary := writer.Boundary()
	closeStr := fmt.Sprintf("\r\n--%s--\r\n", boundary)

	// Add additional params
	if params != nil {
		for key, value := range params {
			writer.WriteField(key, value)
		}
	}

	// Create the request and set headers
	closeBuf := bytes.NewBufferString(closeStr)
	reader := io.MultiReader(body, file, closeBuf)
	req, err := http.NewRequest("POST", url, reader)
	req.Header.Add("Content-Type", "multipart/form-data; boundary="+boundary)
	req.ContentLength = int64(body.Len()) + int64(closeBuf.Len())
	req.Close = true

	// Post the file
	resp, err := c.Client.Do(req)
	defer resp.Body.Close()
	b, _ := ioutil.ReadAll(resp.Body)

	// Return status and any error code
	if resp.StatusCode > 299 {
		return resp.StatusCode, errors.New(string(b))
	}

	return resp.StatusCode, nil
}
<commit_msg>Fix up call to function once we changed its name.<commit_after>package ezhttp

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"os"
	"strings"
)

// EzClient is a container for the http requests.
type EzClient struct {
	*http.Client
	body io.Reader
}

// NewInsecureClient creates a new SSL client but skips verification.
func NewInsecureClient() *EzClient {
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	return &EzClient{Client: &http.Client{Transport: tr}}
}

// NewClient creates a http connection.
func NewClient() *EzClient {
	return &EzClient{Client: &http.Client{}}
}

// JSONGet performs a get request and decodes the resulting JSON.
func (c *EzClient) JSONGet(url string, out interface{}) (int, error) {
	resp, err := c.Client.Get(url)
	if err != nil {
		return 0, err
	}
	return decodeJSONResponse(resp, out)
}

// FileGet downloads a file and writes it to path.
func (c *EzClient) FileGet(url, path string) (int, error) {
	resp, err := c.Client.Get(url)
	if err != nil {
		return 0, err
	}

	return writeToPath(resp, path)
}

// writeToPath writes the file in the Response to path.
func writeToPath(resp *http.Response, path string) (int, error) {
	defer resp.Body.Close()

	if resp.StatusCode > 299 {
		msg, _ := ioutil.ReadAll(resp.Body)
		return resp.StatusCode, errors.New(string(msg))
	}

	out, err := os.Create(path)
	if err != nil {
		return 0, err
	}
	defer out.Close()

	io.Copy(out, resp.Body)
	return resp.StatusCode, nil
}

// JSONStr takes a string formatted as JSON and makes it available for POST and PUT.
func (c *EzClient) JSONStr(j string) *EzClient {
	c.body = strings.NewReader(j)
	return c
}

// JSON takes an object, marshals it and make it avaiable for POST and PUT.
func (c *EzClient) JSON(j interface{}) *EzClient {
	b, err := json.Marshal(j)
	if err == nil {
		c.body = bytes.NewReader(b)
	}

	return c
}

// JSONPost performs a POST with JSON.
func (c *EzClient) JSONPost(url string, out interface{}) (int, error) {
	resp, err := c.Client.Post(url, "application/json", c.body)
	if err != nil {
		return 0, err
	}
	return decodeJSONResponse(resp, out)
}

// decodeJSONResponse decodes a response containing JSON.
func decodeJSONResponse(resp *http.Response, out interface{}) (int, error) {
	defer resp.Body.Close()
	body, _ := ioutil.ReadAll(resp.Body)

	if resp.StatusCode > 299 {
		return resp.StatusCode, errors.New(string(body))
	}

	err := json.Unmarshal(body, out)
	if err != nil {
		return resp.StatusCode, err
	}
	return resp.StatusCode, nil
}

// PostFile uploads a file.
func (c *EzClient) PostFile(url, filepath, formName string, params map[string]string) (int, error) {
	// Setup body
	body := bytes.NewBufferString("")
	writer := multipart.NewWriter(body)
	defer writer.Close()

	// Create file form
	part, err := writer.CreateFormFile(formName, filepath)
	if err != nil {
		return 0, err
	}

	// Read file into form
	file, err := os.Open(filepath)
	if err != nil {
		return 0, err
	}
	defer file.Close()

	fileContents, err := ioutil.ReadAll(file)
	part.Write(fileContents)

	// Add boundary
	boundary := writer.Boundary()
	closeStr := fmt.Sprintf("\r\n--%s--\r\n", boundary)

	// Add additional params
	if params != nil {
		for key, value := range params {
			writer.WriteField(key, value)
		}
	}

	// Create the request and set headers
	closeBuf := bytes.NewBufferString(closeStr)
	reader := io.MultiReader(body, file, closeBuf)
	req, err := http.NewRequest("POST", url, reader)
	req.Header.Add("Content-Type", "multipart/form-data; boundary="+boundary)
	req.ContentLength = int64(body.Len()) + int64(closeBuf.Len())
	req.Close = true

	// Post the file
	resp, err := c.Client.Do(req)
	defer resp.Body.Close()
	b, _ := ioutil.ReadAll(resp.Body)

	// Return status and any error code
	if resp.StatusCode > 299 {
		return resp.StatusCode, errors.New(string(b))
	}

	return resp.StatusCode, nil
}
<|endoftext|>


<commit_before>package fancylogger

import (
	"io"
	"log"

	"github.com/mgutz/ansi"
)

var (
	// Success is the color of success
	Success = ansi.ColorCode("green+bh")
	// Info is the color of information
	Info = ansi.ColorCode("white+bh")
	// Error is the color of a meltdown
	Error = ansi.ColorCode("red+bh")
	// Warning means gun it
	Warning = ansi.ColorCode("yellow+bh")
	// Reset puts the colors back the way it was
	Reset = ansi.ColorCode("reset")
)

// FancyLogger will give you nicely printed color logs
type FancyLogger struct {
	Trace   *log.Logger
	Info    *log.Logger
	Warning *log.Logger
	Error   *log.Logger
}

// NewFancyLogger will set up a fancy logger instance, but requires all of
// outputs to be defined
func NewFancyLogger(traceHandle io.Writer, infoHandle io.Writer,
	warningHandle io.Writer, errorHandle io.Writer) *FancyLogger {
	return &FancyLogger{
		Trace: log.New(traceHandle, "TRACE: ",
			log.Ldate|log.Ltime|log.Lmicroseconds|log.Lshortfile),
		Info: log.New(infoHandle, Info+"INFO: "+Reset,
			log.Ldate|log.Ltime|log.Lmicroseconds),
		Warning: log.New(warningHandle, Warning+"WARNING: "+Reset,
			log.Ldate|log.Ltime|log.Lmicroseconds|log.Lshortfile),
		Error: log.New(errorHandle, Error+"ERROR: "+Reset,
			log.Ldate|log.Ltime|log.Lmicroseconds|log.Lshortfile),
	}
}
<commit_msg>Add color for Trace and fix yellow for Waring<commit_after>package fancylogger

import (
	"io"
	"log"

	"github.com/mgutz/ansi"
)

var (
	// Success is the color of success
	Success = ansi.ColorCode("green+bh")
	// Info is the color of information
	Info = ansi.ColorCode("white+bh")
	// Error is the color of a meltdown
	Error = ansi.ColorCode("red+bh")
	// Warning means gun it
	Warning = ansi.ColorCode("yellow")
	// Trace will give you tracer rounds
	Trace = ansi.ColorCode("magenta+bh")
	// Reset puts the colors back the way it was
	Reset = ansi.ColorCode("reset")
)

// FancyLogger will give you nicely printed color logs
type FancyLogger struct {
	Trace   *log.Logger
	Info    *log.Logger
	Warning *log.Logger
	Error   *log.Logger
}

// NewFancyLogger will set up a fancy logger instance, but requires all of
// outputs to be defined
func NewFancyLogger(traceHandle io.Writer, infoHandle io.Writer,
	warningHandle io.Writer, errorHandle io.Writer) *FancyLogger {
	return &FancyLogger{
		Trace: log.New(traceHandle, Trace+"TRACE: "+Reset,
			log.Ldate|log.Ltime|log.Lmicroseconds|log.Lshortfile),
		Info: log.New(infoHandle, Info+"INFO: "+Reset,
			log.Ldate|log.Ltime|log.Lmicroseconds),
		Warning: log.New(warningHandle, Warning+"WARNING: "+Reset,
			log.Ldate|log.Ltime|log.Lmicroseconds|log.Lshortfile),
		Error: log.New(errorHandle, Error+"ERROR: "+Reset,
			log.Ldate|log.Ltime|log.Lmicroseconds|log.Lshortfile),
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"github.com/codegangsta/cli"
	"io"
	"io/ioutil"
	"log"
	"os"
	"regexp"
	"text/template"
	"unicode"
)

const (
	SYMBOL = iota
	STRUCT_BEGIN
	STRUCT_END
	DATA_TYPE
	ARRAY_TYPE
	TK_EOF
)

var (
	datatypes = map[string]bool{
		"integer": true,
		"string":  true,
		"bytes":   true,
		"byte":    true,
		"boolean": true,
		"float":   true,
	}

	funcs map[string]lang_type
)

var (
	TOKEN_EOF = &token{typ: TK_EOF}
)

type func_info struct {
	T string `json:"t"` // type
	R string `json:"r"` // read
	W string `json:"w"` // write
}

type lang_type struct {
	Go func_info `json:"go"` // golang
	Cs func_info `json:"cs"` // c#
}
type (
	field_info struct {
		Name  string
		Typ   string
		Array bool
	}
	struct_info struct {
		Name   string
		Fields []field_info
	}
)

type token struct {
	typ     int
	literal string
	r       rune
}

func syntax_error(p *Parser) {
	log.Fatal("syntax error @line:", p.lexer.lineno)
}

type Lexer struct {
	reader *bytes.Buffer
	lineno int
}

func (lex *Lexer) init(r io.Reader) {
	bts, err := ioutil.ReadAll(r)
	if err != nil {
		log.Println(err)
	}

	// 清除注释
	re := regexp.MustCompile("(?m:^#(.*)$)")
	bts = re.ReplaceAllLiteral(bts, nil)
	lex.reader = bytes.NewBuffer(bts)
	lex.lineno = 1
}

func (lex *Lexer) next() (t *token) {
	defer func() {
		//log.Println(t)
	}()
	var r rune
	var err error
	for {
		r, _, err = lex.reader.ReadRune()
		if err == io.EOF {
			return TOKEN_EOF
		} else if unicode.IsSpace(r) {
			if r == '\n' {
				lex.lineno++
			}
			continue
		}
		break
	}

	if r == '=' {
		for k := 0; k < 2; k++ { // check "==="
			r, _, err = lex.reader.ReadRune()
			if err == io.EOF {
				return TOKEN_EOF
			}
			if r != '=' {
				lex.reader.UnreadRune()
				return &token{typ: STRUCT_BEGIN}
			}
		}
		return &token{typ: STRUCT_END}
	} else if unicode.IsLetter(r) {
		var runes []rune
		for {
			runes = append(runes, r)
			r, _, err = lex.reader.ReadRune()
			if err == io.EOF {
				break
			} else if unicode.IsLetter(r) || unicode.IsNumber(r) || r == '_' {
				continue
			} else {
				lex.reader.UnreadRune()
				break
			}
		}

		t := &token{}
		t.literal = string(runes)
		if datatypes[t.literal] {
			t.typ = DATA_TYPE
		} else if t.literal == "array" {
			t.typ = ARRAY_TYPE
		} else {
			t.typ = SYMBOL
		}

		return t
	} else {
		log.Fatal("lex error @line:", lex.lineno)
	}
	return nil
}

func (lex *Lexer) eof() bool {
	for {
		r, _, err := lex.reader.ReadRune()
		if err == io.EOF {
			return true
		} else if unicode.IsSpace(r) {
			if r == '\n' {
				lex.lineno++
			}
			continue
		} else {
			lex.reader.UnreadRune()
			return false
		}
	}
}

//////////////////////////////////////////////////////////////
type Parser struct {
	lexer *Lexer
	info  []struct_info
}

func (p *Parser) init(lex *Lexer) {
	p.lexer = lex
}

func (p *Parser) match(typ int) *token {
	t := p.lexer.next()
	if t.typ != typ {
		syntax_error(p)
	}
	return t
}

func (p *Parser) expr() bool {
	if p.lexer.eof() {
		return false
	}
	info := struct_info{}

	t := p.match(SYMBOL)
	info.Name = t.literal

	p.match(STRUCT_BEGIN)
	p.fields(&info)
	p.info = append(p.info, info)
	return true
}

func (p *Parser) fields(info *struct_info) {
	for {
		t := p.lexer.next()
		if t.typ == STRUCT_END {
			return
		}
		if t.typ != SYMBOL {
			syntax_error(p)
		}

		field := field_info{Name: t.literal}
		t = p.lexer.next()
		if t.typ == ARRAY_TYPE {
			field.Array = true
			t = p.match(SYMBOL)
			field.Typ = t.literal
		} else if t.typ == DATA_TYPE || t.typ == SYMBOL {
			field.Typ = t.literal
		} else {
			syntax_error(p)
		}

		info.Fields = append(info.Fields, field)
	}
}

func main() {
	app := cli.NewApp()
	app.Name = "Protocol Data Structure Generator"
	app.Usage = "handle proto.txt"
	app.Authors = []cli.Author{{Name: "xtaci"}, {Name: "ycs"}}
	app.Version = "1.0"
	app.Flags = []cli.Flag{
		cli.StringFlag{Name: "file,f", Value: "./proto.txt", Usage: "input proto.txt file"},
		cli.StringFlag{Name: "template,t", Value: "./templates/server/proto.tmpl", Usage: "template file"},
	}
	app.Action = func(c *cli.Context) {
		// parse
		file, err := os.Open(c.String("file"))
		if err != nil {
			log.Fatal(err)
		}
		lexer := Lexer{}
		lexer.init(file)
		p := Parser{}
		p.init(&lexer)
		for p.expr() {
		}

		// function mapping
		f, err := os.Open("func_map.json")
		if err != nil {
			log.Fatal(err)
		}
		dec := json.NewDecoder(f)

		// read open bracket
		_, err = dec.Token()
		if err != nil {
			log.Fatal(err)
		}

		for dec.More() {
			// decode an array value (Message)
			err := dec.Decode(&funcs)
			if err != nil {
				log.Fatal(err)
			}
		}

		// read closing bracket
		_, err = dec.Token()
		if err != nil {
			log.Fatal(err)
		}

		funcMap := template.FuncMap{
			"goType": func(t string) string {
				if v, ok := funcs[t]; ok {
					return v.Go.T
				} else {
					return ""
				}
			},
			"goRead": func(t string) string {
				if v, ok := funcs[t]; ok {
					return v.Go.R
				} else {
					return ""
				}
			},
			"goWrite": func(t string) string {
				if v, ok := funcs[t]; ok {
					return v.Go.W
				} else {
					return ""
				}
			},
			"csType": func(t string) string {
				if v, ok := funcs[t]; ok {
					return v.Cs.T
				} else {
					return ""
				}
			},
			"csRead": func(t string) string {
				if v, ok := funcs[t]; ok {
					return v.Cs.R
				} else {
					return ""
				}
			},
			"csWrite": func(t string) string {
				if v, ok := funcs[t]; ok {
					return v.Cs.W
				} else {
					return ""
				}
			},
		}
		tmpl, err := template.New("proto.tmpl").Funcs(funcMap).ParseFiles(c.String("template"))
		if err != nil {
			log.Fatal(err)
		}
		err = tmpl.Execute(os.Stdout, p.info)
		if err != nil {
			log.Fatal(err)
		}
	}
	app.Run(os.Args)
}
<commit_msg>adjust consts names<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"github.com/codegangsta/cli"
	"io"
	"io/ioutil"
	"log"
	"os"
	"regexp"
	"text/template"
	"unicode"
)

const (
	TK_SYMBOL = iota
	TK_STRUCT_BEGIN
	TK_STRUCT_END
	TK_DATA_TYPE
	TK_ARRAY
	TK_EOF
)

var (
	datatypes = map[string]bool{
		"integer": true,
		"string":  true,
		"bytes":   true,
		"byte":    true,
		"boolean": true,
		"float":   true,
	}

	funcs map[string]lang_type
)

var (
	TOKEN_EOF = &token{typ: TK_EOF}
)

type func_info struct {
	T string `json:"t"` // type
	R string `json:"r"` // read
	W string `json:"w"` // write
}

type lang_type struct {
	Go func_info `json:"go"` // golang
	Cs func_info `json:"cs"` // c#
}
type (
	field_info struct {
		Name  string
		Typ   string
		Array bool
	}
	struct_info struct {
		Name   string
		Fields []field_info
	}
)

type token struct {
	typ     int
	literal string
	r       rune
}

func syntax_error(p *Parser) {
	log.Fatal("syntax error @line:", p.lexer.lineno)
}

type Lexer struct {
	reader *bytes.Buffer
	lineno int
}

func (lex *Lexer) init(r io.Reader) {
	bts, err := ioutil.ReadAll(r)
	if err != nil {
		log.Println(err)
	}

	// 清除注释
	re := regexp.MustCompile("(?m:^#(.*)$)")
	bts = re.ReplaceAllLiteral(bts, nil)
	lex.reader = bytes.NewBuffer(bts)
	lex.lineno = 1
}

func (lex *Lexer) next() (t *token) {
	defer func() {
		//log.Println(t)
	}()
	var r rune
	var err error
	for {
		r, _, err = lex.reader.ReadRune()
		if err == io.EOF {
			return TOKEN_EOF
		} else if unicode.IsSpace(r) {
			if r == '\n' {
				lex.lineno++
			}
			continue
		}
		break
	}

	if r == '=' {
		for k := 0; k < 2; k++ { // check "==="
			r, _, err = lex.reader.ReadRune()
			if err == io.EOF {
				return TOKEN_EOF
			}
			if r != '=' {
				lex.reader.UnreadRune()
				return &token{typ: TK_STRUCT_BEGIN}
			}
		}
		return &token{typ: TK_STRUCT_END}
	} else if unicode.IsLetter(r) {
		var runes []rune
		for {
			runes = append(runes, r)
			r, _, err = lex.reader.ReadRune()
			if err == io.EOF {
				break
			} else if unicode.IsLetter(r) || unicode.IsNumber(r) || r == '_' {
				continue
			} else {
				lex.reader.UnreadRune()
				break
			}
		}

		t := &token{}
		t.literal = string(runes)
		if datatypes[t.literal] {
			t.typ = TK_DATA_TYPE
		} else if t.literal == "array" {
			t.typ = TK_ARRAY
		} else {
			t.typ = TK_SYMBOL
		}

		return t
	} else {
		log.Fatal("lex error @line:", lex.lineno)
	}
	return nil
}

func (lex *Lexer) eof() bool {
	for {
		r, _, err := lex.reader.ReadRune()
		if err == io.EOF {
			return true
		} else if unicode.IsSpace(r) {
			if r == '\n' {
				lex.lineno++
			}
			continue
		} else {
			lex.reader.UnreadRune()
			return false
		}
	}
}

//////////////////////////////////////////////////////////////
type Parser struct {
	lexer *Lexer
	info  []struct_info
}

func (p *Parser) init(lex *Lexer) {
	p.lexer = lex
}

func (p *Parser) match(typ int) *token {
	t := p.lexer.next()
	if t.typ != typ {
		syntax_error(p)
	}
	return t
}

func (p *Parser) expr() bool {
	if p.lexer.eof() {
		return false
	}
	info := struct_info{}

	t := p.match(TK_SYMBOL)
	info.Name = t.literal

	p.match(TK_STRUCT_BEGIN)
	p.fields(&info)
	p.info = append(p.info, info)
	return true
}

func (p *Parser) fields(info *struct_info) {
	for {
		t := p.lexer.next()
		if t.typ == TK_STRUCT_END {
			return
		}
		if t.typ != TK_SYMBOL {
			syntax_error(p)
		}

		field := field_info{Name: t.literal}
		t = p.lexer.next()
		if t.typ == TK_ARRAY {
			field.Array = true
			t = p.match(TK_SYMBOL)
			field.Typ = t.literal
		} else if t.typ == TK_DATA_TYPE || t.typ == TK_SYMBOL {
			field.Typ = t.literal
		} else {
			syntax_error(p)
		}

		info.Fields = append(info.Fields, field)
	}
}

func main() {
	app := cli.NewApp()
	app.Name = "Protocol Data Structure Generator"
	app.Usage = "handle proto.txt"
	app.Authors = []cli.Author{{Name: "xtaci"}, {Name: "ycs"}}
	app.Version = "1.0"
	app.Flags = []cli.Flag{
		cli.StringFlag{Name: "file,f", Value: "./proto.txt", Usage: "input proto.txt file"},
		cli.StringFlag{Name: "template,t", Value: "./templates/server/proto.tmpl", Usage: "template file"},
	}
	app.Action = func(c *cli.Context) {
		// parse
		file, err := os.Open(c.String("file"))
		if err != nil {
			log.Fatal(err)
		}
		lexer := Lexer{}
		lexer.init(file)
		p := Parser{}
		p.init(&lexer)
		for p.expr() {
		}

		// function mapping
		f, err := os.Open("func_map.json")
		if err != nil {
			log.Fatal(err)
		}
		dec := json.NewDecoder(f)

		// read open bracket
		_, err = dec.Token()
		if err != nil {
			log.Fatal(err)
		}

		for dec.More() {
			// decode an array value (Message)
			err := dec.Decode(&funcs)
			if err != nil {
				log.Fatal(err)
			}
		}

		// read closing bracket
		_, err = dec.Token()
		if err != nil {
			log.Fatal(err)
		}

		funcMap := template.FuncMap{
			"goType": func(t string) string {
				if v, ok := funcs[t]; ok {
					return v.Go.T
				} else {
					return ""
				}
			},
			"goRead": func(t string) string {
				if v, ok := funcs[t]; ok {
					return v.Go.R
				} else {
					return ""
				}
			},
			"goWrite": func(t string) string {
				if v, ok := funcs[t]; ok {
					return v.Go.W
				} else {
					return ""
				}
			},
			"csType": func(t string) string {
				if v, ok := funcs[t]; ok {
					return v.Cs.T
				} else {
					return ""
				}
			},
			"csRead": func(t string) string {
				if v, ok := funcs[t]; ok {
					return v.Cs.R
				} else {
					return ""
				}
			},
			"csWrite": func(t string) string {
				if v, ok := funcs[t]; ok {
					return v.Cs.W
				} else {
					return ""
				}
			},
		}
		tmpl, err := template.New("proto.tmpl").Funcs(funcMap).ParseFiles(c.String("template"))
		if err != nil {
			log.Fatal(err)
		}
		err = tmpl.Execute(os.Stdout, p.info)
		if err != nil {
			log.Fatal(err)
		}
	}
	app.Run(os.Args)
}
<|endoftext|>