<commit_before>package pinger

import (
	"fmt"
	"log"
	"math/rand"
	"net"
	"sync"
	"syscall"
	"time"
	"unsafe"

	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
)

const ProtocolICMP = 1

type PingStats struct {
	Latency  []time.Duration
	Sent     int
	Received int
	SentTime map[int]time.Time
}

type EchoRequest struct {
	Peer     string
	Count    int
	Deadline time.Time
	Id       int
	Stats    *PingStats
	Done     chan *PingStats
	Recv     chan *EchoResponse

	m      sync.RWMutex
	pinger *Pinger
}

type EchoResponse struct {
	Peer     string
	Id       int
	Seq      int
	Received time.Time
}

func (e *EchoResponse) String() string {
	return fmt.Sprintf("Peer %s, Id: %d, Seq: %d, Recv: %s", e.Peer, e.Id, e.Seq, e.Received)
}

func (e *EchoRequest) Listen() {
	seqRecv := make(map[int]bool)

	timer := time.NewTimer(e.Deadline.Sub(time.Now()))
WAIT:
	for {
		select {
		case <-timer.C:
			// deadline reached.
			log.Printf("go-pinger: deadline reached waiting for repsonse. Peer: %s, Id: %d", e.Peer, e.Id)
			break WAIT
		case resp := <-e.Recv:
			e.m.Lock()
			rtt := resp.Received.Sub(e.Stats.SentTime[resp.Seq])
			if rtt < 0 {
				rtt = 1
			}
			e.Stats.Received++
			e.Stats.Latency = append(e.Stats.Latency, rtt)
			e.m.Unlock()

			seqRecv[resp.Seq] = true
			// we have recieved responses for all pings sent.
			if e.Stats.Received >= e.Stats.Sent {
				break WAIT
			}
		}
	}
	timer.Stop()
	for i := 0; i < e.Count; i++ {
		if _, ok := seqRecv[i]; !ok {
			key := packetKey(e.Peer, e.Id, i)
			e.pinger.DeleteKey(key)
		}
	}

	e.Done <- e.Stats
}

func (e *EchoRequest) Send() {
	for i := 0; i < e.Count; i++ {
		pkt := icmp.Message{
			Type: ipv4.ICMPTypeEcho,
			Code: 0,
			Body: &icmp.Echo{
				ID:   e.Id,
				Seq:  i,
				Data: []byte("raintank/go-pinger"),
			},
		}
		wb, err := pkt.Marshal(nil)
		if err != nil {
			if e.pinger.Debug {
				log.Printf("failed to marshal ICMP Echo packet. %s", err)
			}
			continue
		}
		sentTime, err := e.pinger.WritePkt(wb, e.Peer)
		if err == nil {
			e.Stats.Sent++
			e.Stats.SentTime[i] = sentTime
			if e.pinger.Debug {
				log.Printf("go-pinger: sent pkt. Peer %s, Id: %d, Seq: %d, Sent: %s", e.Peer, e.Id, i, sentTime)
			}
		}
	}
	e.Listen()
}

type Pinger struct {
	queue      map[string]chan *EchoResponse
	m          sync.RWMutex
	running    bool
	conn       net.PacketConn
	Debug      bool
	packetChan chan EchoResponse
	Counter    int
}

func NewPinger() *Pinger {
	rand.Seed(time.Now().UnixNano())
	return &Pinger{queue: make(map[string]chan *EchoResponse), Counter: rand.Intn(0xffff)}
}

func packetKey(addr string, id, seq int) string {
	return fmt.Sprintf("%s-%d-%d", addr, id, seq)
}

func (p *Pinger) Ping(address string, count int, deadline time.Time) (<-chan *PingStats, error) {
	// ensuire the IP address is valid.
	ipAddr := net.ParseIP(address)
	if ipAddr == nil {
		return nil, fmt.Errorf("Failed to parse IP address")
	}

	p.m.Lock()
	defer p.m.Unlock()

	p.Counter++
	if p.Counter > 65535 {
		p.Counter = 0
	}

	req := &EchoRequest{
		Peer:     address,
		Count:    count,
		Deadline: deadline,
		Id:       p.Counter,

		Done:  make(chan *PingStats),
		Recv:  make(chan *EchoResponse, count),
		Stats: &PingStats{Latency: make([]time.Duration, 0), SentTime: make(map[int]time.Time)},

		pinger: p,
	}

	for i := 0; i < count; i++ {
		key := packetKey(req.Peer, req.Id, i)
		p.queue[key] = req.Recv
	}

	if !p.running {
		log.Println("First ping, creating raw Socket.\n")
		p.start()
	}

	go req.Send()

	return req.Done, nil
}

func (p *Pinger) start() {
	c, err := net.ListenPacket("ip4:icmp", "0.0.0.0")
	if err != nil {
		panic(err)
	}
	if p.Debug {
		log.Printf("Listening on socket for icmp packets.")
	}
	if c == nil {
		panic("packetConn is nil\n")
	}
	p.conn = c
	p.running = true
	p.packetChan = make(chan EchoResponse, 10000)
	go p.listenIpv4()
	go p.processPkt()
}

func (p *Pinger) Stop() {
	p.running = false
	p.conn.Close()
	if p.Debug {
		log.Printf("Socket closed.")
	}
	p.conn = nil
	close(p.packetChan)
}

func (p *Pinger) listenIpv4() {
	if p.conn == nil {
		panic("conn doesnt exist")
	}
	log.Printf("starting rawSocket listener\n")
	rb := make([]byte, 1500)
	pkt := EchoResponse{}
	var readErr error
	var data []byte
	ipconn, ok := p.conn.(*net.IPConn)
	if !ok {
		panic("connection is not IPConn")
	}
	file, err := ipconn.File()
	if err != nil {
		panic(err.Error())
	}
	defer file.Close()
	fd := file.Fd()

	var pktTime time.Time
	recvTime := syscall.Timeval{}
	for {
		if p.conn == nil {
			break
		}
		n, peer, err := p.conn.ReadFrom(rb)
		if err != nil {
			readErr = err
			break
		}
		_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), uintptr(syscall.SIOCGSTAMP), uintptr(unsafe.Pointer(&recvTime)))
		err = nil
		if errno != 0 {
			err = errno
		}
		if err == nil {
			pktTime = time.Unix(0, recvTime.Nano())
		} else {
			pktTime = time.Now()
		}

		rm, err := icmp.ParseMessage(ProtocolICMP, rb[:n])
		if err != nil {
			fmt.Println(err.Error())
			continue
		}
		if rm.Type == ipv4.ICMPTypeEchoReply {
			data = rm.Body.(*icmp.Echo).Data
			if len(data) < 9 {
				log.Printf("go-pinger: invalid data payload from %s. Expected at least 9bytes got %d", peer.String(), len(data))
				continue
			}
			pkt = EchoResponse{
				Peer:     peer.String(),
				Seq:      rm.Body.(*icmp.Echo).Seq,
				Id:       rm.Body.(*icmp.Echo).ID,
				Received: pktTime,
			}
			if p.Debug {
				log.Printf("go-pinger: recieved pkt. %s\n", pkt.String())
			}
			select {
			case p.packetChan <- pkt:
			default:
				log.Printf("go-pinger: droped echo response due to blocked packetChan. %s\n", pkt.String())
			}
		}
	}
	if p.Debug {
		log.Printf("listen loop ended.")
	}
	p.m.Lock()
	if p.running {
		log.Println(readErr.Error())
		p.Stop()
		p.start()
	}
	p.m.Unlock()
}

func (p *Pinger) processPkt() {
	for pkt := range p.packetChan {
		key := packetKey(pkt.Peer, pkt.Id, pkt.Seq)
		p.m.RLock()
		req, ok := p.queue[key]
		delete(p.queue, key)
		p.m.RUnlock()
		if ok {
			if p.Debug {
				log.Printf("reply packet matches request packet. %s\n", pkt.String())
			}
			select {
			case req <- &pkt:
			default:
				log.Printf("go-pinger: droped echo response due to blocked response chan. %s", pkt.String())
			}
		} else {
			if p.Debug {
				log.Printf("go-pinger: unexpected echo response. %s\n", pkt.String())
			}
		}
	}
}

func (p *Pinger) DeleteKey(key string) {
	p.m.Lock()
	delete(p.queue, key)
	p.m.Unlock()
}

func (p *Pinger) WritePkt(b []byte, dst string) (time.Time, error) {
	if _, err := p.conn.WriteTo(b, &net.IPAddr{IP: net.ParseIP(dst)}); err != nil {
		return time.Now(), err
	}
	return time.Now(), nil
}
<commit_msg>remove debug log statement<commit_after>package pinger

import (
	"fmt"
	"log"
	"math/rand"
	"net"
	"sync"
	"syscall"
	"time"
	"unsafe"

	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
)

const ProtocolICMP = 1

type PingStats struct {
	Latency  []time.Duration
	Sent     int
	Received int
	SentTime map[int]time.Time
}

type EchoRequest struct {
	Peer     string
	Count    int
	Deadline time.Time
	Id       int
	Stats    *PingStats
	Done     chan *PingStats
	Recv     chan *EchoResponse

	m      sync.RWMutex
	pinger *Pinger
}

type EchoResponse struct {
	Peer     string
	Id       int
	Seq      int
	Received time.Time
}

func (e *EchoResponse) String() string {
	return fmt.Sprintf("Peer %s, Id: %d, Seq: %d, Recv: %s", e.Peer, e.Id, e.Seq, e.Received)
}

func (e *EchoRequest) Listen() {
	seqRecv := make(map[int]bool)

	timer := time.NewTimer(e.Deadline.Sub(time.Now()))
WAIT:
	for {
		select {
		case <-timer.C:
			// deadline reached.
			// log.Printf("go-pinger: deadline reached waiting for repsonse. Peer: %s, Id: %d", e.Peer, e.Id)
			break WAIT
		case resp := <-e.Recv:
			e.m.Lock()
			rtt := resp.Received.Sub(e.Stats.SentTime[resp.Seq])
			if rtt < 0 {
				rtt = 1
			}
			e.Stats.Received++
			e.Stats.Latency = append(e.Stats.Latency, rtt)
			e.m.Unlock()

			seqRecv[resp.Seq] = true
			// we have recieved responses for all pings sent.
			if e.Stats.Received >= e.Stats.Sent {
				break WAIT
			}
		}
	}
	timer.Stop()
	for i := 0; i < e.Count; i++ {
		if _, ok := seqRecv[i]; !ok {
			key := packetKey(e.Peer, e.Id, i)
			e.pinger.DeleteKey(key)
		}
	}

	e.Done <- e.Stats
}

func (e *EchoRequest) Send() {
	for i := 0; i < e.Count; i++ {
		pkt := icmp.Message{
			Type: ipv4.ICMPTypeEcho,
			Code: 0,
			Body: &icmp.Echo{
				ID:   e.Id,
				Seq:  i,
				Data: []byte("raintank/go-pinger"),
			},
		}
		wb, err := pkt.Marshal(nil)
		if err != nil {
			if e.pinger.Debug {
				log.Printf("failed to marshal ICMP Echo packet. %s", err)
			}
			continue
		}
		sentTime, err := e.pinger.WritePkt(wb, e.Peer)
		if err == nil {
			e.Stats.Sent++
			e.Stats.SentTime[i] = sentTime
			if e.pinger.Debug {
				log.Printf("go-pinger: sent pkt. Peer %s, Id: %d, Seq: %d, Sent: %s", e.Peer, e.Id, i, sentTime)
			}
		}
	}
	e.Listen()
}

type Pinger struct {
	queue      map[string]chan *EchoResponse
	m          sync.RWMutex
	running    bool
	conn       net.PacketConn
	Debug      bool
	packetChan chan EchoResponse
	Counter    int
}

func NewPinger() *Pinger {
	rand.Seed(time.Now().UnixNano())
	return &Pinger{queue: make(map[string]chan *EchoResponse), Counter: rand.Intn(0xffff)}
}

func packetKey(addr string, id, seq int) string {
	return fmt.Sprintf("%s-%d-%d", addr, id, seq)
}

func (p *Pinger) Ping(address string, count int, deadline time.Time) (<-chan *PingStats, error) {
	// ensuire the IP address is valid.
	ipAddr := net.ParseIP(address)
	if ipAddr == nil {
		return nil, fmt.Errorf("Failed to parse IP address")
	}

	p.m.Lock()
	defer p.m.Unlock()

	p.Counter++
	if p.Counter > 65535 {
		p.Counter = 0
	}

	req := &EchoRequest{
		Peer:     address,
		Count:    count,
		Deadline: deadline,
		Id:       p.Counter,

		Done:  make(chan *PingStats),
		Recv:  make(chan *EchoResponse, count),
		Stats: &PingStats{Latency: make([]time.Duration, 0), SentTime: make(map[int]time.Time)},

		pinger: p,
	}

	for i := 0; i < count; i++ {
		key := packetKey(req.Peer, req.Id, i)
		p.queue[key] = req.Recv
	}

	if !p.running {
		log.Println("First ping, creating raw Socket.\n")
		p.start()
	}

	go req.Send()

	return req.Done, nil
}

func (p *Pinger) start() {
	c, err := net.ListenPacket("ip4:icmp", "0.0.0.0")
	if err != nil {
		panic(err)
	}
	if p.Debug {
		log.Printf("Listening on socket for icmp packets.")
	}
	if c == nil {
		panic("packetConn is nil\n")
	}
	p.conn = c
	p.running = true
	p.packetChan = make(chan EchoResponse, 10000)
	go p.listenIpv4()
	go p.processPkt()
}

func (p *Pinger) Stop() {
	p.running = false
	p.conn.Close()
	if p.Debug {
		log.Printf("Socket closed.")
	}
	p.conn = nil
	close(p.packetChan)
}

func (p *Pinger) listenIpv4() {
	if p.conn == nil {
		panic("conn doesnt exist")
	}
	log.Printf("starting rawSocket listener\n")
	rb := make([]byte, 1500)
	pkt := EchoResponse{}
	var readErr error
	var data []byte
	ipconn, ok := p.conn.(*net.IPConn)
	if !ok {
		panic("connection is not IPConn")
	}
	file, err := ipconn.File()
	if err != nil {
		panic(err.Error())
	}
	defer file.Close()
	fd := file.Fd()

	var pktTime time.Time
	recvTime := syscall.Timeval{}
	for {
		if p.conn == nil {
			break
		}
		n, peer, err := p.conn.ReadFrom(rb)
		if err != nil {
			readErr = err
			break
		}
		_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), uintptr(syscall.SIOCGSTAMP), uintptr(unsafe.Pointer(&recvTime)))
		err = nil
		if errno != 0 {
			err = errno
		}
		if err == nil {
			pktTime = time.Unix(0, recvTime.Nano())
		} else {
			pktTime = time.Now()
		}

		rm, err := icmp.ParseMessage(ProtocolICMP, rb[:n])
		if err != nil {
			fmt.Println(err.Error())
			continue
		}
		if rm.Type == ipv4.ICMPTypeEchoReply {
			data = rm.Body.(*icmp.Echo).Data
			if len(data) < 9 {
				log.Printf("go-pinger: invalid data payload from %s. Expected at least 9bytes got %d", peer.String(), len(data))
				continue
			}
			pkt = EchoResponse{
				Peer:     peer.String(),
				Seq:      rm.Body.(*icmp.Echo).Seq,
				Id:       rm.Body.(*icmp.Echo).ID,
				Received: pktTime,
			}
			if p.Debug {
				log.Printf("go-pinger: recieved pkt. %s\n", pkt.String())
			}
			select {
			case p.packetChan <- pkt:
			default:
				log.Printf("go-pinger: droped echo response due to blocked packetChan. %s\n", pkt.String())
			}
		}
	}
	if p.Debug {
		log.Printf("listen loop ended.")
	}
	p.m.Lock()
	if p.running {
		log.Println(readErr.Error())
		p.Stop()
		p.start()
	}
	p.m.Unlock()
}

func (p *Pinger) processPkt() {
	for pkt := range p.packetChan {
		key := packetKey(pkt.Peer, pkt.Id, pkt.Seq)
		p.m.RLock()
		req, ok := p.queue[key]
		delete(p.queue, key)
		p.m.RUnlock()
		if ok {
			if p.Debug {
				log.Printf("reply packet matches request packet. %s\n", pkt.String())
			}
			select {
			case req <- &pkt:
			default:
				log.Printf("go-pinger: droped echo response due to blocked response chan. %s", pkt.String())
			}
		} else {
			if p.Debug {
				log.Printf("go-pinger: unexpected echo response. %s\n", pkt.String())
			}
		}
	}
}

func (p *Pinger) DeleteKey(key string) {
	p.m.Lock()
	delete(p.queue, key)
	p.m.Unlock()
}

func (p *Pinger) WritePkt(b []byte, dst string) (time.Time, error) {
	if _, err := p.conn.WriteTo(b, &net.IPAddr{IP: net.ParseIP(dst)}); err != nil {
		return time.Now(), err
	}
	return time.Now(), nil
}
<|endoftext|>


<commit_before>package web_test

import (
	"fmt"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/sclevine/agouti/matchers"

	"github.com/concourse/atc"
)

var _ = Describe("Viewing builds", func() {
	Context("with a job build", func() {
		var build atc.Build

		BeforeEach(func() {
			_, _, err := client.CreateOrUpdatePipelineConfig(pipelineName, "0", atc.Config{
				Jobs: []atc.JobConfig{
					{
						Name: "some-job",
						Plan: atc.PlanSequence{
							{
								Task: "some-task",
								TaskConfig: &atc.TaskConfig{
									Run: atc.TaskRunConfig{
										Path: "sh",
										Args: []string{"-c", "echo hello from some-job"},
									},
								},
							},
						},
					},
				},
			})
			Expect(err).NotTo(HaveOccurred())

			_, err = client.UnpausePipeline(pipelineName)
			Expect(err).NotTo(HaveOccurred())

			build, err = client.CreateJobBuild(pipelineName, "some-job")
			Expect(err).NotTo(HaveOccurred())
		})

		It("can view the build", func() {
			Expect(page.Navigate(atcRoute(build.URL))).To(Succeed())
			Eventually(page).Should(HaveURL(atcRoute(fmt.Sprintf("pipelines/%s/jobs/some-job/builds/%s", pipelineName, build.Name))))
			Eventually(page.Find("h1")).Should(HaveText(fmt.Sprintf("some-job #%s", build.Name)))
			Eventually(page.Find("#builds")).Should(HaveText(build.Name))

			Eventually(page.Find("#page-header.succeeded")).Should(BeFound())
			Eventually(page.Find(".build-times").Text).Should(ContainSubstring("duration"))

			Expect(page.Find(".build-step .header").Click()).To(Succeed())
			Eventually(page.Find("#build-body").Text).Should(ContainSubstring("hello from some-job"))
		})
	})

	Context("with a one-off build", func() {
		var build atc.Build

		BeforeEach(func() {
			var err error

			pf := atc.NewPlanFactory(0)

			build, err = client.CreateBuild(pf.NewPlan(atc.TaskPlan{
				Name: "some-task",
				Config: &atc.TaskConfig{
					Run: atc.TaskRunConfig{
						Path: "sh",
						Args: []string{"-c", "echo hello from one-off"},
					},
				},
			}))
			Expect(err).NotTo(HaveOccurred())
		})

		It("can view the build", func() {
			Expect(page.Navigate(atcRoute(build.URL))).To(Succeed())
			Eventually(page).Should(HaveURL(atcRoute(fmt.Sprintf("/builds/%d", build.ID))))
			Eventually(page.Find("h1")).Should(HaveText(fmt.Sprintf("build #%d", build.ID)))
			Expect(page.Find("#builds").Text()).Should(BeEmpty())

			Eventually(page.Find("#page-header.succeeded")).Should(BeFound())
			Eventually(page.Find(".build-times").Text).Should(ContainSubstring("duration"))

			Eventually(page.Find(".build-step .header").Text).Should(ContainSubstring("some-task"))
			Expect(page.Find(".build-step .header").Click()).To(Succeed())
			Eventually(page.Find("#build-body").Text).Should(ContainSubstring("hello from one-off"))
		})
	})
})
<commit_msg>account for renamed css class build-duration<commit_after>package web_test

import (
	"fmt"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/sclevine/agouti/matchers"

	"github.com/concourse/atc"
)

var _ = Describe("Viewing builds", func() {
	Context("with a job build", func() {
		var build atc.Build

		BeforeEach(func() {
			_, _, err := client.CreateOrUpdatePipelineConfig(pipelineName, "0", atc.Config{
				Jobs: []atc.JobConfig{
					{
						Name: "some-job",
						Plan: atc.PlanSequence{
							{
								Task: "some-task",
								TaskConfig: &atc.TaskConfig{
									Run: atc.TaskRunConfig{
										Path: "sh",
										Args: []string{"-c", "echo hello from some-job"},
									},
								},
							},
						},
					},
				},
			})
			Expect(err).NotTo(HaveOccurred())

			_, err = client.UnpausePipeline(pipelineName)
			Expect(err).NotTo(HaveOccurred())

			build, err = client.CreateJobBuild(pipelineName, "some-job")
			Expect(err).NotTo(HaveOccurred())
		})

		It("can view the build", func() {
			Expect(page.Navigate(atcRoute(build.URL))).To(Succeed())
			Eventually(page).Should(HaveURL(atcRoute(fmt.Sprintf("pipelines/%s/jobs/some-job/builds/%s", pipelineName, build.Name))))
			Eventually(page.Find("h1")).Should(HaveText(fmt.Sprintf("some-job #%s", build.Name)))
			Eventually(page.Find("#builds")).Should(HaveText(build.Name))

			Eventually(page.Find("#page-header.succeeded")).Should(BeFound())
			Eventually(page.Find(".build-duration").Text).Should(ContainSubstring("duration"))

			Expect(page.Find(".build-step .header").Click()).To(Succeed())
			Eventually(page.Find("#build-body").Text).Should(ContainSubstring("hello from some-job"))
		})
	})

	Context("with a one-off build", func() {
		var build atc.Build

		BeforeEach(func() {
			var err error

			pf := atc.NewPlanFactory(0)

			build, err = client.CreateBuild(pf.NewPlan(atc.TaskPlan{
				Name: "some-task",
				Config: &atc.TaskConfig{
					Run: atc.TaskRunConfig{
						Path: "sh",
						Args: []string{"-c", "echo hello from one-off"},
					},
				},
			}))
			Expect(err).NotTo(HaveOccurred())
		})

		It("can view the build", func() {
			Expect(page.Navigate(atcRoute(build.URL))).To(Succeed())
			Eventually(page).Should(HaveURL(atcRoute(fmt.Sprintf("/builds/%d", build.ID))))
			Eventually(page.Find("h1")).Should(HaveText(fmt.Sprintf("build #%d", build.ID)))
			Expect(page.Find("#builds").Text()).Should(BeEmpty())

			Eventually(page.Find("#page-header.succeeded")).Should(BeFound())
			Eventually(page.Find(".build-duration").Text).Should(ContainSubstring("duration"))

			Eventually(page.Find(".build-step .header").Text).Should(ContainSubstring("some-task"))
			Expect(page.Find(".build-step .header").Click()).To(Succeed())
			Eventually(page.Find("#build-body").Text).Should(ContainSubstring("hello from one-off"))
		})
	})
})
<|endoftext|>


<commit_before>package rrdtool

import (
	"encoding/base64"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/rpc"
	"net/rpc/jsonrpc"
	"sync/atomic"
	"time"

	"stathat.com/c/consistent"

	cmodel "github.com/open-falcon/common/model"
	"github.com/open-falcon/graph/g"
	"github.com/open-falcon/graph/store"
)

type Task_t struct {
	Method string
	Key    string
	Done   chan error
	Args   interface{}
	Reply  interface{}
}

const (
	FETCH_S_SUCCESS = iota
	FETCH_S_ERR
	SEND_S_SUCCESS
	SEND_S_ERR
	QUERY_S_SUCCESS
	QUERY_S_ERR
	CONN_S_ERR
	CONN_S_DIAL
	STAT_SIZE
)

var (
	Consistent       *consistent.Consistent
	Task_ch          map[string]chan *Task_t
	clients          map[string][]*rpc.Client
	flushrrd_timeout int32
	stat_cnt         [STAT_SIZE]uint64
)

func init() {
	Consistent = consistent.New()
	Task_ch = make(map[string]chan *Task_t)
	clients = make(map[string][]*rpc.Client)
}

func GetCounter() (ret string) {
	return fmt.Sprintf("FETCH_S_SUCCESS[%d] FETCH_S_ERR[%d] SEND_S_SUCCESS[%d] SEND_S_ERR[%d] QUERY_S_SUCCESS[%d] QUERY_S_ERR[%d] CONN_S_ERR[%d] CONN_S_DIAL[%d]",
		atomic.LoadUint64(&stat_cnt[FETCH_S_SUCCESS]),
		atomic.LoadUint64(&stat_cnt[FETCH_S_ERR]),
		atomic.LoadUint64(&stat_cnt[SEND_S_SUCCESS]),
		atomic.LoadUint64(&stat_cnt[SEND_S_ERR]),
		atomic.LoadUint64(&stat_cnt[QUERY_S_SUCCESS]),
		atomic.LoadUint64(&stat_cnt[QUERY_S_ERR]),
		atomic.LoadUint64(&stat_cnt[CONN_S_ERR]),
		atomic.LoadUint64(&stat_cnt[CONN_S_DIAL]))
}

func migrate_start(cfg *g.GlobalConfig) {
	var err error
	var i int
	var client *rpc.Client
	if cfg.Migrate.Enabled {
		Consistent.NumberOfReplicas = cfg.Migrate.Replicas

		for node, addr := range cfg.Migrate.Cluster {
			Consistent.Add(node)
			Task_ch[node] = make(chan *Task_t, 1)
			clients[node] = make([]*rpc.Client, cfg.Migrate.Concurrency)

			for i = 0; i < cfg.Migrate.Concurrency; i++ {
				if client, err = jsonrpc.Dial("tcp", addr); err != nil {
					log.Fatalf("node:%s addr:%s err:%s\n", node, addr, err)
				}
				clients[node][i] = client
				go task_worker(i, Task_ch[node], client, node, addr)
			}
		}
	}
}

func task_worker(idx int, ch chan *Task_t, client *rpc.Client, node, addr string) {
	var err error
	for {
		select {
		case task := <-ch:
			if task.Method == "Graph.Send" {
				if err = send_data(client, task.Key, addr); err != nil {
					atomic.AddUint64(&stat_cnt[SEND_S_ERR], 1)
				} else {
					atomic.AddUint64(&stat_cnt[SEND_S_SUCCESS], 1)
				}
			} else if task.Method == "Graph.Query" {
				if err = query_data(client, addr, task.Args, task.Reply); err != nil {
					atomic.AddUint64(&stat_cnt[QUERY_S_ERR], 1)
				} else {
					atomic.AddUint64(&stat_cnt[QUERY_S_SUCCESS], 1)
				}
			} else {
				if atomic.LoadInt32(&flushrrd_timeout) != 0 {
					// hope this more faster than fetch_rrd
					if err = send_data(client, task.Key, addr); err != nil {
						atomic.AddUint64(&stat_cnt[SEND_S_ERR], 1)
					} else {
						atomic.AddUint64(&stat_cnt[SEND_S_SUCCESS], 1)
					}
				} else {
					if err = fetch_rrd(client, task.Key, addr); err != nil {
						atomic.AddUint64(&stat_cnt[FETCH_S_ERR], 1)
					} else {
						atomic.AddUint64(&stat_cnt[FETCH_S_SUCCESS], 1)
					}
				}
			}
			if task.Done != nil {
				task.Done <- err
			}
		}
	}
}

func reconnection(client *rpc.Client, addr string) {
	var err error

	atomic.AddUint64(&stat_cnt[CONN_S_ERR], 1)
	client.Close()

	client, err = jsonrpc.Dial("tcp", addr)
	atomic.AddUint64(&stat_cnt[CONN_S_DIAL], 1)

	for err != nil {
		//danger!! block routine
		time.Sleep(time.Millisecond * 500)
		client, err = jsonrpc.Dial("tcp", addr)
		atomic.AddUint64(&stat_cnt[CONN_S_DIAL], 1)
	}
}

func query_data(client *rpc.Client, addr string,
	args interface{}, resp interface{}) error {
	var (
		err error
	)

	err = jsonrpc_call(client, "Graph.Query", args, resp,
		time.Duration(g.Config().CallTimeout)*time.Millisecond)

	if err != nil {
		reconnection(client, addr)
	}
	return err
}

func send_data(client *rpc.Client, key string, addr string) error {
	var (
		err  error
		flag uint32
		resp *cmodel.SimpleRpcResponse
	)

	//remote
	if flag, err = store.GraphItems.GetFlag(key); err != nil {
		return err
	}
	cfg := g.Config()

	store.GraphItems.SetFlag(key, flag|g.GRAPH_F_SENDING)

	items := store.GraphItems.PopAll(key)
	items_size := len(items)
	if items_size == 0 {
		goto out
	}
	resp = &cmodel.SimpleRpcResponse{}

	err = jsonrpc_call(client, "Graph.Send", items, resp,
		time.Duration(cfg.CallTimeout)*time.Millisecond)

	if err != nil {
		store.GraphItems.PushAll(key, items)
		reconnection(client, addr)
		goto err_out
	}
	goto out

err_out:
	flag |= g.GRAPH_F_ERR
out:
	flag &= ^g.GRAPH_F_SENDING
	store.GraphItems.SetFlag(key, flag)
	return err

}

func fetch_rrd(client *rpc.Client, key string, addr string) error {
	var (
		err      error
		flag     uint32
		md5      string
		dsType   string
		filename string
		step     int
		rrdfile  g.File64
		ctx      []byte
	)

	cfg := g.Config()

	if flag, err = store.GraphItems.GetFlag(key); err != nil {
		return err
	}

	store.GraphItems.SetFlag(key, flag|g.GRAPH_F_FETCHING)

	md5, dsType, step, _ = g.SplitRrdCacheKey(key)
	filename = g.RrdFileName(cfg.RRD.Storage, md5, dsType, step)

	items := store.GraphItems.PopAll(key)
	items_size := len(items)
	if items_size == 0 {
		// impossible
		goto out
	}

	err = jsonrpc_call(client, "Graph.GetRrd", key, &rrdfile,
		time.Duration(cfg.CallTimeout)*time.Millisecond)

	if err != nil {
		store.GraphItems.PushAll(key, items)
		reconnection(client, addr)
		goto err_out
	}

	if ctx, err = base64.StdEncoding.DecodeString(rrdfile.Body64); err != nil {
		store.GraphItems.PushAll(key, items)
		goto err_out
	} else {
		if err = ioutil.WriteFile(filename, ctx, 0644); err != nil {
			store.GraphItems.PushAll(key, items)
			goto err_out
		} else {
			flag &= ^g.GRAPH_F_MISS
			Flush(filename, items)
			goto out
		}
	}
	//noneed
	goto out

err_out:
	flag |= g.GRAPH_F_ERR
out:
	flag &= ^g.GRAPH_F_FETCHING
	store.GraphItems.SetFlag(key, flag)
	return err
}

func jsonrpc_call(client *rpc.Client, method string, args interface{},
	reply interface{}, timeout time.Duration) error {
	done := make(chan *rpc.Call, 1)
	client.Go(method, args, reply, done)
	select {
	case <-time.After(timeout):
		return errors.New("timeout")
	case call := <-done:
		if call.Error == nil {
			return nil
		} else {
			return call.Error
		}
	}
}
<commit_msg>support dial timeout<commit_after>package rrdtool

import (
	"encoding/base64"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/rpc"
	"net/rpc/jsonrpc"
	"sync/atomic"
	"time"

	"stathat.com/c/consistent"

	cmodel "github.com/open-falcon/common/model"
	"github.com/open-falcon/graph/g"
	"github.com/open-falcon/graph/store"
)

type Task_t struct {
	Method string
	Key    string
	Done   chan error
	Args   interface{}
	Reply  interface{}
}

const (
	FETCH_S_SUCCESS = iota
	FETCH_S_ERR
	SEND_S_SUCCESS
	SEND_S_ERR
	QUERY_S_SUCCESS
	QUERY_S_ERR
	CONN_S_ERR
	CONN_S_DIAL
	STAT_SIZE
)

var (
	Consistent       *consistent.Consistent
	Task_ch          map[string]chan *Task_t
	clients          map[string][]*rpc.Client
	flushrrd_timeout int32
	stat_cnt         [STAT_SIZE]uint64
)

func init() {
	Consistent = consistent.New()
	Task_ch = make(map[string]chan *Task_t)
	clients = make(map[string][]*rpc.Client)
}

func GetCounter() (ret string) {
	return fmt.Sprintf("FETCH_S_SUCCESS[%d] FETCH_S_ERR[%d] SEND_S_SUCCESS[%d] SEND_S_ERR[%d] QUERY_S_SUCCESS[%d] QUERY_S_ERR[%d] CONN_S_ERR[%d] CONN_S_DIAL[%d]",
		atomic.LoadUint64(&stat_cnt[FETCH_S_SUCCESS]),
		atomic.LoadUint64(&stat_cnt[FETCH_S_ERR]),
		atomic.LoadUint64(&stat_cnt[SEND_S_SUCCESS]),
		atomic.LoadUint64(&stat_cnt[SEND_S_ERR]),
		atomic.LoadUint64(&stat_cnt[QUERY_S_SUCCESS]),
		atomic.LoadUint64(&stat_cnt[QUERY_S_ERR]),
		atomic.LoadUint64(&stat_cnt[CONN_S_ERR]),
		atomic.LoadUint64(&stat_cnt[CONN_S_DIAL]))
}

func dial(network, address string, timeout time.Duration) (*rpc.Client, error) {
	d := net.Dialer{Timeout: timeout}
	conn, err := d.Dial(network, address)
	if err != nil {
		return nil, err
	}
	return jsonrpc.NewClient(conn), err
}

func migrate_start(cfg *g.GlobalConfig) {
	var err error
	var i int
	if cfg.Migrate.Enabled {
		Consistent.NumberOfReplicas = cfg.Migrate.Replicas

		for node, addr := range cfg.Migrate.Cluster {
			Consistent.Add(node)
			Task_ch[node] = make(chan *Task_t, 1)
			clients[node] = make([]*rpc.Client, cfg.Migrate.Concurrency)

			for i = 0; i < cfg.Migrate.Concurrency; i++ {
				if clients[node][i], err = dial("tcp", addr, time.Second); err != nil {
					log.Fatalf("node:%s addr:%s err:%s\n", node, addr, err)
				}
				go task_worker(i, Task_ch[node], clients[node][i], addr)
			}
		}
	}
}

func task_worker(idx int, ch chan *Task_t, client *rpc.Client, addr string) {
	var err error
	for {
		select {
		case task := <-ch:
			if task.Method == "Graph.Send" {
				if err = send_data(client, task.Key, addr); err != nil {
					atomic.AddUint64(&stat_cnt[SEND_S_ERR], 1)
				} else {
					atomic.AddUint64(&stat_cnt[SEND_S_SUCCESS], 1)
				}
			} else if task.Method == "Graph.Query" {
				if err = query_data(client, addr, task.Args, task.Reply); err != nil {
					atomic.AddUint64(&stat_cnt[QUERY_S_ERR], 1)
				} else {
					atomic.AddUint64(&stat_cnt[QUERY_S_SUCCESS], 1)
				}
			} else {
				if atomic.LoadInt32(&flushrrd_timeout) != 0 {
					// hope this more faster than fetch_rrd
					if err = send_data(client, task.Key, addr); err != nil {
						atomic.AddUint64(&stat_cnt[SEND_S_ERR], 1)
					} else {
						atomic.AddUint64(&stat_cnt[SEND_S_SUCCESS], 1)
					}
				} else {
					if err = fetch_rrd(client, task.Key, addr); err != nil {
						atomic.AddUint64(&stat_cnt[FETCH_S_ERR], 1)
					} else {
						atomic.AddUint64(&stat_cnt[FETCH_S_SUCCESS], 1)
					}
				}
			}
			if task.Done != nil {
				task.Done <- err
			}
		}
	}
}

func reconnection(client *rpc.Client, addr string) {
	var err error

	atomic.AddUint64(&stat_cnt[CONN_S_ERR], 1)
	client.Close()

	client, err = dial("tcp", addr, time.Second)
	atomic.AddUint64(&stat_cnt[CONN_S_DIAL], 1)

	for err != nil {
		//danger!! block routine
		time.Sleep(time.Millisecond * 500)
		client, err = dial("tcp", addr, time.Second)
		atomic.AddUint64(&stat_cnt[CONN_S_DIAL], 1)
	}
}

func query_data(client *rpc.Client, addr string,
	args interface{}, resp interface{}) error {
	var (
		err error
	)

	err = jsonrpc_call(client, "Graph.Query", args, resp,
		time.Duration(g.Config().CallTimeout)*time.Millisecond)

	if err != nil {
		reconnection(client, addr)
	}
	return err
}

func send_data(client *rpc.Client, key string, addr string) error {
	var (
		err  error
		flag uint32
		resp *cmodel.SimpleRpcResponse
	)

	//remote
	if flag, err = store.GraphItems.GetFlag(key); err != nil {
		return err
	}
	cfg := g.Config()

	store.GraphItems.SetFlag(key, flag|g.GRAPH_F_SENDING)

	items := store.GraphItems.PopAll(key)
	items_size := len(items)
	if items_size == 0 {
		goto out
	}
	resp = &cmodel.SimpleRpcResponse{}

	err = jsonrpc_call(client, "Graph.Send", items, resp,
		time.Duration(cfg.CallTimeout)*time.Millisecond)

	if err != nil {
		store.GraphItems.PushAll(key, items)
		reconnection(client, addr)
		goto err_out
	}
	goto out

err_out:
	flag |= g.GRAPH_F_ERR
out:
	flag &= ^g.GRAPH_F_SENDING
	store.GraphItems.SetFlag(key, flag)
	return err

}

func fetch_rrd(client *rpc.Client, key string, addr string) error {
	var (
		err      error
		flag     uint32
		md5      string
		dsType   string
		filename string
		step     int
		rrdfile  g.File64
		ctx      []byte
	)

	cfg := g.Config()

	if flag, err = store.GraphItems.GetFlag(key); err != nil {
		return err
	}

	store.GraphItems.SetFlag(key, flag|g.GRAPH_F_FETCHING)

	md5, dsType, step, _ = g.SplitRrdCacheKey(key)
	filename = g.RrdFileName(cfg.RRD.Storage, md5, dsType, step)

	items := store.GraphItems.PopAll(key)
	items_size := len(items)
	if items_size == 0 {
		// impossible
		goto out
	}

	err = jsonrpc_call(client, "Graph.GetRrd", key, &rrdfile,
		time.Duration(cfg.CallTimeout)*time.Millisecond)

	if err != nil {
		store.GraphItems.PushAll(key, items)
		reconnection(client, addr)
		goto err_out
	}

	if ctx, err = base64.StdEncoding.DecodeString(rrdfile.Body64); err != nil {
		store.GraphItems.PushAll(key, items)
		goto err_out
	} else {
		if err = ioutil.WriteFile(filename, ctx, 0644); err != nil {
			store.GraphItems.PushAll(key, items)
			goto err_out
		} else {
			flag &= ^g.GRAPH_F_MISS
			Flush(filename, items)
			goto out
		}
	}
	//noneed
	goto out

err_out:
	flag |= g.GRAPH_F_ERR
out:
	flag &= ^g.GRAPH_F_FETCHING
	store.GraphItems.SetFlag(key, flag)
	return err
}

func jsonrpc_call(client *rpc.Client, method string, args interface{},
	reply interface{}, timeout time.Duration) error {
	done := make(chan *rpc.Call, 1)
	client.Go(method, args, reply, done)
	select {
	case <-time.After(timeout):
		return errors.New("i/o timeout[jsonrpc]")
	case call := <-done:
		if call.Error == nil {
			return nil
		} else {
			return call.Error
		}
	}
}
<|endoftext|>


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

package safehttp

import (
	"errors"
	"net/http"
	"net/textproto"
)

// Header represents the key-value pairs in an HTTP header.
// The keys will be in canonical form, as returned by
// textproto.CanonicalMIMEHeaderKey.
type Header struct {
	wrapped   http.Header
	immutable map[string]bool
}

func newHeader(h http.Header) Header {
	return Header{wrapped: h, immutable: map[string]bool{}}
}

// MarkImmutable marks the header with the given name as immutable.
// The name is first canonicalized using textproto.CanonicalMIMEHeaderKey.
// This header is now read-only. If the header previously was
// undefined, then it will forever be undefined. Other methods in
// the struct can't write to, change or delete the header with this
// name. These methods will instead fail when applied on an immutable
// header.
func (h Header) MarkImmutable(name string) {
	name = textproto.CanonicalMIMEHeaderKey(name)
	h.immutable[name] = true
}

// Set sets the header with the given name to the given value.
// The name is first canonicalized using textproto.CanonicalMIMEHeaderKey.
// This method first removes all other values associated with this
// header before setting the new value. Returns an error when
// applied on immutable headers or on the Set-Cookie header.
func (h Header) Set(name, value string) error {
	name = textproto.CanonicalMIMEHeaderKey(name)
	if err := h.writableHeader(name); err != nil {
		return err
	}
	h.wrapped.Set(name, value)
	return nil
}

// Add adds a new header with the given name and the given value to
// the collection of headers. The name is first canonicalized using
// textproto.CanonicalMIMEHeaderKey. Returns an error when applied
// on immutable headers or on the Set-Cookie header.
func (h Header) Add(name, value string) error {
	name = textproto.CanonicalMIMEHeaderKey(name)
	if err := h.writableHeader(name); err != nil {
		return err
	}
	h.wrapped.Add(name, value)
	return nil
}

// Del deletes all headers with the given name. The name is first
// canonicalized using textproto.CanonicalMIMEHeaderKey. Returns an
// error when applied on immutable headers or on the Set-Cookie
// header.
func (h Header) Del(name string) error {
	name = textproto.CanonicalMIMEHeaderKey(name)
	if err := h.writableHeader(name); err != nil {
		return err
	}
	h.wrapped.Del(name)
	return nil
}

// Get returns the value of the first header with the given name.
// The name is first canonicalized using textproto.CanonicalMIMEHeaderKey.
// If no header exists with the given name then "" is returned.
func (h Header) Get(name string) string {
	return h.wrapped.Get(name)
}

// Values returns all the values of all the headers with the given name.
// The name is first canonicalized using textproto.CanonicalMIMEHeaderKey.
// The values are returned in the same order as they were sent in the request.
// The values are returned as a copy of the original slice of strings in
// the internal header map. This is to prevent modification of the original
// slice. If no header exists with the given name then an empty slice is
// returned.
func (h Header) Values(name string) []string {
	v := h.wrapped.Values(name)
	clone := make([]string, len(v))
	copy(clone, v)
	return clone
}

// SetCookie adds the cookie provided as a Set-Cookie header in the header
// collection. If the cookie is nil or cookie.Name is invalid, no header is
// added. This is the only method that can modify the Set-Cookie header.
// If other methods try to modify the header they will return errors.
// TODO: Replace http.Cookie with safehttp.Cookie.
func (h Header) SetCookie(cookie *http.Cookie) {
	if v := cookie.String(); v != "" {
		h.wrapped.Add("Set-Cookie", v)
	}
}

// TODO: Add Write, WriteSubset and Clone when needed.

// writableHeader assumes that the given name already has been canonicalized
// using textproto.CanonicalMIMEHeaderKey.
func (h Header) writableHeader(name string) error {
	if name == "Set-Cookie" {
		return errors.New("can't write to Set-Cookie header")
	}
	if h.immutable[name] {
		return errors.New("immutable header")
	}
	return nil
}
<commit_msg>Added TODO about legacy conversions<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 	https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package safehttp

import (
	"errors"
	"net/http"
	"net/textproto"
)

// Header represents the key-value pairs in an HTTP header.
// The keys will be in canonical form, as returned by
// textproto.CanonicalMIMEHeaderKey.
type Header struct {
	wrapped   http.Header
	immutable map[string]bool
}

func newHeader(h http.Header) Header {
	return Header{wrapped: h, immutable: map[string]bool{}}
}

// MarkImmutable marks the header with the given name as immutable.
// The name is first canonicalized using textproto.CanonicalMIMEHeaderKey.
// This header is now read-only. If the header previously was
// undefined, then it will forever be undefined. Other methods in
// the struct can't write to, change or delete the header with this
// name. These methods will instead fail when applied on an immutable
// header.
func (h Header) MarkImmutable(name string) {
	name = textproto.CanonicalMIMEHeaderKey(name)
	h.immutable[name] = true
}

// Set sets the header with the given name to the given value.
// The name is first canonicalized using textproto.CanonicalMIMEHeaderKey.
// This method first removes all other values associated with this
// header before setting the new value. Returns an error when
// applied on immutable headers or on the Set-Cookie header.
func (h Header) Set(name, value string) error {
	name = textproto.CanonicalMIMEHeaderKey(name)
	if err := h.writableHeader(name); err != nil {
		return err
	}
	h.wrapped.Set(name, value)
	return nil
}

// Add adds a new header with the given name and the given value to
// the collection of headers. The name is first canonicalized using
// textproto.CanonicalMIMEHeaderKey. Returns an error when applied
// on immutable headers or on the Set-Cookie header.
func (h Header) Add(name, value string) error {
	name = textproto.CanonicalMIMEHeaderKey(name)
	if err := h.writableHeader(name); err != nil {
		return err
	}
	h.wrapped.Add(name, value)
	return nil
}

// Del deletes all headers with the given name. The name is first
// canonicalized using textproto.CanonicalMIMEHeaderKey. Returns an
// error when applied on immutable headers or on the Set-Cookie
// header.
func (h Header) Del(name string) error {
	name = textproto.CanonicalMIMEHeaderKey(name)
	if err := h.writableHeader(name); err != nil {
		return err
	}
	h.wrapped.Del(name)
	return nil
}

// Get returns the value of the first header with the given name.
// The name is first canonicalized using textproto.CanonicalMIMEHeaderKey.
// If no header exists with the given name then "" is returned.
func (h Header) Get(name string) string {
	return h.wrapped.Get(name)
}

// Values returns all the values of all the headers with the given name.
// The name is first canonicalized using textproto.CanonicalMIMEHeaderKey.
// The values are returned in the same order as they were sent in the request.
// The values are returned as a copy of the original slice of strings in
// the internal header map. This is to prevent modification of the original
// slice. If no header exists with the given name then an empty slice is
// returned.
func (h Header) Values(name string) []string {
	v := h.wrapped.Values(name)
	clone := make([]string, len(v))
	copy(clone, v)
	return clone
}

// SetCookie adds the cookie provided as a Set-Cookie header in the header
// collection. If the cookie is nil or cookie.Name is invalid, no header is
// added. This is the only method that can modify the Set-Cookie header.
// If other methods try to modify the header they will return errors.
// TODO: Replace http.Cookie with safehttp.Cookie.
func (h Header) SetCookie(cookie *http.Cookie) {
	if v := cookie.String(); v != "" {
		h.wrapped.Add("Set-Cookie", v)
	}
}

// TODO: Add Write, WriteSubset and Clone when needed.

// writableHeader assumes that the given name already has been canonicalized
// using textproto.CanonicalMIMEHeaderKey.
func (h Header) writableHeader(name string) error {
	// TODO(@mattiasgrenfeldt, @kele, @empijei): Think about how this should
	// work during legacy conversions.
	if name == "Set-Cookie" {
		return errors.New("can't write to Set-Cookie header")
	}
	if h.immutable[name] {
		return errors.New("immutable header")
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path"
	"strings"

	log "github.com/Sirupsen/logrus"
)

type (
	Config struct {
		Username   string
		Password   string
		Email      string
		Registry   string
		Folder     string
		AlwaysAuth bool
		SkipVerify bool
	}

	NpmPackage struct {
		Name    string `json:"name"`
		Version string `json:"version"`
	}

	Plugin struct {
		Config Config
	}
)

const GlobalRegistry = "https://registry.npmjs.org"

func (p Plugin) Exec() error {
	// read the package
	packagePath := path.Join(p.Config.Folder, "package.json")

	npmPackage, err := readPackageFile(packagePath)

	if err != nil {
		log.WithFields(log.Fields{
			"error": err,
		}).Error("Could not read package.json")
		return err
	}

	log.WithFields(log.Fields{
		"name":    npmPackage.Name,
		"version": npmPackage.Version,
	}).Info("Found package")

	// see if the package should be published
	publish, err := shouldPublishPackage(p.Config, npmPackage)

	if publish {
		log.Info("Attempting to publish package")

		// write the npmrc file
		err := writeNpmrcFile(p.Config)

		if err != nil {
			log.WithFields(log.Fields{
				"error": err,
			}).Error("Problem creating npmrc file")

			return err
		}

		var cmds []*exec.Cmd

		// write the version command
		cmds = append(cmds, versionCommand())

		// write registry command
		if p.Config.Registry != GlobalRegistry {
			cmds = append(cmds, registryCommand(p.Config.Registry))
		}

		// write auth command
		if p.Config.AlwaysAuth {
			cmds = append(cmds, alwaysAuthCommand())
		}

		// write skip verify command
		if p.Config.SkipVerify {
			cmds = append(cmds, skipVerifyCommand())
		}

		// write the publish command
		cmds = append(cmds, publishCommand())

		// run the commands
		err = runCommands(cmds)

		if err != nil {
			log.WithFields(log.Fields{
				"error": err,
			}).Error("Package was not published")
		}
	} else {
		log.WithFields(log.Fields{
			"reason": err,
		}).Info("Package will not be published")
	}

	return nil
}

/// Reads the package file at the given path
func readPackageFile(path string) (*NpmPackage, error) {
	// read the file
	file, err := ioutil.ReadFile(path)

	if err != nil {
		return nil, err
	}

	// unmarshal the json data
	npmPackage := NpmPackage{}
	err = json.Unmarshal(file, &npmPackage)

	if err != nil {
		return nil, err
	}

	// make sure values are present
	if len(npmPackage.Name) == 0 {
		return nil, errors.New("No package name present")
	}

	if len(npmPackage.Version) == 0 {
		return nil, errors.New("No package version present")
	}

	return &npmPackage, nil
}

/// Determines if the package should be published
func shouldPublishPackage(config Config, npmPackage *NpmPackage) (bool, error) {
	// get the url for the package
	//
	// encoding the portion of the url for the package name in case of scopes
	packageUrl := fmt.Sprintf("%s/%s/%s", config.Registry, url.QueryEscape(npmPackage.Name), npmPackage.Version)

	log.WithFields(log.Fields{
		"url": packageUrl,
	}).Info("Requesting package information")

	// create a request for the package
	req, err := http.NewRequest("GET", packageUrl, nil)

	if err != nil {
		return false, err
	}

	// set authentication if needed
	// NPM uses basic http auth
	if config.AlwaysAuth {
		req.SetBasicAuth(config.Username, config.Password)
	}

	// skip verify if necessary
	if config.SkipVerify {
		http.DefaultTransport = &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}

		log.Warning("Skipping SSL verification")
	}

	// get the response
	resp, err := http.DefaultClient.Do(req)

	if err != nil {
		return false, err
	}

	defer resp.Body.Close()

	statusCode := resp.StatusCode

	// look for a 404 to see if the package should be published
	if statusCode == http.StatusNotFound {
		return true, nil
	} else if statusCode == http.StatusOK {
		return false, errors.New("Package already published")
	} else {
		contents, err := ioutil.ReadAll(resp.Body)

		if err != nil {
			return false, err
		}

		return false, errors.New(fmt.Sprintf("Error Occurred. Status %d\nBody:%s", statusCode, contents))
	}
}

/// Writes the npmrc file
func writeNpmrcFile(config Config) error {
	// get the base64 encoded string
	authString := fmt.Sprintf("%s:%s", config.Username, config.Password)
	encoded := base64.StdEncoding.EncodeToString([]byte(authString))

	// create the file contents
	contents := fmt.Sprintf("_auth = %s\nemail = %s", encoded, config.Email)

	// write the file
	return ioutil.WriteFile("/root/.npmrc", []byte(contents), 0644)
}

// Gets the npm version
func versionCommand() *exec.Cmd {
	return exec.Command("npm", "--version")
}

// Sets the npm registry
func registryCommand(registry string) *exec.Cmd {
	return exec.Command("npm", "config", "set", "registry", registry)
}

// Sets the always off flag
func alwaysAuthCommand() *exec.Cmd {
	return exec.Command("npm", "config", "set", "always-auth", "true")
}

// Skip ssl verification
func skipVerifyCommand() *exec.Cmd {
	return exec.Command("npm", "config", "set", "strict-ssl", "false")
}

// Publishes the package
func publishCommand() *exec.Cmd {
	return exec.Command("npm", "publish")
}

// Trace writes each command to standard error (preceded by a ‘$ ’) before it
// is executed. Used for debugging your build.
func trace(cmd *exec.Cmd) {
	fmt.Fprintf(os.Stdout, "+ %s\n", strings.Join(cmd.Args, " "))
}

func runCommands(cmds []*exec.Cmd) error {
	for _, cmd := range cmds {
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		trace(cmd)

		err := cmd.Run()
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>Adding some sanity checking for config<commit_after>package main

import (
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path"
	"strings"

	log "github.com/Sirupsen/logrus"
)

type (
	Config struct {
		Username   string
		Password   string
		Email      string
		Registry   string
		Folder     string
		AlwaysAuth bool
    SkipVerify bool
	}

	NpmPackage struct {
		Name    string `json:"name"`
		Version string `json:"version"`
	}

	Plugin struct {
		Config Config
	}
)

const GlobalRegistry = "https://registry.npmjs.org"

func (p Plugin) Exec() error {
	// check for a username
	if len(p.Config.Username) == 0 {
		log.Error("No username provided")
		return errors.New("No username provided")
	}

	// check for an email
	if len(p.Config.Email) == 0 {
		log.Error("No email address provided")
		return errors.New("No email address provided")
	}

	// check for a password
	if len(p.Config.Password) == 0 {
		log.Warning("No password provided")
	}

	log.WithFields(log.Fields{
		"username": p.Config.Username,
		"email": p.Config.Email,
	}).Info("Specified credentials")

	// read the package
	packagePath := path.Join(p.Config.Folder, "package.json")

	npmPackage, err := readPackageFile(packagePath)

	if err != nil {
		log.WithFields(log.Fields{
			"error": err,
		}).Error("Could not read package.json")
		return err
	}

	log.WithFields(log.Fields{
		"name":    npmPackage.Name,
		"version": npmPackage.Version,
	}).Info("Found package")

	// see if the package should be published
	publish, err := shouldPublishPackage(p.Config, npmPackage)

	if publish {
		log.Info("Attempting to publish package")

		// write the npmrc file
		err := writeNpmrcFile(p.Config)

		if err != nil {
			log.WithFields(log.Fields{
				"error": err,
			}).Error("Problem creating npmrc file")

			return err
		}

		var cmds []*exec.Cmd

    // write the version command
    cmds = append(cmds, versionCommand())

		// write registry command
		if p.Config.Registry != GlobalRegistry {
			cmds = append(cmds, registryCommand(p.Config.Registry))
		}

		// write auth command
		if p.Config.AlwaysAuth {
			cmds = append(cmds, alwaysAuthCommand())
		}

    // write skip verify command
    if p.Config.SkipVerify {
      cmds = append(cmds, skipVerifyCommand())
    }

		// write the publish command
		cmds = append(cmds, publishCommand())

		// run the commands
		err = runCommands(cmds)

		if err != nil {
			log.WithFields(log.Fields{
				"error": err,
			}).Error("Package was not published")
		}
	} else {
		log.WithFields(log.Fields{
			"reason": err,
		}).Info("Package will not be published")
	}

	return nil
}

/// Reads the package file at the given path
func readPackageFile(path string) (*NpmPackage, error) {
	// read the file
	file, err := ioutil.ReadFile(path)

	if err != nil {
		return nil, err
	}

	// unmarshal the json data
	npmPackage := NpmPackage{}
	err = json.Unmarshal(file, &npmPackage)

	if err != nil {
		return nil, err
	}

	// make sure values are present
	if len(npmPackage.Name) == 0 {
		return nil, errors.New("No package name present")
	}

	if len(npmPackage.Version) == 0 {
		return nil, errors.New("No package version present")
	}

	return &npmPackage, nil
}

/// Determines if the package should be published
func shouldPublishPackage(config Config, npmPackage *NpmPackage) (bool, error) {
	// get the url for the package
	//
	// encoding the portion of the url for the package name in case of scopes
	packageUrl := fmt.Sprintf("%s/%s/%s", config.Registry, url.QueryEscape(npmPackage.Name), npmPackage.Version)

	log.WithFields(log.Fields{
		"url": packageUrl,
	}).Info("Requesting package information")

	// create a request for the package
	req, err := http.NewRequest("GET", packageUrl, nil)

	if err != nil {
		return false, err
	}

	// set authentication if needed
	// NPM uses basic http auth
	if config.AlwaysAuth {
		req.SetBasicAuth(config.Username, config.Password)
	}

  // skip verify if necessary
  if config.SkipVerify {
    http.DefaultTransport = &http.Transport{
      TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
    }

    log.Warning("Skipping SSL verification")
  }

	// get the response
	resp, err := http.DefaultClient.Do(req)

	if err != nil {
		return false, err
	}

	defer resp.Body.Close()

	statusCode := resp.StatusCode

	// look for a 404 to see if the package should be published
	if statusCode == http.StatusNotFound {
		return true, nil
	} else if statusCode == http.StatusOK {
		return false, errors.New("Package already published")
	} else {
		contents, err := ioutil.ReadAll(resp.Body)

		if err != nil {
			return false, err
		}

		return false, errors.New(fmt.Sprintf("Error Occurred. Status %d\nBody:%s", statusCode, contents))
	}
}

/// Writes the npmrc file
func writeNpmrcFile(config Config) error {
	// get the base64 encoded string
	authString := fmt.Sprintf("%s:%s", config.Username, config.Password)
	encoded := base64.StdEncoding.EncodeToString([]byte(authString))

	// create the file contents
	contents := fmt.Sprintf("_auth = %s\nemail = %s", encoded, config.Email)

	// write the file
	return ioutil.WriteFile("/root/.npmrc", []byte(contents), 0644)
}

// Gets the npm version
func versionCommand() *exec.Cmd {
  return exec.Command("npm", "--version")
}

// Sets the npm registry
func registryCommand(registry string) *exec.Cmd {
	return exec.Command("npm", "config", "set", "registry", registry)
}

// Sets the always off flag
func alwaysAuthCommand() *exec.Cmd {
	return exec.Command("npm", "config", "set", "always-auth", "true")
}

// Skip ssl verification
func skipVerifyCommand() *exec.Cmd {
  return exec.Command("npm", "config", "set", "ca=\"\"")
}

// Publishes the package
func publishCommand() *exec.Cmd {
	return exec.Command("npm", "publish")
}

// Trace writes each command to standard error (preceded by a ‘$ ’) before it
// is executed. Used for debugging your build.
func trace(cmd *exec.Cmd) {
	fmt.Fprintf(os.Stdout, "+ %s\n", strings.Join(cmd.Args, " "))
}

func runCommands(cmds []*exec.Cmd) error {
	for _, cmd := range cmds {
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		trace(cmd)

		err := cmd.Run()
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strings"

	"github.com/bluele/slack"
)

type (
	Repo struct {
		Owner string
		Name  string
	}

	Build struct {
		Tag     string
		Event   string
		Number  int
		Commit  string
		Ref     string
		Branch  string
		Author  string
		Message string
		Status  string
		Link    string
		Started int64
		Created int64
	}

	Config struct {
		Webhook   string
		Channel   string
		Recipient string
		Username  string
		Template  string
		ImageURL  string
		IconURL   string
		IconEmoji string
	}

	Job struct {
		Started int64
	}

	Plugin struct {
		Repo   Repo
		Build  Build
		Config Config
		Job    Job
	}
)

func (p Plugin) Exec() error {
	attachment := slack.Attachment{
		Text:       message(p.Repo, p.Build),
		Fallback:   fallback(p.Repo, p.Build),
		Color:      color(p.Build),
		MarkdownIn: []string{"text", "fallback"},
		ImageURL:   p.Config.ImageURL,
	}

	payload := slack.WebHookPostPayload{}
	payload.Username = p.Config.Username
	payload.Attachments = []*slack.Attachment{&attachment}
	payload.IconUrl = p.Config.IconURL
	payload.IconEmoji = p.Config.IconEmoji

	if p.Config.Recipient != "" {
		payload.Channel = prepend("@", p.Config.Recipient)
	} else if p.Config.Channel != "" {
		payload.Channel = prepend("#", p.Config.Channel)
	}

	if p.Config.Template != "" {
		txt, err := RenderTrim(p.Config.Template, p)
		if err != nil {
			return err
		}
		attachment.Text = txt
	}

	client := slack.NewWebHook(p.Config.Webhook)
	return client.PostMessage(&payload)
}

func message(repo Repo, build Build) string {
	return fmt.Sprintf("*%s* <%s|%s/%s#%s> (%s) by %s",
		build.Status,
		build.Link,
		repo.Owner,
		repo.Name,
		build.Commit[:8],
		build.Branch,
		build.Author,
		build.Message,
	)
}

func fallback(repo Repo, build Build) string {
	return fmt.Sprintf("%s %s/%s#%s (%s) by %s",
		build.Status,
		repo.Owner,
		repo.Name,
		build.Commit[:8],
		build.Branch,
		build.Author,
		build.Message,
	)
}

func color(build Build) string {
	switch build.Status {
	case "success":
		return "good"
	case "failure", "error", "killed":
		return "danger"
	default:
		return "warning"
	}
}

func prepend(prefix, s string) string {
	if !strings.HasPrefix(s, prefix) {
		return prefix + s
	}
	return s
}
<commit_msg>Added commit.message to build vars<commit_after>package main

import (
	"fmt"
	"strings"

	"github.com/bluele/slack"
)

type (
	Repo struct {
		Owner string
		Name  string
	}

	Build struct {
		Tag     string
		Event   string
		Number  int
		Commit  string
		Ref     string
		Branch  string
		Author  string
		Message string
		Status  string
		Link    string
		Started int64
		Created int64
	}

	Config struct {
		Webhook   string
		Channel   string
		Recipient string
		Username  string
		Template  string
		ImageURL  string
		IconURL   string
		IconEmoji string
	}

	Job struct {
		Started int64
	}

	Plugin struct {
		Repo   Repo
		Build  Build
		Config Config
		Job    Job
	}
)

func (p Plugin) Exec() error {
	attachment := slack.Attachment{
		Text:       message(p.Repo, p.Build),
		Fallback:   fallback(p.Repo, p.Build),
		Color:      color(p.Build),
		MarkdownIn: []string{"text", "fallback"},
		ImageURL:   p.Config.ImageURL,
	}

	payload := slack.WebHookPostPayload{}
	payload.Username = p.Config.Username
	payload.Attachments = []*slack.Attachment{&attachment}
	payload.IconUrl = p.Config.IconURL
	payload.IconEmoji = p.Config.IconEmoji

	if p.Config.Recipient != "" {
		payload.Channel = prepend("@", p.Config.Recipient)
	} else if p.Config.Channel != "" {
		payload.Channel = prepend("#", p.Config.Channel)
	}

	if p.Config.Template != "" {
		txt, err := RenderTrim(p.Config.Template, p)
		if err != nil {
			return err
		}
		attachment.Text = txt
	}

	client := slack.NewWebHook(p.Config.Webhook)
	return client.PostMessage(&payload)
}

func message(repo Repo, build Build) string {
	return fmt.Sprintf("*%s* <%s|%s/%s#%s> (%s) by %s. Changes: %s",
		build.Status,
		build.Link,
		repo.Owner,
		repo.Name,
		build.Commit[:8],
		build.Branch,
		build.Author,
		build.Message,
	)
}

func fallback(repo Repo, build Build) string {
	return fmt.Sprintf("%s %s/%s#%s (%s) by %s. Changes: %s",
		build.Status,
		repo.Owner,
		repo.Name,
		build.Commit[:8],
		build.Branch,
		build.Author,
		build.Message,
	)
}

func color(build Build) string {
	switch build.Status {
	case "success":
		return "good"
	case "failure", "error", "killed":
		return "danger"
	default:
		return "warning"
	}
}

func prepend(prefix, s string) string {
	if !strings.HasPrefix(s, prefix) {
		return prefix + s
	}
	return s
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
)

// Plugin defines the Helm plugin parameters.
type Plugin struct {
	Debug        bool     `envconfig:"DEBUG"`
	ShowEnv      bool     `envconfig:"SHOW_ENV"`
	Wait         bool     `envconfig:"WAIT"`
	Recreate     bool     `envconfig:"RECREATE_PODS" default:"false"`
	WaitTimeout  uint32   `envconfig:"WAIT_TIMEOUT" default:"300"`
	Actions      []string `envconfig:"ACTIONS" required:"true"`
	AuthKey      string   `envconfig:"AUTH_KEY"`
	KeyPath      string   `envconfig:"KEY_PATH"`
	Zone         string   `envconfig:"ZONE"`
	Cluster      string   `envconfig:"CLUSTER"`
	Project      string   `envconfig:"PROJECT"`
	Namespace    string   `envconfig:"NAMESPACE"`
	ChartRepo    string   `envconfig:"CHART_REPO"`
	Bucket       string   `envconfig:"BUCKET"`
	ChartPath    string   `envconfig:"CHART_PATH" required:"true"`
	ChartVersion string   `envconfig:"CHART_VERSION"`
	Release      string   `envconfig:"RELEASE"`
	Package      string   `envconfig:"PACKAGE"`
	Values       []string `envconfig:"VALUES"`
	ValueFiles   []string `envconfig:"VALUE_FILES"`
}

const (
	gcloudBin  = "/opt/google-cloud-sdk/bin/gcloud"
	gsutilBin  = "/opt/google-cloud-sdk/bin/gsutil"
	kubectlBin = "/opt/google-cloud-sdk/bin/kubectl"
	helmBin    = "/opt/google-cloud-sdk/bin/helm"

	lintPkg   = "lint"
	createPkg = "create"
	pushPkg   = "push"
	pullPkg   = "pull"
	deployPkg = "deploy"
	testPkg   = "test"
)

var reVersions = regexp.MustCompile(`(?P<realm>Client|Server): &version.Version.SemVer:"(?P<semver>.*?)".*?GitCommit:"(?P<commit>.*?)".*?GitTreeState:"(?P<treestate>.*?)"`)

// Exec executes the plugin step.
func (p Plugin) Exec() error {
	// only setup project when needed args are provided
	if p.Project != "" && p.Cluster != "" && p.Zone != "" {
		if err := setupProject(p.Project, p.Cluster, p.Zone, p.Debug); err != nil {
			return err
		}
		if err := helmInit(p.Debug); err != nil {
			return err
		}
	}

	for _, a := range p.Actions {
		switch a {
		case lintPkg:
			if err := p.lintPackage(); err != nil {
				return err
			}
		case createPkg:
			if err := p.createPackage(); err != nil {
				return err
			}
		case pushPkg:
			if err := p.pushPackage(); err != nil {
				return err
			}
		case pullPkg:
			if err := p.pullPackage(); err != nil {
				return err
			}
		case deployPkg:
			if err := p.deployPackage(); err != nil {
				return err
			}
		case testPkg:
			if err := p.testPackage(); err != nil {
				return err
			}
		default:
			return errors.New("unknown action")
		}
	}

	return nil
}

// setupProject setups gcloud project.
func setupProject(project, cluster, zone string, debug bool) error {
	// project configuration
	cmd := exec.Command(gcloudBin, "config", "set", "project", project)
	if err := run(cmd, debug); err != nil {
		return fmt.Errorf("could not the configure the project with glcoud: %v", err)
	}

	// cluster configuration
	cmd = exec.Command(gcloudBin, "container", "clusters", "get-credentials", cluster, "--zone", zone)
	if err := run(cmd, debug); err != nil {
		return fmt.Errorf("could not configure the cluster with glcoud: %v", err)
	}

	return nil
}

// setupAuth configures gcloud to use the given authFile
func setupAuth(authFile string, debug bool) error {
	if err := os.Setenv("GOOGLE_APPLICATION_CREDENTIALS", authFile); err != nil {
		return fmt.Errorf("could not set GOOGLE_APPLICATION_CREDENTIALS env variable: %v", err)
	}

	// authorization
	cmd := exec.Command(gcloudBin, "auth", "activate-service-account", fmt.Sprintf("--key-file=%s", authFile))
	if err := run(cmd, debug); err != nil {
		return fmt.Errorf("could not authorize with glcoud: %v", err)
	}
	return nil
}

// createPackage creates Helm package for Kubernetes.
// helm package --version $PLUGIN_CHART_VERSION $PLUGIN_CHART_PATH
func (p Plugin) createPackage() error {
	return run(exec.Command(helmBin, "package", "--version", p.ChartVersion, p.ChartPath), p.Debug)
}

// cpPackage copies a file from SOURCE to DEST
// gsutil cp SOURCE DEST
func (p Plugin) cpPackage(source string, dest string) error {
	return run(exec.Command(gsutilBin, "cp", source, dest), p.Debug)
}

// cpPackage pulls helm chart from Google Storage to local
// gsutil cp $PACKAGE-$PLUGIN_CHART_VERSION.tgz gs://$PLUGIN_BUCKET
func (p Plugin) pullPackage() error {
	return p.cpPackage(
		fmt.Sprintf("gs://%s/%s-%s.tgz", p.Bucket, p.Package, p.ChartVersion),
		fmt.Sprintf("%s-%s.tgz", p.Package, p.ChartVersion),
	)
}

// pushPackage pushes Helm package to the Google Storage.
// gsutil cp $PACKAGE-$PLUGIN_CHART_VERSION.tgz gs://$PLUGIN_BUCKET
func (p Plugin) pushPackage() error {
	return p.cpPackage(
		fmt.Sprintf("%s-%s.tgz", p.Package, p.ChartVersion),
		fmt.Sprintf("gs://%s", p.Bucket),
	)
}

// helm lint $CHARTPATH -i
func (p Plugin) lintPackage() error {
	return run(exec.Command(helmBin, "lint", p.ChartPath), p.Debug)
}

// helm upgrade $PACKAGE $PACKAGE-$PLUGIN_CHART_VERSION.tgz -i
func (p Plugin) deployPackage() error {
	args := []string{
		helmBin,
		"upgrade",
		p.Release,
		fmt.Sprintf("%s-%s.tgz", p.Package, p.ChartVersion),
	}
	if len(p.ValueFiles) > 0 {
		for _, f := range p.ValueFiles {
			args = append(args, "-f", f)
		}
	}
	if len(p.Values) > 0 {
		args = append(args, "--set", strings.Join(p.Values, ","))
	}
	if p.Recreate {
		args = append(args, "--recreate-pods")
	}
	args = append(args, "--install")
	args = append(args, "--namespace", p.Namespace)

	if p.Wait {
		args = append(args, "--wait", "--timeout", strconv.Itoa(int(p.WaitTimeout)))
	}
	return run(exec.Command("/bin/sh", "-c", strings.Join(args, " ")), p.Debug)
}

// helm test $PACKAGE
func (p Plugin) testPackage() error {
	args := []string{helmBin, "test", p.Release, "--cleanup", "--timeout", strconv.Itoa(int(p.WaitTimeout))}
	return run(exec.Command("/bin/sh", "-c", strings.Join(args, " ")), p.Debug)
}

// fetchHelmVersions returns helm and tiller versions as map
func fetchHelmVersions(debug bool) (map[string]map[string]string, error) {
	var out bytes.Buffer
	var stderr bytes.Buffer
	cmd := exec.Command(helmBin, "version")
	cmd.Stdout = &out
	cmd.Stderr = &stderr

	if debug {
		log.Printf("running: %s", strings.Join(cmd.Args, " "))
	}
	if err := cmd.Run(); err != nil {
		return nil, errors.New(stderr.String())
	}
	if debug {
		log.Printf("%s", out.String())
	}

	lines := strings.Split(out.String(), "\n")
	versions := make(map[string]map[string]string)

	// we just care about the first two lines
	for _, line := range lines[:2] {
		entry, reErr := scanNamed(line, reVersions)
		if reErr != nil {
			return nil, reErr
		}
		versions[strings.ToLower(entry["realm"])] = entry
	}

	return versions, nil
}

// helmInit inits Triller on Kubernetes cluster.
func helmInit(debug bool) error {
	args := []string{"init"}

	ver, err := fetchHelmVersions(debug)
	if err == nil {
		switch strings.Compare(ver["client"]["semver"], ver["server"]["semver"]) {
		case -1: // client is older than tiller
			return fmt.Errorf("helm client is out of date")
		case 1: // client is newer than tiller
			args = append(args, "--upgrade")
		default: // client and tiller are at the same version
			args = append(args, "--client-only")
		}
	}

	if err := run(exec.Command(helmBin, args...), debug); err != nil {
		return err
	}

	// poll for tiller (call helm version 10 times)
	return pollTiller(10, debug)
}

// pollTiller repeatedly calls helm version and checks its exit code
func pollTiller(retries int, debug bool) error {
	var err error
	for i := 0; i < retries; i++ {
		if err = run(exec.Command(helmBin, "version"), debug); err == nil {
			return nil
		}
	}
	return err
}

func (p Plugin) movePkg() error {
	if err := os.Mkdir(p.Bucket, os.ModeDir); err != nil {
		return err
	}
	return cp(
		fmt.Sprintf("%s-%s.tgz", p.Package, p.ChartVersion),
		fmt.Sprintf("%s/%s-%s.tgz", p.Bucket, p.Package, p.ChartVersion),
	)
}

// cp copies file
func cp(src, dst string) error {
	s, err := os.Open(src)
	if err != nil {
		return err
	}
	// no need to check errors on read only file, we already got everything
	// we need from the filesystem, so nothing can go wrong now.
	defer s.Close()
	d, err := os.Create(dst)
	if err != nil {
		return err
	}
	if _, err := io.Copy(d, s); err != nil {
		d.Close()
		return err
	}
	return d.Close()
}

// scanNamed maps named regex groups to a golang map
func scanNamed(str string, rg *regexp.Regexp) (map[string]string, error) {
	result := make(map[string]string)
	for _, match := range rg.FindAllStringSubmatch(str, -1) {
		for i, name := range rg.SubexpNames() {
			if i != 0 && match[i] != "" && name != "" {
				result[name] = match[i]
			}
		}
	}

	if len(result) == 0 {
		return nil, errors.New("emtpy resultset")
	}

	return result, nil
}

func run(cmd *exec.Cmd, debug bool) error {
	if debug {
		log.Printf("running: %s", strings.Join(cmd.Args, " "))
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
	}
	return cmd.Run()
}
<commit_msg>add namespace argument to test action<commit_after>package main

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
)

// Plugin defines the Helm plugin parameters.
type Plugin struct {
	Debug        bool     `envconfig:"DEBUG"`
	ShowEnv      bool     `envconfig:"SHOW_ENV"`
	Wait         bool     `envconfig:"WAIT"`
	Recreate     bool     `envconfig:"RECREATE_PODS" default:"false"`
	WaitTimeout  uint32   `envconfig:"WAIT_TIMEOUT" default:"300"`
	Actions      []string `envconfig:"ACTIONS" required:"true"`
	AuthKey      string   `envconfig:"AUTH_KEY"`
	KeyPath      string   `envconfig:"KEY_PATH"`
	Zone         string   `envconfig:"ZONE"`
	Cluster      string   `envconfig:"CLUSTER"`
	Project      string   `envconfig:"PROJECT"`
	Namespace    string   `envconfig:"NAMESPACE"`
	ChartRepo    string   `envconfig:"CHART_REPO"`
	Bucket       string   `envconfig:"BUCKET"`
	ChartPath    string   `envconfig:"CHART_PATH" required:"true"`
	ChartVersion string   `envconfig:"CHART_VERSION"`
	Release      string   `envconfig:"RELEASE"`
	Package      string   `envconfig:"PACKAGE"`
	Values       []string `envconfig:"VALUES"`
	ValueFiles   []string `envconfig:"VALUE_FILES"`
}

const (
	gcloudBin  = "/opt/google-cloud-sdk/bin/gcloud"
	gsutilBin  = "/opt/google-cloud-sdk/bin/gsutil"
	kubectlBin = "/opt/google-cloud-sdk/bin/kubectl"
	helmBin    = "/opt/google-cloud-sdk/bin/helm"

	lintPkg   = "lint"
	createPkg = "create"
	pushPkg   = "push"
	pullPkg   = "pull"
	deployPkg = "deploy"
	testPkg   = "test"
)

var reVersions = regexp.MustCompile(`(?P<realm>Client|Server): &version.Version.SemVer:"(?P<semver>.*?)".*?GitCommit:"(?P<commit>.*?)".*?GitTreeState:"(?P<treestate>.*?)"`)

// Exec executes the plugin step.
func (p Plugin) Exec() error {
	// only setup project when needed args are provided
	if p.Project != "" && p.Cluster != "" && p.Zone != "" {
		if err := setupProject(p.Project, p.Cluster, p.Zone, p.Debug); err != nil {
			return err
		}
		if err := helmInit(p.Debug); err != nil {
			return err
		}
	}

	for _, a := range p.Actions {
		switch a {
		case lintPkg:
			if err := p.lintPackage(); err != nil {
				return err
			}
		case createPkg:
			if err := p.createPackage(); err != nil {
				return err
			}
		case pushPkg:
			if err := p.pushPackage(); err != nil {
				return err
			}
		case pullPkg:
			if err := p.pullPackage(); err != nil {
				return err
			}
		case deployPkg:
			if err := p.deployPackage(); err != nil {
				return err
			}
		case testPkg:
			if err := p.testPackage(); err != nil {
				return err
			}
		default:
			return errors.New("unknown action")
		}
	}

	return nil
}

// setupProject setups gcloud project.
func setupProject(project, cluster, zone string, debug bool) error {
	// project configuration
	cmd := exec.Command(gcloudBin, "config", "set", "project", project)
	if err := run(cmd, debug); err != nil {
		return fmt.Errorf("could not the configure the project with glcoud: %v", err)
	}

	// cluster configuration
	cmd = exec.Command(gcloudBin, "container", "clusters", "get-credentials", cluster, "--zone", zone)
	if err := run(cmd, debug); err != nil {
		return fmt.Errorf("could not configure the cluster with glcoud: %v", err)
	}

	return nil
}

// setupAuth configures gcloud to use the given authFile
func setupAuth(authFile string, debug bool) error {
	if err := os.Setenv("GOOGLE_APPLICATION_CREDENTIALS", authFile); err != nil {
		return fmt.Errorf("could not set GOOGLE_APPLICATION_CREDENTIALS env variable: %v", err)
	}

	// authorization
	cmd := exec.Command(gcloudBin, "auth", "activate-service-account", fmt.Sprintf("--key-file=%s", authFile))
	if err := run(cmd, debug); err != nil {
		return fmt.Errorf("could not authorize with glcoud: %v", err)
	}
	return nil
}

// createPackage creates Helm package for Kubernetes.
// helm package --version $PLUGIN_CHART_VERSION $PLUGIN_CHART_PATH
func (p Plugin) createPackage() error {
	return run(exec.Command(helmBin, "package", "--version", p.ChartVersion, p.ChartPath), p.Debug)
}

// cpPackage copies a file from SOURCE to DEST
// gsutil cp SOURCE DEST
func (p Plugin) cpPackage(source string, dest string) error {
	return run(exec.Command(gsutilBin, "cp", source, dest), p.Debug)
}

// cpPackage pulls helm chart from Google Storage to local
// gsutil cp $PACKAGE-$PLUGIN_CHART_VERSION.tgz gs://$PLUGIN_BUCKET
func (p Plugin) pullPackage() error {
	return p.cpPackage(
		fmt.Sprintf("gs://%s/%s-%s.tgz", p.Bucket, p.Package, p.ChartVersion),
		fmt.Sprintf("%s-%s.tgz", p.Package, p.ChartVersion),
	)
}

// pushPackage pushes Helm package to the Google Storage.
// gsutil cp $PACKAGE-$PLUGIN_CHART_VERSION.tgz gs://$PLUGIN_BUCKET
func (p Plugin) pushPackage() error {
	return p.cpPackage(
		fmt.Sprintf("%s-%s.tgz", p.Package, p.ChartVersion),
		fmt.Sprintf("gs://%s", p.Bucket),
	)
}

// helm lint $CHARTPATH -i
func (p Plugin) lintPackage() error {
	return run(exec.Command(helmBin, "lint", p.ChartPath), p.Debug)
}

// helm upgrade $PACKAGE $PACKAGE-$PLUGIN_CHART_VERSION.tgz -i
func (p Plugin) deployPackage() error {
	args := []string{
		helmBin,
		"upgrade",
		p.Release,
		fmt.Sprintf("%s-%s.tgz", p.Package, p.ChartVersion),
	}
	if len(p.ValueFiles) > 0 {
		for _, f := range p.ValueFiles {
			args = append(args, "-f", f)
		}
	}
	if len(p.Values) > 0 {
		args = append(args, "--set", strings.Join(p.Values, ","))
	}
	if p.Recreate {
		args = append(args, "--recreate-pods")
	}
	args = append(args, "--install")
	args = append(args, "--namespace", p.Namespace)

	if p.Wait {
		args = append(args, "--wait", "--timeout", strconv.Itoa(int(p.WaitTimeout)))
	}
	return run(exec.Command("/bin/sh", "-c", strings.Join(args, " ")), p.Debug)
}

// helm test $PACKAGE
func (p Plugin) testPackage() error {
	args := []string{helmBin, "test", p.Release, "--cleanup", "--namespace", p.Namespace, "--timeout", strconv.Itoa(int(p.WaitTimeout))}
	return run(exec.Command("/bin/sh", "-c", strings.Join(args, " ")), p.Debug)
}

// fetchHelmVersions returns helm and tiller versions as map
func fetchHelmVersions(debug bool) (map[string]map[string]string, error) {
	var out bytes.Buffer
	var stderr bytes.Buffer
	cmd := exec.Command(helmBin, "version")
	cmd.Stdout = &out
	cmd.Stderr = &stderr

	if debug {
		log.Printf("running: %s", strings.Join(cmd.Args, " "))
	}
	if err := cmd.Run(); err != nil {
		return nil, errors.New(stderr.String())
	}
	if debug {
		log.Printf("%s", out.String())
	}

	lines := strings.Split(out.String(), "\n")
	versions := make(map[string]map[string]string)

	// we just care about the first two lines
	for _, line := range lines[:2] {
		entry, reErr := scanNamed(line, reVersions)
		if reErr != nil {
			return nil, reErr
		}
		versions[strings.ToLower(entry["realm"])] = entry
	}

	return versions, nil
}

// helmInit inits Triller on Kubernetes cluster.
func helmInit(debug bool) error {
	args := []string{"init"}

	ver, err := fetchHelmVersions(debug)
	if err == nil {
		switch strings.Compare(ver["client"]["semver"], ver["server"]["semver"]) {
		case -1: // client is older than tiller
			return fmt.Errorf("helm client is out of date")
		case 1: // client is newer than tiller
			args = append(args, "--upgrade")
		default: // client and tiller are at the same version
			args = append(args, "--client-only")
		}
	}

	if err := run(exec.Command(helmBin, args...), debug); err != nil {
		return err
	}

	// poll for tiller (call helm version 10 times)
	return pollTiller(10, debug)
}

// pollTiller repeatedly calls helm version and checks its exit code
func pollTiller(retries int, debug bool) error {
	var err error
	for i := 0; i < retries; i++ {
		if err = run(exec.Command(helmBin, "version"), debug); err == nil {
			return nil
		}
	}
	return err
}

func (p Plugin) movePkg() error {
	if err := os.Mkdir(p.Bucket, os.ModeDir); err != nil {
		return err
	}
	return cp(
		fmt.Sprintf("%s-%s.tgz", p.Package, p.ChartVersion),
		fmt.Sprintf("%s/%s-%s.tgz", p.Bucket, p.Package, p.ChartVersion),
	)
}

// cp copies file
func cp(src, dst string) error {
	s, err := os.Open(src)
	if err != nil {
		return err
	}
	// no need to check errors on read only file, we already got everything
	// we need from the filesystem, so nothing can go wrong now.
	defer s.Close()
	d, err := os.Create(dst)
	if err != nil {
		return err
	}
	if _, err := io.Copy(d, s); err != nil {
		d.Close()
		return err
	}
	return d.Close()
}

// scanNamed maps named regex groups to a golang map
func scanNamed(str string, rg *regexp.Regexp) (map[string]string, error) {
	result := make(map[string]string)
	for _, match := range rg.FindAllStringSubmatch(str, -1) {
		for i, name := range rg.SubexpNames() {
			if i != 0 && match[i] != "" && name != "" {
				result[name] = match[i]
			}
		}
	}

	if len(result) == 0 {
		return nil, errors.New("emtpy resultset")
	}

	return result, nil
}

func run(cmd *exec.Cmd, debug bool) error {
	if debug {
		log.Printf("running: %s", strings.Join(cmd.Args, " "))
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
	}
	return cmd.Run()
}
<|endoftext|>


<commit_before>
package ssh

import (
	"crypto"
	"errors"
	"fmt"
	"net"
)

func DialASlurp(network, addr string, config *ClientConfig) ([]byte, error) {
	conn, err := net.Dial(network, addr)
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	return Slurp(conn, config)
}

func Slurp(c net.Conn, config *ClientConfig) ([]byte, error) {
	conn := &ClientConn{
		transport:     newTransport(c, config.rand()),
		config:        config,
		globalRequest: globalRequest{response: make(chan interface{}, 1)},
	}
	defer conn.Close()
	return conn.slurpHostKey()
}

// SlurpHostKey acts like its going to handshake, but it just plays along long enough to get the host key.
func (c *ClientConn) slurpHostKey() ([]byte, error) {
	var magics handshakeMagics

	if _, err := c.Write(clientVersion); err != nil {
		return nil, err
	}
	if err := c.Flush(); err != nil {
		return nil, err
	}
	magics.clientVersion = clientVersion[:len(clientVersion)-2]

	// read remote server version
	version, err := readVersion(c)
	if err != nil {
		return nil, err
	}
	magics.serverVersion = version
	clientKexInit := kexInitMsg{
		KexAlgos:                supportedKexAlgos,
		ServerHostKeyAlgos:      supportedHostKeyAlgos,
		CiphersClientServer:     c.config.Crypto.ciphers(),
		CiphersServerClient:     c.config.Crypto.ciphers(),
		MACsClientServer:        c.config.Crypto.macs(),
		MACsServerClient:        c.config.Crypto.macs(),
		CompressionClientServer: supportedCompressions,
		CompressionServerClient: supportedCompressions,
	}
	kexInitPacket := marshal(msgKexInit, clientKexInit)
	magics.clientKexInit = kexInitPacket

	if err := c.writePacket(kexInitPacket); err != nil {
		return nil, err
	}
	packet, err := c.readPacket()
	if err != nil {
		return nil, err
	}

	magics.serverKexInit = packet

	var serverKexInit kexInitMsg
	if err = unmarshal(&serverKexInit, packet, msgKexInit); err != nil {
		return nil, err
	}

	kexAlgo, hostKeyAlgo, ok := findAgreedAlgorithms(c.transport, &clientKexInit, &serverKexInit)
	if !ok {
		return nil, errors.New("ssh: no common algorithms")
	}

	if serverKexInit.FirstKexFollows && kexAlgo != serverKexInit.KexAlgos[0] {
		// The server sent a Kex message for the wrong algorithm,
		// which we have to ignore.
		if _, err := c.readPacket(); err != nil {
			return nil, err
		}
	}

	var hostKey []byte
	var hashFunc crypto.Hash
	switch kexAlgo {
	case kexAlgoDH14SHA1:
		hashFunc = crypto.SHA1
		dhGroup14Once.Do(initDHGroup14)
		_, _, hostKey, err = c.kexDH(dhGroup14, hashFunc, &magics, hostKeyAlgo)
	case keyAlgoDH1SHA1:
		hashFunc = crypto.SHA1
		dhGroup1Once.Do(initDHGroup1)
		_, _, hostKey, err = c.kexDH(dhGroup1, hashFunc, &magics, hostKeyAlgo)
	default:
		err = fmt.Errorf("ssh: unexpected key exchange algorithm %v", kexAlgo)
	}
	if err != nil {
		return nil, err
	}

	return hostKey, err
}
<commit_msg>Add public key known_hosts formatting, fingerprint.<commit_after>package ssh

import (
	"bytes"
	"crypto"
	"crypto/md5"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"net"
)

func DialASlurp(network, addr string, config *ClientConfig) ([]byte, error) {
	conn, err := net.Dial(network, addr)
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	return Slurp(conn, config)
}

func Slurp(c net.Conn, config *ClientConfig) ([]byte, error) {
	conn := &ClientConn{
		transport:     newTransport(c, config.rand()),
		config:        config,
		globalRequest: globalRequest{response: make(chan interface{}, 1)},
	}
	defer conn.Close()
	return conn.slurpHostKey()
}

// SlurpHostKey acts like its going to handshake, but it just plays along long enough to get the host key.
func (c *ClientConn) slurpHostKey() ([]byte, error) {
	var magics handshakeMagics

	if _, err := c.Write(clientVersion); err != nil {
		return nil, err
	}
	if err := c.Flush(); err != nil {
		return nil, err
	}
	magics.clientVersion = clientVersion[:len(clientVersion)-2]

	// read remote server version
	version, err := readVersion(c)
	if err != nil {
		return nil, err
	}
	magics.serverVersion = version
	clientKexInit := kexInitMsg{
		KexAlgos:                supportedKexAlgos,
		ServerHostKeyAlgos:      supportedHostKeyAlgos,
		CiphersClientServer:     c.config.Crypto.ciphers(),
		CiphersServerClient:     c.config.Crypto.ciphers(),
		MACsClientServer:        c.config.Crypto.macs(),
		MACsServerClient:        c.config.Crypto.macs(),
		CompressionClientServer: supportedCompressions,
		CompressionServerClient: supportedCompressions,
	}
	kexInitPacket := marshal(msgKexInit, clientKexInit)
	magics.clientKexInit = kexInitPacket

	if err := c.writePacket(kexInitPacket); err != nil {
		return nil, err
	}
	packet, err := c.readPacket()
	if err != nil {
		return nil, err
	}

	magics.serverKexInit = packet

	var serverKexInit kexInitMsg
	if err = unmarshal(&serverKexInit, packet, msgKexInit); err != nil {
		return nil, err
	}

	kexAlgo, hostKeyAlgo, ok := findAgreedAlgorithms(c.transport, &clientKexInit, &serverKexInit)
	if !ok {
		return nil, errors.New("ssh: no common algorithms")
	}

	if serverKexInit.FirstKexFollows && kexAlgo != serverKexInit.KexAlgos[0] {
		// The server sent a Kex message for the wrong algorithm,
		// which we have to ignore.
		if _, err := c.readPacket(); err != nil {
			return nil, err
		}
	}

	var hostKey []byte
	var hashFunc crypto.Hash
	switch kexAlgo {
	case kexAlgoDH14SHA1:
		hashFunc = crypto.SHA1
		dhGroup14Once.Do(initDHGroup14)
		_, _, hostKey, err = c.kexDH(dhGroup14, hashFunc, &magics, hostKeyAlgo)
	case keyAlgoDH1SHA1:
		hashFunc = crypto.SHA1
		dhGroup1Once.Do(initDHGroup1)
		_, _, hostKey, err = c.kexDH(dhGroup1, hashFunc, &magics, hostKeyAlgo)
	default:
		err = fmt.Errorf("ssh: unexpected key exchange algorithm %v", kexAlgo)
	}
	if err != nil {
		return nil, err
	}

	return hostKey, err
}

type SlurpedHostKey struct {
	Key interface{}
}

var HostKeyParseErr error = errors.New("Parse error on host key")

func ParseHostKey(buf []byte) (*SlurpedHostKey, error) {
	key, _, ok := parsePubKey(buf)
	if !ok {
		return nil, HostKeyParseErr
	}
	return &SlurpedHostKey{key}, nil
}

func (pk *SlurpedHostKey) String() string {
	armor := base64.StdEncoding.EncodeToString(serializePublickey(pk.Key))
	return fmt.Sprintf("%s %s", algoName(pk.Key), armor)
}

func (pk *SlurpedHostKey) Fingerprint() string {
	h := md5.New()
	h.Write(serializePublickey(pk.Key))
	fp := h.Sum(nil)
	result := bytes.NewBuffer([]byte{})
	for i := 0; i < len(fp); i++ {
		if i > 0 {
			result.WriteByte(byte(':'))
		}
		io.WriteString(result, fmt.Sprintf("%02x", fp[i]))
	}
	return string(result.Bytes())
}
<|endoftext|>


<commit_before>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package sh

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

type Prog struct {
	Stmts []Node
}

func (p Prog) String() string {
	return nodeJoin(p.Stmts, "; ")
}

type Node interface {
	fmt.Stringer
}

func nodeJoin(ns []Node, sep string) string {
	var b bytes.Buffer
	for i, n := range ns {
		if i > 0 {
			io.WriteString(&b, sep)
		}
		io.WriteString(&b, n.String())
	}
	return b.String()
}

func stmtJoin(ns []Node) string {
	return nodeJoin(ns, "; ")
}

func wordJoin(ns []Node) string {
	return nodeJoin(ns, " ")
}

type Command struct {
	Args []Node

	Background bool
}

func (c Command) String() string {
	nodes := make([]Node, 0, len(c.Args))
	for _, l := range c.Args {
		nodes = append(nodes, l)
	}
	suffix := ""
	if c.Background {
		suffix += " &"
	}
	return wordJoin(nodes) + suffix
}

type Redirect struct {
	Op  Token
	Obj Node
}

func (r Redirect) String() string {
	return r.Op.String() + r.Obj.String()
}

type Subshell struct {
	Stmts []Node
}

func (s Subshell) String() string {
	return "( " + stmtJoin(s.Stmts) + "; )"
}

type Block struct {
	Stmts []Node
}

func (b Block) String() string {
	return "{ " + stmtJoin(b.Stmts) + "; }"
}

type IfStmt struct {
	Cond      Node
	ThenStmts []Node
	Elifs     []Node
	ElseStmts []Node
}

func (s IfStmt) String() string {
	var b bytes.Buffer
	io.WriteString(&b, "if ")
	io.WriteString(&b, s.Cond.String())
	io.WriteString(&b, "; then ")
	io.WriteString(&b, stmtJoin(s.ThenStmts))
	for _, n := range s.Elifs {
		e := n.(Elif)
		io.WriteString(&b, "; ")
		io.WriteString(&b, e.String())
	}
	if len(s.ElseStmts) > 0 {
		io.WriteString(&b, "; else ")
		io.WriteString(&b, stmtJoin(s.ElseStmts))
	}
	io.WriteString(&b, "; fi")
	return b.String()
}

type Elif struct {
	Cond      Node
	ThenStmts []Node
}

func (e Elif) String() string {
	var b bytes.Buffer
	io.WriteString(&b, "elif ")
	io.WriteString(&b, e.Cond.String())
	io.WriteString(&b, "; then ")
	io.WriteString(&b, stmtJoin(e.ThenStmts))
	return b.String()
}

type WhileStmt struct {
	Cond    Node
	DoStmts []Node
}

func (w WhileStmt) String() string {
	var b bytes.Buffer
	io.WriteString(&b, "while ")
	io.WriteString(&b, w.Cond.String())
	io.WriteString(&b, "; do ")
	io.WriteString(&b, stmtJoin(w.DoStmts))
	io.WriteString(&b, "; done")
	return b.String()
}

type ForStmt struct {
	Name     Node
	WordList []Node
	DoStmts  []Node
}

func (w ForStmt) String() string {
	var b bytes.Buffer
	io.WriteString(&b, "for ")
	io.WriteString(&b, w.Name.String())
	io.WriteString(&b, " in ")
	io.WriteString(&b, wordJoin(w.WordList))
	io.WriteString(&b, "; do ")
	io.WriteString(&b, stmtJoin(w.DoStmts))
	io.WriteString(&b, "; done")
	return b.String()
}

type BinaryExpr struct {
	X, Y Node
	Op   Token
}

func (b BinaryExpr) String() string {
	return fmt.Sprintf("%s %s %s", b.X, b.Op, b.Y)
}

type Comment struct {
	Text string
}

func (c Comment) String() string {
	return "#" + c.Text
}

type FuncDecl struct {
	Name Lit
	Body Node
}

func (f FuncDecl) String() string {
	return fmt.Sprintf("%s() %s", f.Name, f.Body)
}

type Word struct {
	Parts []Node
}

func (w Word) String() string {
	return nodeJoin(w.Parts, "")
}

type Lit struct {
	Val string
}

func (l Lit) String() string {
	return l.Val
}
<commit_msg>Remove now unnecessary args slicing<commit_after>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package sh

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

type Prog struct {
	Stmts []Node
}

func (p Prog) String() string {
	return nodeJoin(p.Stmts, "; ")
}

type Node interface {
	fmt.Stringer
}

func nodeJoin(ns []Node, sep string) string {
	var b bytes.Buffer
	for i, n := range ns {
		if i > 0 {
			io.WriteString(&b, sep)
		}
		io.WriteString(&b, n.String())
	}
	return b.String()
}

func stmtJoin(ns []Node) string {
	return nodeJoin(ns, "; ")
}

func wordJoin(ns []Node) string {
	return nodeJoin(ns, " ")
}

type Command struct {
	Args []Node

	Background bool
}

func (c Command) String() string {
	suffix := ""
	if c.Background {
		suffix += " &"
	}
	return wordJoin(c.Args) + suffix
}

type Redirect struct {
	Op  Token
	Obj Node
}

func (r Redirect) String() string {
	return r.Op.String() + r.Obj.String()
}

type Subshell struct {
	Stmts []Node
}

func (s Subshell) String() string {
	return "( " + stmtJoin(s.Stmts) + "; )"
}

type Block struct {
	Stmts []Node
}

func (b Block) String() string {
	return "{ " + stmtJoin(b.Stmts) + "; }"
}

type IfStmt struct {
	Cond      Node
	ThenStmts []Node
	Elifs     []Node
	ElseStmts []Node
}

func (s IfStmt) String() string {
	var b bytes.Buffer
	io.WriteString(&b, "if ")
	io.WriteString(&b, s.Cond.String())
	io.WriteString(&b, "; then ")
	io.WriteString(&b, stmtJoin(s.ThenStmts))
	for _, n := range s.Elifs {
		e := n.(Elif)
		io.WriteString(&b, "; ")
		io.WriteString(&b, e.String())
	}
	if len(s.ElseStmts) > 0 {
		io.WriteString(&b, "; else ")
		io.WriteString(&b, stmtJoin(s.ElseStmts))
	}
	io.WriteString(&b, "; fi")
	return b.String()
}

type Elif struct {
	Cond      Node
	ThenStmts []Node
}

func (e Elif) String() string {
	var b bytes.Buffer
	io.WriteString(&b, "elif ")
	io.WriteString(&b, e.Cond.String())
	io.WriteString(&b, "; then ")
	io.WriteString(&b, stmtJoin(e.ThenStmts))
	return b.String()
}

type WhileStmt struct {
	Cond    Node
	DoStmts []Node
}

func (w WhileStmt) String() string {
	var b bytes.Buffer
	io.WriteString(&b, "while ")
	io.WriteString(&b, w.Cond.String())
	io.WriteString(&b, "; do ")
	io.WriteString(&b, stmtJoin(w.DoStmts))
	io.WriteString(&b, "; done")
	return b.String()
}

type ForStmt struct {
	Name     Node
	WordList []Node
	DoStmts  []Node
}

func (w ForStmt) String() string {
	var b bytes.Buffer
	io.WriteString(&b, "for ")
	io.WriteString(&b, w.Name.String())
	io.WriteString(&b, " in ")
	io.WriteString(&b, wordJoin(w.WordList))
	io.WriteString(&b, "; do ")
	io.WriteString(&b, stmtJoin(w.DoStmts))
	io.WriteString(&b, "; done")
	return b.String()
}

type BinaryExpr struct {
	X, Y Node
	Op   Token
}

func (b BinaryExpr) String() string {
	return fmt.Sprintf("%s %s %s", b.X, b.Op, b.Y)
}

type Comment struct {
	Text string
}

func (c Comment) String() string {
	return "#" + c.Text
}

type FuncDecl struct {
	Name Lit
	Body Node
}

func (f FuncDecl) String() string {
	return fmt.Sprintf("%s() %s", f.Name, f.Body)
}

type Word struct {
	Parts []Node
}

func (w Word) String() string {
	return nodeJoin(w.Parts, "")
}

type Lit struct {
	Val string
}

func (l Lit) String() string {
	return l.Val
}
<|endoftext|>


<commit_before>package casbin

import (
	"github.com/hsluoyz/casbin/util"
	"log"
)

// Initialize the roles in RBAC.
func (model Model) BuildRoleLinks() {
	for _, ast := range model["g"] {
		ast.buildRoleLinks()
	}
}

// Print the policy.
func (model Model) PrintPolicy() {
	log.Print("Policy:")
	for key, ast := range model["p"] {
		log.Print(key, ": ", ast.Value, ": ", ast.Policy)
	}

	for key, ast := range model["g"] {
		log.Print(key, ": ", ast.Value, ": ", ast.Policy)
	}
}

// Clear all current policy.
func (model Model) ClearPolicy() {
	for _, ast := range model["p"] {
		ast.Policy = nil
	}

	for _, ast := range model["g"] {
		ast.Policy = nil
	}
}

// Get all rules in a policy.
func (model Model) GetPolicy(sec string, ptype string) [][]string {
	return model[sec][ptype].Policy
}

// Get rules based on a field filter from a policy.
func (model Model) GetFilteredPolicy(sec string, ptype string, fieldIndex int, fieldValue string) [][]string {
	res := [][]string{}

	for _, v := range model[sec][ptype].Policy {
		if v[fieldIndex] == fieldValue {
			res = append(res, v)
		}
	}

	return res
}

// Determine whether a model has the specified policy rule.
func (model Model) HasPolicy(sec string, ptype string, policy []string) bool {
	for _, rule := range model[sec][ptype].Policy {
		if util.ArrayEquals(policy, rule) {
			return true
		}
	}

	return false
}

// Add a policy rule to the model.
func (model Model) AddPolicy(sec string, ptype string, policy []string) bool {
	if !model.HasPolicy(sec, ptype, policy) {
		model[sec][ptype].Policy = append(model[sec][ptype].Policy, policy)
		return true
	} else {
		return false
	}
}

// Remove a policy rule from the model.
func (model Model) RemovePolicy(sec string, ptype string, policy []string) bool {
	for i, rule := range model[sec][ptype].Policy {
		if util.ArrayEquals(policy, rule) {
			model[sec][ptype].Policy = append(model[sec][ptype].Policy[:i], model[sec][ptype].Policy[i+1:]...)
			return true
		}
	}

	return false
}

// Remove policy rules based on a field filter from the model.
func (model Model) RemoveFilteredPolicy(sec string, ptype string, fieldIndex int, fieldValue string) bool {
	res := false
	for i := range model[sec][ptype].Policy {
		if model[sec][ptype].Policy[i][fieldIndex] == fieldValue {
			model[sec][ptype].Policy = append(model[sec][ptype].Policy[:i], model[sec][ptype].Policy[i+1:]...)
			i -= 1
			res = true
		}
	}

	return res
}

// Get all values for a field for all rules in a policy, duplicated values are removed.
func (model Model) GetValuesForFieldInPolicy(sec string, ptype string, fieldIndex int) []string {
	users := []string{}

	for _, rule := range model[sec][ptype].Policy {
		users = append(users, rule[fieldIndex])
	}

	util.ArrayRemoveDuplicates(&users)
	// sort.Strings(users)

	return users
}
<commit_msg>Fix the bug in removing policy with a filter function.<commit_after>package casbin

import (
	"github.com/hsluoyz/casbin/util"
	"log"
)

// Initialize the roles in RBAC.
func (model Model) BuildRoleLinks() {
	for _, ast := range model["g"] {
		ast.buildRoleLinks()
	}
}

// Print the policy.
func (model Model) PrintPolicy() {
	log.Print("Policy:")
	for key, ast := range model["p"] {
		log.Print(key, ": ", ast.Value, ": ", ast.Policy)
	}

	for key, ast := range model["g"] {
		log.Print(key, ": ", ast.Value, ": ", ast.Policy)
	}
}

// Clear all current policy.
func (model Model) ClearPolicy() {
	for _, ast := range model["p"] {
		ast.Policy = nil
	}

	for _, ast := range model["g"] {
		ast.Policy = nil
	}
}

// Get all rules in a policy.
func (model Model) GetPolicy(sec string, ptype string) [][]string {
	return model[sec][ptype].Policy
}

// Get rules based on a field filter from a policy.
func (model Model) GetFilteredPolicy(sec string, ptype string, fieldIndex int, fieldValue string) [][]string {
	res := [][]string{}

	for _, v := range model[sec][ptype].Policy {
		if v[fieldIndex] == fieldValue {
			res = append(res, v)
		}
	}

	return res
}

// Determine whether a model has the specified policy rule.
func (model Model) HasPolicy(sec string, ptype string, policy []string) bool {
	for _, rule := range model[sec][ptype].Policy {
		if util.ArrayEquals(policy, rule) {
			return true
		}
	}

	return false
}

// Add a policy rule to the model.
func (model Model) AddPolicy(sec string, ptype string, policy []string) bool {
	if !model.HasPolicy(sec, ptype, policy) {
		model[sec][ptype].Policy = append(model[sec][ptype].Policy, policy)
		return true
	} else {
		return false
	}
}

// Remove a policy rule from the model.
func (model Model) RemovePolicy(sec string, ptype string, policy []string) bool {
	for i, rule := range model[sec][ptype].Policy {
		if util.ArrayEquals(policy, rule) {
			model[sec][ptype].Policy = append(model[sec][ptype].Policy[:i], model[sec][ptype].Policy[i+1:]...)
			return true
		}
	}

	return false
}

// Remove policy rules based on a field filter from the model.
func (model Model) RemoveFilteredPolicy(sec string, ptype string, fieldIndex int, fieldValue string) bool {
	tmp := [][]string{}
	res := false
	for _, rule := range model[sec][ptype].Policy {
		if rule[fieldIndex] != fieldValue {
			tmp = append(tmp, rule)
		} else {
			res = true
		}
	}

	model[sec][ptype].Policy = tmp
	return res
}

// Get all values for a field for all rules in a policy, duplicated values are removed.
func (model Model) GetValuesForFieldInPolicy(sec string, ptype string, fieldIndex int) []string {
	users := []string{}

	for _, rule := range model[sec][ptype].Policy {
		users = append(users, rule[fieldIndex])
	}

	util.ArrayRemoveDuplicates(&users)
	// sort.Strings(users)

	return users
}
<|endoftext|>