<commit_before>package robots

import (
	"fmt"
	"math/rand"
	"strings"
	"time"
)

type DecideBot struct {
}

func init() {
	d := &DecideBot{}
	RegisterRobot("decide", d)
}

func (d DecideBot) Run(p *Payload) (slashCommandImmediateReturn string) {
	go d.DeferredAction(p)
	text := strings.TrimSpace(p.Text)
	if text == "" {
		return "I need something to decide on!"
	} else {
		return ""
	}
}

func (d DecideBot) DeferredAction(p *Payload) {
	response := &IncomingWebhook{
		Domain:      p.TeamDomain,
		Channel:     p.ChannelID,
		Username:    "Fate Bot",
		IconEmoji:   ":ghost:",
		UnfurlLinks: true,
		Parse:       ParseStyleFull,
	}
	text := strings.TrimSpace(p.Text)
	if text != "" {
		split := strings.Split(text, " ")
		response.Text = fmt.Sprintf("@%s: Deciding between: (%s)", p.UserName, strings.Join(split, ", "))
		go response.Send()
		response.Text = fmt.Sprintf("@%s: Decided on: %s", p.UserName, Decide(split))
		response.Send()
	}
}

func (d DecideBot) Description() (description string) {
	return "Decides your fate!\n\tUsage: /decide Life Death ...\n\tExpected Response: Deciding on (Life, Death, ...)\n\tDecided on Life!"
}

func Decide(Fates []string) (result string) {
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	n := len(Fates)
	if n > 0 {
		return Fates[r.Intn(n)]
	} else {
		return fmt.Sprintf("Error")
	}
}
<commit_msg>fix decide bot<commit_after>package robots

import (
	"fmt"
	"math/rand"
	"strings"
	"time"
)

type DecideBot struct {
}

func init() {
	d := &DecideBot{}
	RegisterRobot("decide", d)
}

func (d DecideBot) Run(p *Payload) (slashCommandImmediateReturn string) {
	go d.DeferredAction(p)
	text := strings.TrimSpace(p.Text)
	if text == "" {
		return "I need something to decide on!"
	} else {
		return ""
	}
}

func (d DecideBot) DeferredAction(p *Payload) {
	response := IncomingWebhook{
		Domain:      p.TeamDomain,
		Channel:     p.ChannelID,
		Username:    "Fate Bot",
		IconEmoji:   ":ghost:",
		UnfurlLinks: true,
		Parse:       ParseStyleFull,
	}
	text := strings.TrimSpace(p.Text)
	if text != "" {
		split := strings.Split(text, " ")
		response.Text = fmt.Sprintf("@%s: Deciding between: (%s) -> %s", p.UserName, strings.Join(split, ", "), Decide(split))
		response.Send()
	}
}

func (d DecideBot) Description() (description string) {
	return "Decides your fate!\n\tUsage: /decide Life Death ...\n\tExpected Response: Deciding on (Life, Death, ...)\n\tDecided on Life!"
}

func Decide(Fates []string) (result string) {
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	n := len(Fates)
	if n > 0 {
		return Fates[r.Intn(n)]
	} else {
		return fmt.Sprintf("Error")
	}
}
<|endoftext|>


<commit_before>package router

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"strings"
	"sync"

	"golang.org/x/crypto/ssh"

	vhost "github.com/inconshreveable/go-vhost"
	"github.com/miekg/dns"
)

type Protocol int

const (
	ProtocolHTTP Protocol = iota
	ProtocolHTTPS
	ProtocolSSH
	ProtocolDNS
)

type Director func(protocol Protocol, host string) (*net.TCPAddr, error)

type proxyRouter struct {
	sync.Mutex

	keyPath      string
	director     Director
	closed       bool
	httpListener net.Listener
	udpDnsServer *dns.Server
	tcpDnsServer *dns.Server
	sshListener  net.Listener
	sshConfig    *ssh.ServerConfig
}

func (r *proxyRouter) Listen(httpAddr, dnsAddr, sshAddr string) {
	l, err := net.Listen("tcp", httpAddr)
	if err != nil {
		log.Fatal(err)
	}
	r.httpListener = l
	go func() {
		for !r.closed {
			conn, err := r.httpListener.Accept()
			if err != nil {
				continue
			}
			go r.handleConnection(conn)
		}
	}()

	dnsMux := dns.NewServeMux()
	dnsMux.HandleFunc(".", r.dnsRequest)
	r.udpDnsServer = &dns.Server{Addr: dnsAddr, Net: "udp", Handler: dnsMux}
	r.tcpDnsServer = &dns.Server{Addr: dnsAddr, Net: "tcp", Handler: dnsMux}

	wg := sync.WaitGroup{}
	wg.Add(2)

	r.udpDnsServer.NotifyStartedFunc = func() {
		wg.Done()
	}
	r.tcpDnsServer.NotifyStartedFunc = func() {
		wg.Done()
	}
	go r.udpDnsServer.ListenAndServe()
	go r.tcpDnsServer.ListenAndServe()
	wg.Wait()

	lssh, err := net.Listen("tcp", sshAddr)
	if err != nil {
		log.Fatal("failed to listen for connection: ", err)
	}
	r.sshListener = lssh
	go func() {
		for {
			nConn, err := lssh.Accept()
			if err != nil {
				log.Fatal("failed to accept incoming connection: ", err)
			}

			go r.sshHandle(nConn)
		}
	}()
}
func (r *proxyRouter) ListenAndWait(httpAddr, dnsAddr, sshAddr string) {
	listenWG := sync.WaitGroup{}

	l, err := net.Listen("tcp", httpAddr)
	if err != nil {
		log.Fatal(err)
	}
	r.httpListener = l
	listenWG.Add(1)
	go func() {
		for !r.closed {
			conn, err := r.httpListener.Accept()
			if err != nil {
				continue
			}
			go r.handleConnection(conn)
		}
		listenWG.Done()
	}()

	dnsMux := dns.NewServeMux()
	dnsMux.HandleFunc(".", r.dnsRequest)
	r.udpDnsServer = &dns.Server{Addr: dnsAddr, Net: "udp", Handler: dnsMux}
	r.tcpDnsServer = &dns.Server{Addr: dnsAddr, Net: "tcp", Handler: dnsMux}

	wg := sync.WaitGroup{}
	wg.Add(2)

	r.udpDnsServer.NotifyStartedFunc = func() {
		wg.Done()
	}
	r.tcpDnsServer.NotifyStartedFunc = func() {
		wg.Done()
	}
	go r.udpDnsServer.ListenAndServe()
	go r.tcpDnsServer.ListenAndServe()
	wg.Wait()

	lssh, err := net.Listen("tcp", sshAddr)
	if err != nil {
		log.Fatal("failed to listen for connection: ", err)
	}
	r.sshListener = lssh
	listenWG.Add(1)
	go func() {
		for {
			nConn, err := lssh.Accept()
			if err != nil {
				log.Fatal("failed to accept incoming connection: ", err)
			}

			go r.sshHandle(nConn)
		}
		listenWG.Done()
	}()
	listenWG.Wait()
}

func (r *proxyRouter) sshHandle(nConn net.Conn) {
	sshCon, chans, reqs, err := ssh.NewServerConn(nConn, r.sshConfig)
	if err != nil {
		nConn.Close()
		return
	}

	dstHost, err := r.director(ProtocolSSH, sshCon.User())
	if err != nil {
		nConn.Close()
		return
	}

	// The incoming Request channel must be serviced.
	go ssh.DiscardRequests(reqs)

	newChannel := <-chans
	if newChannel == nil {
		sshCon.Close()
		return
	}

	if newChannel.ChannelType() != "session" {
		newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
		return
	}

	channel, requests, err := newChannel.Accept()
	if err != nil {
		log.Fatalf("Could not accept channel: %v", err)
	}

	stderr := channel.Stderr()

	fmt.Fprintf(stderr, "Connecting to %s\r\n", dstHost.String())

	clientConfig := &ssh.ClientConfig{
		User: "root",
		Auth: []ssh.AuthMethod{
			ssh.Password("root"),
		},
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
	}

	client, err := ssh.Dial("tcp", dstHost.String(), clientConfig)
	if err != nil {
		fmt.Fprintf(stderr, "Connect failed: %v\r\n", err)
		channel.Close()
		return
	}

	go func() {
		for newChannel = range chans {
			if newChannel == nil {
				return
			}

			channel2, reqs2, err := client.OpenChannel(newChannel.ChannelType(), newChannel.ExtraData())
			if err != nil {
				x, ok := err.(*ssh.OpenChannelError)
				if ok {
					newChannel.Reject(x.Reason, x.Message)
				} else {
					newChannel.Reject(ssh.Prohibited, "remote server denied channel request")
				}
				continue
			}

			channel, reqs, err := newChannel.Accept()
			if err != nil {
				channel2.Close()
				continue
			}
			go proxySsh(reqs, reqs2, channel, channel2)
		}
	}()

	// Forward the session channel
	channel2, reqs2, err := client.OpenChannel("session", []byte{})
	if err != nil {
		fmt.Fprintf(stderr, "Remote session setup failed: %v\r\n", err)
		channel.Close()
		return
	}

	maskedReqs := make(chan *ssh.Request, 1)
	go func() {
		for req := range requests {
			if req.Type == "auth-agent-req@openssh.com" {
				continue
			}
			maskedReqs <- req
		}
	}()
	proxySsh(maskedReqs, reqs2, channel, channel2)
}

func (r *proxyRouter) dnsRequest(w dns.ResponseWriter, req *dns.Msg) {
	if len(req.Question) > 0 {
		question := req.Question[0].Name

		if question == "localhost." {
			log.Printf("Asked for [localhost.] returning automatically [127.0.0.1]\n")
			m := new(dns.Msg)
			m.SetReply(req)
			m.Authoritative = true
			m.RecursionAvailable = true
			a, err := dns.NewRR(fmt.Sprintf("%s 60 IN A 127.0.0.1", question))
			if err != nil {
				log.Fatal(err)
			}
			m.Answer = append(m.Answer, a)
			w.WriteMsg(m)
			return
		}

		dstHost, err := r.director(ProtocolDNS, strings.TrimSuffix(question, "."))
		if err != nil {
			// Director couldn't resolve it, try to lookup in the system's DNS
			ips, err := net.LookupIP(question)
			if err != nil {
				// we have no information about this and we are not a recursive dns server, so we just fail so the client can fallback to the next dns server it has configured
				w.Close()
				// dns.HandleFailed(w, r)
				return
			}
			m := new(dns.Msg)
			m.SetReply(req)
			m.Authoritative = true
			m.RecursionAvailable = true
			for _, ip := range ips {
				ipv4 := ip.To4()
				if ipv4 == nil {
					a, err := dns.NewRR(fmt.Sprintf("%s 60 IN AAAA %s", question, ip.String()))
					if err != nil {
						log.Fatal(err)
					}
					m.Answer = append(m.Answer, a)
				} else {
					a, err := dns.NewRR(fmt.Sprintf("%s 60 IN A %s", question, ipv4.String()))
					if err != nil {
						log.Fatal(err)
					}
					m.Answer = append(m.Answer, a)
				}
			}
			w.WriteMsg(m)
			return
		}

		m := new(dns.Msg)
		m.SetReply(req)
		m.Authoritative = true
		m.RecursionAvailable = true
		a, err := dns.NewRR(fmt.Sprintf("%s 60 IN A %s", question, dstHost.IP))
		if err != nil {
			log.Println(err)
			w.Close()
			// dns.HandleFailed(w, r)
			return
		}
		m.Answer = append(m.Answer, a)
		w.WriteMsg(m)
		return
	}
}

func (r *proxyRouter) Close() {
	r.Lock()
	defer r.Unlock()

	if r.httpListener != nil {
		r.httpListener.Close()
	}
	if r.udpDnsServer != nil {
		r.udpDnsServer.Shutdown()
	}
	r.closed = true
}

func (r *proxyRouter) ListenHttpAddress() string {
	if r.httpListener != nil {
		return r.httpListener.Addr().String()
	}
	return ""
}

func (r *proxyRouter) ListenDnsUdpAddress() string {
	if r.udpDnsServer != nil && r.udpDnsServer.PacketConn != nil {
		return r.udpDnsServer.PacketConn.LocalAddr().String()
	}

	return ""
}
func (r *proxyRouter) ListenDnsTcpAddress() string {
	if r.tcpDnsServer != nil && r.tcpDnsServer.Listener != nil {
		return r.tcpDnsServer.Listener.Addr().String()
	}

	return ""
}

func (r *proxyRouter) ListenSshAddress() string {
	if r.sshListener != nil {
		return r.sshListener.Addr().String()
	}

	return ""
}

func (r *proxyRouter) handleConnection(c net.Conn) {
	defer c.Close()
	// first try tls
	vhostConn, err := vhost.TLS(c)

	if err == nil {
		// It is a TLS connection
		defer vhostConn.Close()
		host := vhostConn.ClientHelloMsg.ServerName
		dstHost, err := r.director(ProtocolHTTPS, host)
		if err != nil {
			log.Printf("Error directing request: %v\n", err)
			return
		}
		d, err := net.Dial("tcp", dstHost.String())
		if err != nil {
			log.Printf("Error dialing backend %s: %v\n", dstHost.String(), err)
			return
		}

		proxyConn(vhostConn, d)
	} else {
		// it is not TLS
		// treat it as an http connection

		req, err := http.ReadRequest(bufio.NewReader(vhostConn))
		if err != nil {
			// It is not http neither. So just close the connection.
			return
		}
		host := req.Header.Get("X-Forwarded-Host")
		if host == "" {
			host = req.Host
		}
		dstHost, err := r.director(ProtocolHTTP, host)
		if err != nil {
			log.Printf("Error directing request: %v\n", err)
			return
		}
		d, err := net.Dial("tcp", dstHost.String())
		if err != nil {
			log.Printf("Error dialing backend %s: %v\n", dstHost.String(), err)
			return
		}
		defer d.Close()
		err = req.Write(d)
		if err != nil {
			log.Printf("Error requesting backend %s: %v\n", dstHost.String(), err)
			return
		}
		proxyConn(c, d)
	}
}

func proxySsh(reqs1, reqs2 <-chan *ssh.Request, channel1, channel2 ssh.Channel) {
	var closer sync.Once
	closeFunc := func() {
		channel1.Close()
		channel2.Close()
	}

	defer closer.Do(closeFunc)

	closerChan := make(chan bool, 1)

	go func() {
		io.Copy(channel1, channel2)
		closerChan <- true
	}()

	go func() {
		io.Copy(channel2, channel1)
		closerChan <- true
	}()

	for {
		select {
		case req := <-reqs1:
			if req == nil {
				return
			}
			b, err := channel2.SendRequest(req.Type, req.WantReply, req.Payload)
			if err != nil {
				return
			}
			req.Reply(b, nil)

		case req := <-reqs2:
			if req == nil {
				return
			}
			b, err := channel1.SendRequest(req.Type, req.WantReply, req.Payload)
			if err != nil {
				return
			}
			req.Reply(b, nil)
		case <-closerChan:
			return
		}
	}
}

func proxyConn(src, dst net.Conn) {
	errc := make(chan error, 2)
	cp := func(dst io.Writer, src io.Reader) {
		_, err := io.Copy(dst, src)
		errc <- err
	}
	go cp(src, dst)
	go cp(dst, src)
	<-errc
}

func NewRouter(director Director, keyPath string) *proxyRouter {
	var sshConfig = &ssh.ServerConfig{
		PublicKeyCallback: func(c ssh.ConnMetadata, pubKey ssh.PublicKey) (*ssh.Permissions, error) {
			return nil, nil
		},
	}
	privateBytes, err := ioutil.ReadFile(keyPath)
	if err != nil {
		log.Fatal("Failed to load private key: ", err)
	}

	private, err := ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		log.Fatal("Failed to parse private key: ", err)
	}

	sshConfig.AddHostKey(private)

	return &proxyRouter{director: director, sshConfig: sshConfig}
}
<commit_msg>Small refactor to avoid code duplication<commit_after>package router

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"strings"
	"sync"

	"golang.org/x/crypto/ssh"

	vhost "github.com/inconshreveable/go-vhost"
	"github.com/miekg/dns"
)

type Protocol int

const (
	ProtocolHTTP Protocol = iota
	ProtocolHTTPS
	ProtocolSSH
	ProtocolDNS
)

type Director func(protocol Protocol, host string) (*net.TCPAddr, error)

type proxyRouter struct {
	sync.Mutex

	keyPath      string
	director     Director
	closed       bool
	httpListener net.Listener
	udpDnsServer *dns.Server
	tcpDnsServer *dns.Server
	sshListener  net.Listener
	sshConfig    *ssh.ServerConfig
}

func (r *proxyRouter) Listen(httpAddr, dnsAddr, sshAddr string) {
	r.listen(&sync.WaitGroup{}, httpAddr, dnsAddr, sshAddr)
}

func (r *proxyRouter) ListenAndWait(httpAddr, dnsAddr, sshAddr string) {
	wg := sync.WaitGroup{}
	r.listen(&wg, httpAddr, dnsAddr, sshAddr)
	wg.Wait()
}

func (r *proxyRouter) listen(wg *sync.WaitGroup, httpAddr, dnsAddr, sshAddr string) {

	l, err := net.Listen("tcp", httpAddr)
	if err != nil {
		log.Fatal(err)
	}
	r.httpListener = l
	wg.Add(1)
	go func() {
		for !r.closed {
			conn, err := r.httpListener.Accept()
			if err != nil {
				continue
			}
			go r.handleConnection(conn)
		}
		wg.Done()
	}()

	dnsMux := dns.NewServeMux()
	dnsMux.HandleFunc(".", r.dnsRequest)
	r.udpDnsServer = &dns.Server{Addr: dnsAddr, Net: "udp", Handler: dnsMux}
	r.tcpDnsServer = &dns.Server{Addr: dnsAddr, Net: "tcp", Handler: dnsMux}

	wgStarted := sync.WaitGroup{}
	wgStarted.Add(2)

	r.udpDnsServer.NotifyStartedFunc = func() {
		wgStarted.Done()
	}
	r.tcpDnsServer.NotifyStartedFunc = func() {
		wgStarted.Done()
	}
	go r.udpDnsServer.ListenAndServe()
	go r.tcpDnsServer.ListenAndServe()
	wgStarted.Wait()

	lssh, err := net.Listen("tcp", sshAddr)
	if err != nil {
		log.Fatal("failed to listen for connection: ", err)
	}
	r.sshListener = lssh
	wg.Add(1)
	go func() {
		for {
			nConn, err := lssh.Accept()
			if err != nil {
				log.Fatal("failed to accept incoming connection: ", err)
			}

			go r.sshHandle(nConn)
		}
		wg.Done()
	}()
}

func (r *proxyRouter) sshHandle(nConn net.Conn) {
	sshCon, chans, reqs, err := ssh.NewServerConn(nConn, r.sshConfig)
	if err != nil {
		nConn.Close()
		return
	}

	dstHost, err := r.director(ProtocolSSH, sshCon.User())
	if err != nil {
		nConn.Close()
		return
	}

	// The incoming Request channel must be serviced.
	go ssh.DiscardRequests(reqs)

	newChannel := <-chans
	if newChannel == nil {
		sshCon.Close()
		return
	}

	if newChannel.ChannelType() != "session" {
		newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
		return
	}

	channel, requests, err := newChannel.Accept()
	if err != nil {
		log.Fatalf("Could not accept channel: %v", err)
	}

	stderr := channel.Stderr()

	fmt.Fprintf(stderr, "Connecting to %s\r\n", dstHost.String())

	clientConfig := &ssh.ClientConfig{
		User: "root",
		Auth: []ssh.AuthMethod{
			ssh.Password("root"),
		},
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
	}

	client, err := ssh.Dial("tcp", dstHost.String(), clientConfig)
	if err != nil {
		fmt.Fprintf(stderr, "Connect failed: %v\r\n", err)
		channel.Close()
		return
	}

	go func() {
		for newChannel = range chans {
			if newChannel == nil {
				return
			}

			channel2, reqs2, err := client.OpenChannel(newChannel.ChannelType(), newChannel.ExtraData())
			if err != nil {
				x, ok := err.(*ssh.OpenChannelError)
				if ok {
					newChannel.Reject(x.Reason, x.Message)
				} else {
					newChannel.Reject(ssh.Prohibited, "remote server denied channel request")
				}
				continue
			}

			channel, reqs, err := newChannel.Accept()
			if err != nil {
				channel2.Close()
				continue
			}
			go proxySsh(reqs, reqs2, channel, channel2)
		}
	}()

	// Forward the session channel
	channel2, reqs2, err := client.OpenChannel("session", []byte{})
	if err != nil {
		fmt.Fprintf(stderr, "Remote session setup failed: %v\r\n", err)
		channel.Close()
		return
	}

	maskedReqs := make(chan *ssh.Request, 1)
	go func() {
		for req := range requests {
			if req.Type == "auth-agent-req@openssh.com" {
				continue
			}
			maskedReqs <- req
		}
	}()
	proxySsh(maskedReqs, reqs2, channel, channel2)
}

func (r *proxyRouter) dnsRequest(w dns.ResponseWriter, req *dns.Msg) {
	if len(req.Question) > 0 {
		question := req.Question[0].Name

		if question == "localhost." {
			log.Printf("Asked for [localhost.] returning automatically [127.0.0.1]\n")
			m := new(dns.Msg)
			m.SetReply(req)
			m.Authoritative = true
			m.RecursionAvailable = true
			a, err := dns.NewRR(fmt.Sprintf("%s 60 IN A 127.0.0.1", question))
			if err != nil {
				log.Fatal(err)
			}
			m.Answer = append(m.Answer, a)
			w.WriteMsg(m)
			return
		}

		dstHost, err := r.director(ProtocolDNS, strings.TrimSuffix(question, "."))
		if err != nil {
			// Director couldn't resolve it, try to lookup in the system's DNS
			ips, err := net.LookupIP(question)
			if err != nil {
				// we have no information about this and we are not a recursive dns server, so we just fail so the client can fallback to the next dns server it has configured
				w.Close()
				// dns.HandleFailed(w, r)
				return
			}
			m := new(dns.Msg)
			m.SetReply(req)
			m.Authoritative = true
			m.RecursionAvailable = true
			for _, ip := range ips {
				ipv4 := ip.To4()
				if ipv4 == nil {
					a, err := dns.NewRR(fmt.Sprintf("%s 60 IN AAAA %s", question, ip.String()))
					if err != nil {
						log.Fatal(err)
					}
					m.Answer = append(m.Answer, a)
				} else {
					a, err := dns.NewRR(fmt.Sprintf("%s 60 IN A %s", question, ipv4.String()))
					if err != nil {
						log.Fatal(err)
					}
					m.Answer = append(m.Answer, a)
				}
			}
			w.WriteMsg(m)
			return
		}

		m := new(dns.Msg)
		m.SetReply(req)
		m.Authoritative = true
		m.RecursionAvailable = true
		a, err := dns.NewRR(fmt.Sprintf("%s 60 IN A %s", question, dstHost.IP))
		if err != nil {
			log.Println(err)
			w.Close()
			// dns.HandleFailed(w, r)
			return
		}
		m.Answer = append(m.Answer, a)
		w.WriteMsg(m)
		return
	}
}

func (r *proxyRouter) Close() {
	r.Lock()
	defer r.Unlock()

	if r.httpListener != nil {
		r.httpListener.Close()
	}
	if r.udpDnsServer != nil {
		r.udpDnsServer.Shutdown()
	}
	r.closed = true
}

func (r *proxyRouter) ListenHttpAddress() string {
	if r.httpListener != nil {
		return r.httpListener.Addr().String()
	}
	return ""
}

func (r *proxyRouter) ListenDnsUdpAddress() string {
	if r.udpDnsServer != nil && r.udpDnsServer.PacketConn != nil {
		return r.udpDnsServer.PacketConn.LocalAddr().String()
	}

	return ""
}
func (r *proxyRouter) ListenDnsTcpAddress() string {
	if r.tcpDnsServer != nil && r.tcpDnsServer.Listener != nil {
		return r.tcpDnsServer.Listener.Addr().String()
	}

	return ""
}

func (r *proxyRouter) ListenSshAddress() string {
	if r.sshListener != nil {
		return r.sshListener.Addr().String()
	}

	return ""
}

func (r *proxyRouter) handleConnection(c net.Conn) {
	defer c.Close()
	// first try tls
	vhostConn, err := vhost.TLS(c)

	if err == nil {
		// It is a TLS connection
		defer vhostConn.Close()
		host := vhostConn.ClientHelloMsg.ServerName
		dstHost, err := r.director(ProtocolHTTPS, host)
		if err != nil {
			log.Printf("Error directing request: %v\n", err)
			return
		}
		d, err := net.Dial("tcp", dstHost.String())
		if err != nil {
			log.Printf("Error dialing backend %s: %v\n", dstHost.String(), err)
			return
		}

		proxyConn(vhostConn, d)
	} else {
		// it is not TLS
		// treat it as an http connection

		req, err := http.ReadRequest(bufio.NewReader(vhostConn))
		if err != nil {
			// It is not http neither. So just close the connection.
			return
		}
		host := req.Header.Get("X-Forwarded-Host")
		if host == "" {
			host = req.Host
		}
		dstHost, err := r.director(ProtocolHTTP, host)
		if err != nil {
			log.Printf("Error directing request: %v\n", err)
			return
		}
		d, err := net.Dial("tcp", dstHost.String())
		if err != nil {
			log.Printf("Error dialing backend %s: %v\n", dstHost.String(), err)
			return
		}
		defer d.Close()
		err = req.Write(d)
		if err != nil {
			log.Printf("Error requesting backend %s: %v\n", dstHost.String(), err)
			return
		}
		proxyConn(c, d)
	}
}

func proxySsh(reqs1, reqs2 <-chan *ssh.Request, channel1, channel2 ssh.Channel) {
	var closer sync.Once
	closeFunc := func() {
		channel1.Close()
		channel2.Close()
	}

	defer closer.Do(closeFunc)

	closerChan := make(chan bool, 1)

	go func() {
		io.Copy(channel1, channel2)
		closerChan <- true
	}()

	go func() {
		io.Copy(channel2, channel1)
		closerChan <- true
	}()

	for {
		select {
		case req := <-reqs1:
			if req == nil {
				return
			}
			b, err := channel2.SendRequest(req.Type, req.WantReply, req.Payload)
			if err != nil {
				return
			}
			req.Reply(b, nil)

		case req := <-reqs2:
			if req == nil {
				return
			}
			b, err := channel1.SendRequest(req.Type, req.WantReply, req.Payload)
			if err != nil {
				return
			}
			req.Reply(b, nil)
		case <-closerChan:
			return
		}
	}
}

func proxyConn(src, dst net.Conn) {
	errc := make(chan error, 2)
	cp := func(dst io.Writer, src io.Reader) {
		_, err := io.Copy(dst, src)
		errc <- err
	}
	go cp(src, dst)
	go cp(dst, src)
	<-errc
}

func NewRouter(director Director, keyPath string) *proxyRouter {
	var sshConfig = &ssh.ServerConfig{
		PublicKeyCallback: func(c ssh.ConnMetadata, pubKey ssh.PublicKey) (*ssh.Permissions, error) {
			return nil, nil
		},
	}
	privateBytes, err := ioutil.ReadFile(keyPath)
	if err != nil {
		log.Fatal("Failed to load private key: ", err)
	}

	private, err := ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		log.Fatal("Failed to parse private key: ", err)
	}

	sshConfig.AddHostKey(private)

	return &proxyRouter{director: director, sshConfig: sshConfig}
}
<|endoftext|>


<commit_before>package routes

import (
	"github.com/DVI-GI-2017/Jira__backend/handlers"
	"log"
)

func InitRouter(r *router) {
	const signup = "/signup"
	err := r.Post(signup, handlers.RegisterUser)
	if err != nil {
		log.Panicf("can not init route '%s': %v", signup, err)
	}

	const signin = "/signin"
	err = r.Post(signin, handlers.Login)
	if err != nil {
		log.Panicf("can not init route '%s': %v", signin, err)
	}

	const test = "/test"
	err = r.Post(test, handlers.Test)
	if err != nil {
		log.Panicf("can not init route '%s': %v", test, err)
	}
}
<commit_msg>add all users url<commit_after>package routes

import (
	"github.com/DVI-GI-2017/Jira__backend/handlers"
	"log"
)

func InitRouter(r *router) {
	const signup = "/signup"
	err := r.Post(signup, handlers.RegisterUser)
	if err != nil {
		log.Panicf("can not init route '%s': %v", signup, err)
	}

	const signin = "/signin"
	err = r.Post(signin, handlers.Login)
	if err != nil {
		log.Panicf("can not init route '%s': %v", signin, err)
	}

	const users = "/users"
	err = r.Get(users, handlers.AllUsers)
	if err != nil {
		log.Panicf("can not init route '%s': %v", users, err)
	}

	const test = "/test"
	err = r.Post(test, handlers.Test)
	if err != nil {
		log.Panicf("can not init route '%s': %v", test, err)
	}
}
<|endoftext|>


<commit_before>package rpcout

import (
	"errors"
	"fmt"

	"github.com/sigma/vmw-guestinfo/message"
)

var ErrRpciFormat = errors.New("invalid format for RPCI command result")

const RPCI_PROTOCOL_NUM uint32 = 0x49435052

func SendOne(format string, a ...interface{}) (reply []byte, err error) {
	request := fmt.Sprintf(format, a...)
	return SendOneRaw([]byte(request))
}

func SendOneRaw(request []byte) (reply []byte, err error) {
	out := &RpcOut{}
	if err = out.Start(); err != nil {
		return
	}
	if reply, err = out.Send(request); err != nil {
		return
	}
	if err = out.Stop(); err != nil {
		return
	}
	return
}

type RpcOut struct {
	channel message.Channel
}

func (out *RpcOut) Start() error {
	channel, err := message.Open(RPCI_PROTOCOL_NUM)
	if err != nil {
		return err
	}
	out.channel = channel
	return nil
}

func (out *RpcOut) Stop() error {
	err := message.Close(out.channel)
	out.channel = nil
	return err
}

func (out *RpcOut) Send(request []byte) (reply []byte, err error) {
	if err = message.Send(out.channel, request); err != nil {
		return
	}

	var resp []byte
	if resp, err = message.Receive(out.channel); err != nil {
		return
	}

	switch string(resp[:2]) {
	case "0 ":
		err = fmt.Errorf("RPCI error: %s", resp[2:])
	case "1 ":
		reply = resp[2:]
	default:
		err = ErrRpciFormat
	}
	return
}
<commit_msg>rpcout: add OK return value<commit_after>package rpcout

import (
	"errors"
	"fmt"

	"github.com/sigma/vmw-guestinfo/message"
)

var ErrRpciFormat = errors.New("invalid format for RPCI command result")

const RPCI_PROTOCOL_NUM uint32 = 0x49435052

func SendOne(format string, a ...interface{}) (reply []byte, ok bool, err error) {
	request := fmt.Sprintf(format, a...)
	return SendOneRaw([]byte(request))
}

func SendOneRaw(request []byte) (reply []byte, ok bool, err error) {
	out := &RpcOut{}
	if err = out.Start(); err != nil {
		return
	}
	if reply, ok, err = out.Send(request); err != nil {
		return
	}
	if err = out.Stop(); err != nil {
		return
	}
	return
}

type RpcOut struct {
	channel message.Channel
}

func (out *RpcOut) Start() error {
	channel, err := message.Open(RPCI_PROTOCOL_NUM)
	if err != nil {
		return err
	}
	out.channel = channel
	return nil
}

func (out *RpcOut) Stop() error {
	err := message.Close(out.channel)
	out.channel = nil
	return err
}

func (out *RpcOut) Send(request []byte) (reply []byte, ok bool, err error) {
	if err = message.Send(out.channel, request); err != nil {
		return
	}

	var resp []byte
	if resp, err = message.Receive(out.channel); err != nil {
		return
	}

	switch string(resp[:2]) {
	case "0 ":
		reply = resp[2:]
	case "1 ":
		reply = resp[2:]
		ok = true
	default:
		err = ErrRpciFormat
	}
	return
}
<|endoftext|>


<commit_before>package outputhttp

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tsaikd/gogstash/config"
	"github.com/tsaikd/gogstash/config/goglog"
	"github.com/tsaikd/gogstash/config/logevent"
)

func init() {
	goglog.Logger.SetLevel(logrus.DebugLevel)
	config.RegistOutputHandler(ModuleName, InitHandler)
}

func Test_output_http_module(t *testing.T) {
	assert := assert.New(t)
	assert.NotNil(assert)
	require := require.New(t)
	require.NotNil(require)

	serverRecvMsg := make(chan []byte, 1)
	handler := func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		data, err := io.ReadAll(r.Body)
		require.NoError(err)
		serverRecvMsg <- data
	}

	server := httptest.NewServer(http.HandlerFunc(handler))
	defer server.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	conf, err := config.LoadFromYAML([]byte(strings.TrimSpace(`
debugch: true
output:
  - type: http
    urls: ["` + server.URL + `"]
	`)))
	require.NoError(err)
	require.NoError(conf.Start(ctx))

	conf.TestInputEvent(logevent.LogEvent{
		Timestamp: time.Now(),
		Message:   "outputhttp test message",
	})

	select {
	case <-ctx.Done():
		t.Fatal("timeout")
	case msg := <-serverRecvMsg:
		require.Contains(string(msg), `"message":"outputhttp test message"`)
	}
}

func TestMapFromInts(t *testing.T) {
	input := []int{100, 200, 300}
	myMap := MapFromInts(input)
	// check length
	if len(myMap) != len(input) {
		t.Error("Length incorrect")
	}
	// check if first elem is in list
	_, ok := myMap[input[0]]
	if ok {
		t.Error("First element not in list")
	}
	// check element not in list
	_, ok = myMap[-1]
	if ok {
		t.Error("Found element not in list")
	}
}
<commit_msg>Fixed testing of maps to int<commit_after>package outputhttp

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tsaikd/gogstash/config"
	"github.com/tsaikd/gogstash/config/goglog"
	"github.com/tsaikd/gogstash/config/logevent"
)

func init() {
	goglog.Logger.SetLevel(logrus.DebugLevel)
	config.RegistOutputHandler(ModuleName, InitHandler)
}

func Test_output_http_module(t *testing.T) {
	assert := assert.New(t)
	assert.NotNil(assert)
	require := require.New(t)
	require.NotNil(require)

	serverRecvMsg := make(chan []byte, 1)
	handler := func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		data, err := io.ReadAll(r.Body)
		require.NoError(err)
		serverRecvMsg <- data
	}

	server := httptest.NewServer(http.HandlerFunc(handler))
	defer server.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	conf, err := config.LoadFromYAML([]byte(strings.TrimSpace(`
debugch: true
output:
  - type: http
    urls: ["` + server.URL + `"]
	`)))
	require.NoError(err)
	require.NoError(conf.Start(ctx))

	conf.TestInputEvent(logevent.LogEvent{
		Timestamp: time.Now(),
		Message:   "outputhttp test message",
	})

	select {
	case <-ctx.Done():
		t.Fatal("timeout")
	case msg := <-serverRecvMsg:
		require.Contains(string(msg), `"message":"outputhttp test message"`)
	}
}

func TestMapFromInts(t *testing.T) {
	input := []int{100, 200, 300}
	myMap := MapFromInts(input)
	// check length
	if len(myMap) != len(input) {
		t.Error("Length incorrect")
	}
	// check if first elem is in list
	if _, ok := myMap[input[0]]; !ok {
		t.Error("First element not in list")
	}
	// check element not in list
	if _, ok := myMap[-1]; ok {
		t.Error("Found element not in list")
	}
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"net/url"
	"strconv"

	"github.com/hashicorp/terraform/helper/schema"

	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/sqs"
)

var AttributeMap = map[string]string{
	"delay_seconds":              "DelaySeconds",
	"max_message_size":           "MaximumMessageSize",
	"message_retention_seconds":  "MessageRetentionPeriod",
	"receive_wait_time_seconds":  "ReceiveMessageWaitTimeSeconds",
	"visibility_timeout_seconds": "VisibilityTimeout",
	"policy":                     "Policy",
	"redrive_policy":             "RedrivePolicy",
	"arn":                        "QueueArn",
}

// A number of these are marked as computed because if you don't
// provide a value, SQS will provide you with defaults (which are the
// default values specified below)
func resourceAwsSqsQueue() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsSqsQueueCreate,
		Read:   resourceAwsSqsQueueRead,
		Update: resourceAwsSqsQueueUpdate,
		Delete: resourceAwsSqsQueueDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"delay_seconds": {
				Type:     schema.TypeInt,
				Optional: true,
				Default:  0,
			},
			"max_message_size": {
				Type:     schema.TypeInt,
				Optional: true,
				Default:  262144,
			},
			"message_retention_seconds": {
				Type:     schema.TypeInt,
				Optional: true,
				Default:  345600,
			},
			"receive_wait_time_seconds": {
				Type:     schema.TypeInt,
				Optional: true,
				Default:  0,
			},
			"visibility_timeout_seconds": {
				Type:     schema.TypeInt,
				Optional: true,
				Default:  30,
			},
			"policy": {
				Type:             schema.TypeString,
				Optional:         true,
				DiffSuppressFunc: suppressEquivalentAwsPolicyDiffs,
				Computed:         true,
			},
			"redrive_policy": {
				Type:      schema.TypeString,
				Optional:  true,
				StateFunc: normalizeJson,
			},
			"arn": {
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAwsSqsQueueCreate(d *schema.ResourceData, meta interface{}) error {
	sqsconn := meta.(*AWSClient).sqsconn

	name := d.Get("name").(string)

	log.Printf("[DEBUG] SQS queue create: %s", name)

	req := &sqs.CreateQueueInput{
		QueueName: aws.String(name),
	}

	attributes := make(map[string]*string)

	resource := *resourceAwsSqsQueue()

	for k, s := range resource.Schema {
		if attrKey, ok := AttributeMap[k]; ok {
			if value, ok := d.GetOk(k); ok {
				if s.Type == schema.TypeInt {
					attributes[attrKey] = aws.String(strconv.Itoa(value.(int)))
				} else {
					attributes[attrKey] = aws.String(value.(string))
				}
			}

		}
	}

	if len(attributes) > 0 {
		req.Attributes = attributes
	}

	output, err := sqsconn.CreateQueue(req)
	if err != nil {
		return fmt.Errorf("Error creating SQS queue: %s", err)
	}

	d.SetId(*output.QueueUrl)

	return resourceAwsSqsQueueUpdate(d, meta)
}

func resourceAwsSqsQueueUpdate(d *schema.ResourceData, meta interface{}) error {
	sqsconn := meta.(*AWSClient).sqsconn
	attributes := make(map[string]*string)

	resource := *resourceAwsSqsQueue()

	for k, s := range resource.Schema {
		if attrKey, ok := AttributeMap[k]; ok {
			if d.HasChange(k) {
				log.Printf("[DEBUG] Updating %s", attrKey)
				_, n := d.GetChange(k)
				if s.Type == schema.TypeInt {
					attributes[attrKey] = aws.String(strconv.Itoa(n.(int)))
				} else {
					attributes[attrKey] = aws.String(n.(string))
				}
			}
		}
	}

	if len(attributes) > 0 {
		req := &sqs.SetQueueAttributesInput{
			QueueUrl:   aws.String(d.Id()),
			Attributes: attributes,
		}
		if _, err := sqsconn.SetQueueAttributes(req); err != nil {
			return fmt.Errorf("[ERR] Error updating SQS attributes: %s", err)
		}
	}

	return resourceAwsSqsQueueRead(d, meta)
}

func resourceAwsSqsQueueRead(d *schema.ResourceData, meta interface{}) error {
	sqsconn := meta.(*AWSClient).sqsconn

	attributeOutput, err := sqsconn.GetQueueAttributes(&sqs.GetQueueAttributesInput{
		QueueUrl:       aws.String(d.Id()),
		AttributeNames: []*string{aws.String("All")},
	})

	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			log.Printf("ERROR Found %s", awsErr.Code())
			if "AWS.SimpleQueueService.NonExistentQueue" == awsErr.Code() {
				d.SetId("")
				log.Printf("[DEBUG] SQS Queue (%s) not found", d.Get("name").(string))
				return nil
			}
		}
		return err
	}

	name, err := extractNameFromSqsQueueUrl(d.Id())
	if err != nil {
		return err
	}
	d.Set("name", name)

	if attributeOutput.Attributes != nil && len(attributeOutput.Attributes) > 0 {
		attrmap := attributeOutput.Attributes
		resource := *resourceAwsSqsQueue()
		// iKey = internal struct key, oKey = AWS Attribute Map key
		for iKey, oKey := range AttributeMap {
			if attrmap[oKey] != nil {
				if resource.Schema[iKey].Type == schema.TypeInt {
					value, err := strconv.Atoi(*attrmap[oKey])
					if err != nil {
						return err
					}
					d.Set(iKey, value)
					log.Printf("[DEBUG] Reading %s => %s -> %d", iKey, oKey, value)
				} else {
					log.Printf("[DEBUG] Reading %s => %s -> %s", iKey, oKey, *attrmap[oKey])
					d.Set(iKey, *attrmap[oKey])
				}
			}
		}
	}

	return nil
}

func resourceAwsSqsQueueDelete(d *schema.ResourceData, meta interface{}) error {
	sqsconn := meta.(*AWSClient).sqsconn

	log.Printf("[DEBUG] SQS Delete Queue: %s", d.Id())
	_, err := sqsconn.DeleteQueue(&sqs.DeleteQueueInput{
		QueueUrl: aws.String(d.Id()),
	})
	if err != nil {
		return err
	}
	return nil
}

func extractNameFromSqsQueueUrl(queue string) (string, error) {
	//http://sqs.us-west-2.amazonaws.com/123456789012/queueName
	u, err := url.Parse(queue)
	if err != nil {
		return "", err
	}
	segments := strings.Split(u.Path, "/")
	if len(segments) != 3 {
		return "", fmt.Errorf("SQS Url not parsed correctly")
	}

	return segments[2], nil

}
<commit_msg>Add JSON validation to the aws_sqs_queue resource.<commit_after>package aws

import (
	"fmt"
	"log"
	"net/url"
	"strconv"

	"github.com/hashicorp/terraform/helper/schema"

	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/sqs"
)

var AttributeMap = map[string]string{
	"delay_seconds":              "DelaySeconds",
	"max_message_size":           "MaximumMessageSize",
	"message_retention_seconds":  "MessageRetentionPeriod",
	"receive_wait_time_seconds":  "ReceiveMessageWaitTimeSeconds",
	"visibility_timeout_seconds": "VisibilityTimeout",
	"policy":                     "Policy",
	"redrive_policy":             "RedrivePolicy",
	"arn":                        "QueueArn",
}

// A number of these are marked as computed because if you don't
// provide a value, SQS will provide you with defaults (which are the
// default values specified below)
func resourceAwsSqsQueue() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsSqsQueueCreate,
		Read:   resourceAwsSqsQueueRead,
		Update: resourceAwsSqsQueueUpdate,
		Delete: resourceAwsSqsQueueDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"delay_seconds": {
				Type:     schema.TypeInt,
				Optional: true,
				Default:  0,
			},
			"max_message_size": {
				Type:     schema.TypeInt,
				Optional: true,
				Default:  262144,
			},
			"message_retention_seconds": {
				Type:     schema.TypeInt,
				Optional: true,
				Default:  345600,
			},
			"receive_wait_time_seconds": {
				Type:     schema.TypeInt,
				Optional: true,
				Default:  0,
			},
			"visibility_timeout_seconds": {
				Type:     schema.TypeInt,
				Optional: true,
				Default:  30,
			},
			"policy": {
				Type:             schema.TypeString,
				Optional:         true,
				Computed:         true,
				ValidateFunc:     validateJsonString,
				DiffSuppressFunc: suppressEquivalentAwsPolicyDiffs,
			},
			"redrive_policy": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateJsonString,
				StateFunc: func(v interface{}) string {
					json, _ := normalizeJsonString(v)
					return json
				},
			},
			"arn": {
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAwsSqsQueueCreate(d *schema.ResourceData, meta interface{}) error {
	sqsconn := meta.(*AWSClient).sqsconn

	name := d.Get("name").(string)

	log.Printf("[DEBUG] SQS queue create: %s", name)

	req := &sqs.CreateQueueInput{
		QueueName: aws.String(name),
	}

	attributes := make(map[string]*string)

	resource := *resourceAwsSqsQueue()

	for k, s := range resource.Schema {
		if attrKey, ok := AttributeMap[k]; ok {
			if value, ok := d.GetOk(k); ok {
				if s.Type == schema.TypeInt {
					attributes[attrKey] = aws.String(strconv.Itoa(value.(int)))
				} else {
					attributes[attrKey] = aws.String(value.(string))
				}
			}

		}
	}

	if len(attributes) > 0 {
		req.Attributes = attributes
	}

	output, err := sqsconn.CreateQueue(req)
	if err != nil {
		return fmt.Errorf("Error creating SQS queue: %s", err)
	}

	d.SetId(*output.QueueUrl)

	return resourceAwsSqsQueueUpdate(d, meta)
}

func resourceAwsSqsQueueUpdate(d *schema.ResourceData, meta interface{}) error {
	sqsconn := meta.(*AWSClient).sqsconn
	attributes := make(map[string]*string)

	resource := *resourceAwsSqsQueue()

	for k, s := range resource.Schema {
		if attrKey, ok := AttributeMap[k]; ok {
			if d.HasChange(k) {
				log.Printf("[DEBUG] Updating %s", attrKey)
				_, n := d.GetChange(k)
				if s.Type == schema.TypeInt {
					attributes[attrKey] = aws.String(strconv.Itoa(n.(int)))
				} else {
					attributes[attrKey] = aws.String(n.(string))
				}
			}
		}
	}

	if len(attributes) > 0 {
		req := &sqs.SetQueueAttributesInput{
			QueueUrl:   aws.String(d.Id()),
			Attributes: attributes,
		}
		if _, err := sqsconn.SetQueueAttributes(req); err != nil {
			return fmt.Errorf("[ERR] Error updating SQS attributes: %s", err)
		}
	}

	return resourceAwsSqsQueueRead(d, meta)
}

func resourceAwsSqsQueueRead(d *schema.ResourceData, meta interface{}) error {
	sqsconn := meta.(*AWSClient).sqsconn

	attributeOutput, err := sqsconn.GetQueueAttributes(&sqs.GetQueueAttributesInput{
		QueueUrl:       aws.String(d.Id()),
		AttributeNames: []*string{aws.String("All")},
	})

	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			log.Printf("ERROR Found %s", awsErr.Code())
			if "AWS.SimpleQueueService.NonExistentQueue" == awsErr.Code() {
				d.SetId("")
				log.Printf("[DEBUG] SQS Queue (%s) not found", d.Get("name").(string))
				return nil
			}
		}
		return err
	}

	name, err := extractNameFromSqsQueueUrl(d.Id())
	if err != nil {
		return err
	}
	d.Set("name", name)

	if attributeOutput.Attributes != nil && len(attributeOutput.Attributes) > 0 {
		attrmap := attributeOutput.Attributes
		resource := *resourceAwsSqsQueue()
		// iKey = internal struct key, oKey = AWS Attribute Map key
		for iKey, oKey := range AttributeMap {
			if attrmap[oKey] != nil {
				if resource.Schema[iKey].Type == schema.TypeInt {
					value, err := strconv.Atoi(*attrmap[oKey])
					if err != nil {
						return err
					}
					d.Set(iKey, value)
					log.Printf("[DEBUG] Reading %s => %s -> %d", iKey, oKey, value)
				} else {
					log.Printf("[DEBUG] Reading %s => %s -> %s", iKey, oKey, *attrmap[oKey])
					d.Set(iKey, *attrmap[oKey])
				}
			}
		}
	}

	return nil
}

func resourceAwsSqsQueueDelete(d *schema.ResourceData, meta interface{}) error {
	sqsconn := meta.(*AWSClient).sqsconn

	log.Printf("[DEBUG] SQS Delete Queue: %s", d.Id())
	_, err := sqsconn.DeleteQueue(&sqs.DeleteQueueInput{
		QueueUrl: aws.String(d.Id()),
	})
	if err != nil {
		return err
	}
	return nil
}

func extractNameFromSqsQueueUrl(queue string) (string, error) {
	//http://sqs.us-west-2.amazonaws.com/123456789012/queueName
	u, err := url.Parse(queue)
	if err != nil {
		return "", err
	}
	segments := strings.Split(u.Path, "/")
	if len(segments) != 3 {
		return "", fmt.Errorf("SQS Url not parsed correctly")
	}

	return segments[2], nil

}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"
	"testing"
	"math/rand"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccDataflowCreate(t *testing.T) {

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckDataflowDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccDataflow,
				Check: resource.ComposeTestCheckFunc(
					testAccDataflowExists(
						"googlecli_dataflow.foobar"),
				),
			},
		},
	})
}

var disallowedDeletedStates = map[string]bool {
	"JOB_STATE_RUNNING": true,
	"JOB_STATE_UNKNOWN": true,
	"JOB_STATE_FAILED": true,
	"JOB_STATE_UPDATED": true,
}
func testAccCheckDataflowDestroy(s *terraform.State) error {
	for _, rs := range s.RootModule().Resources {
		if rs.Type != "googlecli_dataflow" {
			continue
		}

		jobdesc, err := ReadDataflow(rs.Primary.ID, "hx-test")
		if err != nil {
			return fmt.Errorf("Failed to read dataflow list")
		}

		if jobdesc.CurrentState == "" {
			return fmt.Errorf("Dataflow jobs never started ")
		}

		if _, ok := disallowedDeletedStates[jobdesc.CurrentState]; ok {
			return fmt.Errorf("Dataflow job in disallowed state: %q", jobdesc.CurrentState)
		}
	}

	return nil
}

var disallowedCreatedStates = map[string] bool {
	"JOB_STATE_FAILED": true,
	"JOB_STATE_STOPPED": true,
	"JOB_STATE_UPDATED": true,
	"JOB_STATE_UNKNOWN": true,
}

func testAccDataflowExists(n string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No ID is set")
		}

		jobdesc, err := ReadDataflow(rs.Primary.Attributes["jobids.0"], "hx-test")
		if err != nil {
			return fmt.Errorf("In test: Command line read has errored: %q with rs.Primary hash: %q", err, rs.Primary)
		}

		if jobdesc.CurrentState == "" {
			return fmt.Errorf("Dataflow jobs never started")
		}

		if _, ok := disallowedCreatedStates[jobdesc.CurrentState]; ok {
			return fmt.Errorf("Dataflow job in disallowed state: %q", jobdesc.CurrentState)
		}

		return nil
	}
}

var randDataflowInt = rand.New(rand.NewSource(time.Now().UnixNano())).Int()

var testAccDataflow = fmt.Sprintf(`
resource "google_storage_bucket" "bucket" {
	name = "tf-test-bucket-%d"
	force_destroy = true
}
resource "googlecli_dataflow" "foobar" {
	name = "foobar-%d"
	classpath = "./test-fixtures/google-cloud-dataflow-java-examples-all-bundled-1.1.1-SNAPSHOT.jar"
	class = "com.google.cloud.dataflow.examples.WordCount"
	optional_args = {
		stagingLocation = "gs://${google_storage_bucket.bucket.name}"
		runner = "DataflowPipelineRunner"
	}
}
`, randDataflowInt, randDataflowInt)
<commit_msg>grab the project name from the environment<commit_after>package main

import (
	"os"
	"fmt"
	"time"
	"testing"
	"math/rand"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccDataflowCreate(t *testing.T) {

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckDataflowDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccDataflow,
				Check: resource.ComposeTestCheckFunc(
					testAccDataflowExists(
						"googlecli_dataflow.foobar"),
				),
			},
		},
	})
}

var disallowedDeletedStates = map[string]bool {
	"JOB_STATE_RUNNING": true,
	"JOB_STATE_UNKNOWN": true,
	"JOB_STATE_FAILED": true,
	"JOB_STATE_UPDATED": true,
}

func testAccCheckDataflowDestroy(s *terraform.State) error {
	projectName := os.Getenv("GOOGLE_PROJECT")
	for _, rs := range s.RootModule().Resources {
		if rs.Type != "googlecli_dataflow" {
			continue
		}

		jobdesc, err := ReadDataflow(rs.Primary.ID, projectName)
		if err != nil {
			return fmt.Errorf("Failed to read dataflow list")
		}

		if jobdesc.CurrentState == "" {
			return fmt.Errorf("Dataflow jobs never started ")
		}

		if _, ok := disallowedDeletedStates[jobdesc.CurrentState]; ok {
			return fmt.Errorf("Dataflow job in disallowed state: %q", jobdesc.CurrentState)
		}
	}

	return nil
}

var disallowedCreatedStates = map[string] bool {
	"JOB_STATE_FAILED": true,
	"JOB_STATE_STOPPED": true,
	"JOB_STATE_UPDATED": true,
	"JOB_STATE_UNKNOWN": true,
}

func testAccDataflowExists(n string) resource.TestCheckFunc {
	projectName := os.Getenv("GOOGLE_PROJECT")
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No ID is set")
		}

		jobdesc, err := ReadDataflow(rs.Primary.Attributes["jobids.0"], projectName)
		if err != nil {
			return fmt.Errorf("In test: Command line read has errored: %q with rs.Primary hash: %q", err, rs.Primary)
		}

		if jobdesc.CurrentState == "" {
			return fmt.Errorf("Dataflow jobs never started")
		}

		if _, ok := disallowedCreatedStates[jobdesc.CurrentState]; ok {
			return fmt.Errorf("Dataflow job in disallowed state: %q", jobdesc.CurrentState)
		}

		return nil
	}
}

var randDataflowInt = rand.New(rand.NewSource(time.Now().UnixNano())).Int()

var testAccDataflow = fmt.Sprintf(`
resource "google_storage_bucket" "bucket" {
	name = "tf-test-bucket-%d"
	force_destroy = true
}
resource "googlecli_dataflow" "foobar" {
	name = "foobar-%d"
	classpath = "./test-fixtures/google-cloud-dataflow-java-examples-all-bundled-1.1.1-SNAPSHOT.jar"
	class = "com.google.cloud.dataflow.examples.WordCount"
	optional_args = {
		stagingLocation = "gs://${google_storage_bucket.bucket.name}"
		runner = "DataflowPipelineRunner"
	}
}
`, randDataflowInt, randDataflowInt)
<|endoftext|>


<commit_before>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License.  You may obtain a copy of the License at

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

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

package obcca

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"database/sql"
	"encoding/pem"
	"errors"
	"io/ioutil"
	"math/big"
	"os"
	"strconv"
	"time"

	_ "github.com/mattn/go-sqlite3"
	"golang.org/x/crypto/sha3"
	"github.com/spf13/viper"
)

// CA is the base certificate authority.
//
type CA struct {
	db *sql.DB

	path string

	priv *ecdsa.PrivateKey
	cert *x509.Certificate
	raw []byte
}

// NewCA sets up a new CA.
//
func NewCA(name string) *CA {
	ca := new(CA)
	ca.path = viper.GetString("server.rootpath")+"/"+viper.GetString("server.cadir")

	if _, err := os.Stat(ca.path); err != nil {
		Info.Println("Fresh start; creating databases, key pairs, and certificates.")

		if err := os.MkdirAll(ca.path, 0755); err != nil {
			Panic.Panicln(err)
		}
	}

	// open or create certificate database
	db, err := sql.Open("sqlite3", ca.path+"/"+name+".db")
	if err != nil {
		Panic.Panicln(err)
	}

	if err := db.Ping(); err != nil {
		Panic.Panicln(err)
	}
	if _, err := db.Exec("CREATE TABLE IF NOT EXISTS Certificates (row INTEGER PRIMARY KEY, id VARCHAR(64), timestamp INTEGER, usage INTEGER, cert BLOB, hash BLOB, kdfkey BLOB)"); err != nil {
		Panic.Panicln(err)
	}
	if _, err := db.Exec("CREATE TABLE IF NOT EXISTS Users (row INTEGER PRIMARY KEY, id VARCHAR(64), role INTEGER, token BLOB, state INTEGER, key BLOB)"); err != nil {
		Panic.Panicln(err)
	}
	ca.db = db

	// read or create signing key pair
	priv, err := ca.readCAPrivateKey(name)
	if err != nil {
		priv = ca.createCAKeyPair(name)
	}
	ca.priv = priv

	// read CA certificate, or create a self-signed CA certificate
	raw, err := ca.readCACertificate(name)
	if err != nil {
		raw = ca.createCACertificate(name, &ca.priv.PublicKey)
	}
	cert, err := x509.ParseCertificate(raw)
	if err != nil {
		Panic.Panicln(err)
	}

	ca.raw = raw
	ca.cert = cert

	return ca
}

// Close closes down the CA.
//
func (ca *CA) Close() {
	ca.db.Close()
}

func (ca *CA) createCAKeyPair(name string) *ecdsa.PrivateKey {
	Trace.Println("Creating CA key pair.")

	curve := elliptic.P384()

	priv, err := ecdsa.GenerateKey(curve, rand.Reader)
	if err == nil {
		raw, _ := x509.MarshalECPrivateKey(priv)
		cooked := pem.EncodeToMemory(
			&pem.Block{
				Type:  "ECDSA PRIVATE KEY",
				Bytes: raw,
			})
		err := ioutil.WriteFile(ca.path+"/"+name+".priv", cooked, 0644)
		if err != nil {
			Panic.Panicln(err)
		}

		raw, _ = x509.MarshalPKIXPublicKey(&priv.PublicKey)
		cooked = pem.EncodeToMemory(
			&pem.Block{
				Type:  "ECDSA PUBLIC KEY",
				Bytes: raw,
			})
		err = ioutil.WriteFile(ca.path+"/"+name+".pub", cooked, 0644)
		if err != nil {
			Panic.Panicln(err)
		}
	}
	if err != nil {
		Panic.Panicln(err)
	}

	return priv
}

func (ca *CA) readCAPrivateKey(name string) (*ecdsa.PrivateKey, error) {
	Trace.Println("Reading CA private key.")

	cooked, err := ioutil.ReadFile(ca.path + "/" + name + ".priv")
	if err != nil {
		return nil, err
	}

	block, _ := pem.Decode(cooked)
	return x509.ParseECPrivateKey(block.Bytes)
}

func (ca *CA) createCACertificate(name string, pub *ecdsa.PublicKey) []byte {
	Trace.Println("Creating CA certificate.")

	raw, err := ca.newCertificate(pub, x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign, nil)
	if err != nil {
		Panic.Panicln(err)
	}

	cooked := pem.EncodeToMemory(
		&pem.Block{
			Type:  "CERTIFICATE",
			Bytes: raw,
		})
	err = ioutil.WriteFile(ca.path+"/"+name+".cert", cooked, 0644)
	if err != nil {
		Panic.Panicln(err)
	}

	return raw
}

func (ca *CA) readCACertificate(name string) ([]byte, error) {
	Trace.Println("Reading CA certificate.")

	cooked, err := ioutil.ReadFile(ca.path+"/"+name+".cert")
	if err != nil {
		return nil, err
	}

	block, _ := pem.Decode(cooked)
	return block.Bytes, nil
}

func (ca *CA) createCertificate(id string, pub interface{}, usage x509.KeyUsage, timestamp int64, kdfKey []byte, opt ...pkix.Extension) ([]byte, error) {
	Trace.Println("Creating certificate for "+id+".")

	raw, err := ca.newCertificate(pub, usage, opt)
	if err != nil {
		Error.Println(err)
		return nil, err
	}

	hash := sha3.New384()
	hash.Write(raw)
	if _, err = ca.db.Exec("INSERT INTO Certificates (id, timestamp, usage, cert, hash, kdfkey) VALUES (?, ?, ?, ?, ?, ?)", id, timestamp, usage, raw, hash.Sum(nil), kdfKey); err != nil {
		Error.Println(err)
	}

	return raw, err
}

func (ca *CA) newCertificate(pub interface{}, usage x509.KeyUsage, ext []pkix.Extension) ([]byte, error) {
	notBefore := time.Now()
	notAfter := notBefore.Add(time.Hour * 24 * 90)

	parent := ca.cert
	isCA := parent == nil

	tmpl := x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			CommonName:   "OBC",
			Organization: []string{"IBM"},
			Country:      []string{"US"},
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		SubjectKeyId:       []byte{1, 2, 3, 4},
		SignatureAlgorithm: x509.ECDSAWithSHA384,
		KeyUsage:           usage,

		BasicConstraintsValid: true,
		IsCA: isCA,
	}

	if len(ext) > 0 {
		tmpl.Extensions = ext
		tmpl.ExtraExtensions = ext
	}
	if isCA {
		parent = &tmpl
	}

	raw, err := x509.CreateCertificate(
		rand.Reader,
		&tmpl,
		parent,
		pub,
		ca.priv,
	)
	if isCA && err != nil {
		Panic.Panicln(err)
	}

	return raw, err
}

func (ca *CA) readCertificate(id string, usage x509.KeyUsage) ([]byte, error) {
	Trace.Println("Reading certificate for "+id+".")

	var raw []byte
	err := ca.db.QueryRow("SELECT cert FROM Certificates WHERE id=? AND usage=?", id, usage).Scan(&raw)

	return raw, err
}

func (ca *CA) readCertificate1(id string, ts int64) ([]byte, error) {
	Trace.Println("Reading certificate for "+id+".")

	var raw []byte
	err := ca.db.QueryRow("SELECT cert FROM Certificates WHERE id=? AND timestamp=?", id, ts).Scan(&raw)

	return raw, err
}

func (ca *CA) readCertificates(id string, opt ...int64) (*sql.Rows, error) {
	Trace.Println("Reading certificatess for "+id+".")

	if len(opt) > 0 && opt[0] != 0 {
		return ca.db.Query("SELECT cert, kdfkey FROM Certificates ORDER BY usage WHERE id=? AND timestamp=?", id, opt[0])
	}

	return ca.db.Query("SELECT cert, kdfkey FROM Certificates WHERE id=?", id)
}

func (ca *CA) readCertificateByHash(hash []byte) ([]byte, error) {
	Trace.Println("Reading certificate for hash "+string(hash)+".")

	var raw []byte
	row := ca.db.QueryRow("SELECT cert FROM Certificates WHERE hash=?", hash)
	err := row.Scan(&raw)

	return raw, err
}

func (ca *CA) registerUser(id string, role int, opt ...string) (string, error) {
	Trace.Println("Registering user "+id+" as "+strconv.FormatInt(int64(role), 2)+".")

	var row int
	err := ca.db.QueryRow("SELECT row FROM Users WHERE id=?", id).Scan(&row)
	if err == nil {
		return "", errors.New("user is already registered")
	}

	var tok string
	if len(opt) > 0 && len(opt[0]) > 0 {
		tok = opt[0]
	} else {
		tok = randomString(12)
	}

	_, err = ca.db.Exec("INSERT INTO Users (id, role, token, state) VALUES (?, ?, ?, ?)", id, role, tok, 0)
	if err != nil {
		Error.Println(err)

	}

	return tok, err
}

func (ca *CA) deleteUser(id string) (error) {
	Trace.Println("Deleting user "+id+".")

	var row int
	err := ca.db.QueryRow("SELECT row FROM Users WHERE id=?", id).Scan(&row)
	if err == nil {
		_, err = ca.db.Exec("DELETE FROM Certificates Where id=?", id)
		if err != nil {
			Error.Println(err)
		}

		_, err = ca.db.Exec("DELETE FROM Users WHERE row=?", row)
		if err != nil {
			Error.Println(err)
		}
	}

	return err
}

func (ca *CA) readToken(id string) *sql.Row {
	Trace.Println("Reading token for "+id+".")

	return ca.db.QueryRow("SELECT token, state, key FROM Users WHERE id=?", id)
}

func (ca *CA) readRole(id string) int {
	Trace.Println("Reading role for "+id+".")

	var role int
	err := ca.db.QueryRow("SELECT role FROM Users WHERE id=?", id).Scan(&role)
	if err != nil {
		Panic.Panicln(err)
	}

	return role
}
<commit_msg>Common name set to user identity.<commit_after>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License.  You may obtain a copy of the License at

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

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

package obcca

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"database/sql"
	"encoding/pem"
	"errors"
	"io/ioutil"
	"math/big"
	"os"
	"strconv"
	"time"

	_ "github.com/mattn/go-sqlite3"
	"golang.org/x/crypto/sha3"
	"github.com/spf13/viper"
)

// CA is the base certificate authority.
//
type CA struct {
	db *sql.DB

	path string

	priv *ecdsa.PrivateKey
	cert *x509.Certificate
	raw []byte
}

// NewCA sets up a new CA.
//
func NewCA(name string) *CA {
	ca := new(CA)
	ca.path = viper.GetString("server.rootpath")+"/"+viper.GetString("server.cadir")

	if _, err := os.Stat(ca.path); err != nil {
		Info.Println("Fresh start; creating databases, key pairs, and certificates.")

		if err := os.MkdirAll(ca.path, 0755); err != nil {
			Panic.Panicln(err)
		}
	}

	// open or create certificate database
	db, err := sql.Open("sqlite3", ca.path+"/"+name+".db")
	if err != nil {
		Panic.Panicln(err)
	}

	if err := db.Ping(); err != nil {
		Panic.Panicln(err)
	}
	if _, err := db.Exec("CREATE TABLE IF NOT EXISTS Certificates (row INTEGER PRIMARY KEY, id VARCHAR(64), timestamp INTEGER, usage INTEGER, cert BLOB, hash BLOB, kdfkey BLOB)"); err != nil {
		Panic.Panicln(err)
	}
	if _, err := db.Exec("CREATE TABLE IF NOT EXISTS Users (row INTEGER PRIMARY KEY, id VARCHAR(64), role INTEGER, token BLOB, state INTEGER, key BLOB)"); err != nil {
		Panic.Panicln(err)
	}
	ca.db = db

	// read or create signing key pair
	priv, err := ca.readCAPrivateKey(name)
	if err != nil {
		priv = ca.createCAKeyPair(name)
	}
	ca.priv = priv

	// read CA certificate, or create a self-signed CA certificate
	raw, err := ca.readCACertificate(name)
	if err != nil {
		raw = ca.createCACertificate(name, &ca.priv.PublicKey)
	}
	cert, err := x509.ParseCertificate(raw)
	if err != nil {
		Panic.Panicln(err)
	}

	ca.raw = raw
	ca.cert = cert

	return ca
}

// Close closes down the CA.
//
func (ca *CA) Close() {
	ca.db.Close()
}

func (ca *CA) createCAKeyPair(name string) *ecdsa.PrivateKey {
	Trace.Println("Creating CA key pair.")

	curve := elliptic.P384()

	priv, err := ecdsa.GenerateKey(curve, rand.Reader)
	if err == nil {
		raw, _ := x509.MarshalECPrivateKey(priv)
		cooked := pem.EncodeToMemory(
			&pem.Block{
				Type:  "ECDSA PRIVATE KEY",
				Bytes: raw,
			})
		err := ioutil.WriteFile(ca.path+"/"+name+".priv", cooked, 0644)
		if err != nil {
			Panic.Panicln(err)
		}

		raw, _ = x509.MarshalPKIXPublicKey(&priv.PublicKey)
		cooked = pem.EncodeToMemory(
			&pem.Block{
				Type:  "ECDSA PUBLIC KEY",
				Bytes: raw,
			})
		err = ioutil.WriteFile(ca.path+"/"+name+".pub", cooked, 0644)
		if err != nil {
			Panic.Panicln(err)
		}
	}
	if err != nil {
		Panic.Panicln(err)
	}

	return priv
}

func (ca *CA) readCAPrivateKey(name string) (*ecdsa.PrivateKey, error) {
	Trace.Println("Reading CA private key.")

	cooked, err := ioutil.ReadFile(ca.path + "/" + name + ".priv")
	if err != nil {
		return nil, err
	}

	block, _ := pem.Decode(cooked)
	return x509.ParseECPrivateKey(block.Bytes)
}

func (ca *CA) createCACertificate(name string, pub *ecdsa.PublicKey) []byte {
	Trace.Println("Creating CA certificate.")

	raw, err := ca.newCertificate(name, pub, x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign, nil)
	if err != nil {
		Panic.Panicln(err)
	}

	cooked := pem.EncodeToMemory(
		&pem.Block{
			Type:  "CERTIFICATE",
			Bytes: raw,
		})
	err = ioutil.WriteFile(ca.path+"/"+name+".cert", cooked, 0644)
	if err != nil {
		Panic.Panicln(err)
	}

	return raw
}

func (ca *CA) readCACertificate(name string) ([]byte, error) {
	Trace.Println("Reading CA certificate.")

	cooked, err := ioutil.ReadFile(ca.path+"/"+name+".cert")
	if err != nil {
		return nil, err
	}

	block, _ := pem.Decode(cooked)
	return block.Bytes, nil
}

func (ca *CA) createCertificate(id string, pub interface{}, usage x509.KeyUsage, timestamp int64, kdfKey []byte, opt ...pkix.Extension) ([]byte, error) {
	Trace.Println("Creating certificate for "+id+".")

	raw, err := ca.newCertificate(id, pub, usage, opt)
	if err != nil {
		Error.Println(err)
		return nil, err
	}

	hash := sha3.New384()
	hash.Write(raw)
	if _, err = ca.db.Exec("INSERT INTO Certificates (id, timestamp, usage, cert, hash, kdfkey) VALUES (?, ?, ?, ?, ?, ?)", id, timestamp, usage, raw, hash.Sum(nil), kdfKey); err != nil {
		Error.Println(err)
	}

	return raw, err
}

func (ca *CA) newCertificate(id string, pub interface{}, usage x509.KeyUsage, ext []pkix.Extension) ([]byte, error) {
	notBefore := time.Now()
	notAfter := notBefore.Add(time.Hour * 24 * 90)

	parent := ca.cert
	isCA := parent == nil

	tmpl := x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			CommonName:   id,
			Organization: []string{"IBM"},
			Country:      []string{"US"},
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		SubjectKeyId:       []byte{1, 2, 3, 4},
		SignatureAlgorithm: x509.ECDSAWithSHA384,
		KeyUsage:           usage,

		BasicConstraintsValid: true,
		IsCA: isCA,
	}

	if len(ext) > 0 {
		tmpl.Extensions = ext
		tmpl.ExtraExtensions = ext
	}
	if isCA {
		parent = &tmpl
	}

	raw, err := x509.CreateCertificate(
		rand.Reader,
		&tmpl,
		parent,
		pub,
		ca.priv,
	)
	if isCA && err != nil {
		Panic.Panicln(err)
	}

	return raw, err
}

func (ca *CA) readCertificate(id string, usage x509.KeyUsage) ([]byte, error) {
	Trace.Println("Reading certificate for "+id+".")

	var raw []byte
	err := ca.db.QueryRow("SELECT cert FROM Certificates WHERE id=? AND usage=?", id, usage).Scan(&raw)

	return raw, err
}

func (ca *CA) readCertificate1(id string, ts int64) ([]byte, error) {
	Trace.Println("Reading certificate for "+id+".")

	var raw []byte
	err := ca.db.QueryRow("SELECT cert FROM Certificates WHERE id=? AND timestamp=?", id, ts).Scan(&raw)

	return raw, err
}

func (ca *CA) readCertificates(id string, opt ...int64) (*sql.Rows, error) {
	Trace.Println("Reading certificatess for "+id+".")

	if len(opt) > 0 && opt[0] != 0 {
		return ca.db.Query("SELECT cert, kdfkey FROM Certificates ORDER BY usage WHERE id=? AND timestamp=?", id, opt[0])
	}

	return ca.db.Query("SELECT cert, kdfkey FROM Certificates WHERE id=?", id)
}

func (ca *CA) readCertificateByHash(hash []byte) ([]byte, error) {
	Trace.Println("Reading certificate for hash "+string(hash)+".")

	var raw []byte
	row := ca.db.QueryRow("SELECT cert FROM Certificates WHERE hash=?", hash)
	err := row.Scan(&raw)

	return raw, err
}

func (ca *CA) registerUser(id string, role int, opt ...string) (string, error) {
	Trace.Println("Registering user "+id+" as "+strconv.FormatInt(int64(role), 2)+".")

	var row int
	err := ca.db.QueryRow("SELECT row FROM Users WHERE id=?", id).Scan(&row)
	if err == nil {
		return "", errors.New("user is already registered")
	}

	var tok string
	if len(opt) > 0 && len(opt[0]) > 0 {
		tok = opt[0]
	} else {
		tok = randomString(12)
	}

	_, err = ca.db.Exec("INSERT INTO Users (id, role, token, state) VALUES (?, ?, ?, ?)", id, role, tok, 0)
	if err != nil {
		Error.Println(err)

	}

	return tok, err
}

func (ca *CA) deleteUser(id string) (error) {
	Trace.Println("Deleting user "+id+".")

	var row int
	err := ca.db.QueryRow("SELECT row FROM Users WHERE id=?", id).Scan(&row)
	if err == nil {
		_, err = ca.db.Exec("DELETE FROM Certificates Where id=?", id)
		if err != nil {
			Error.Println(err)
		}

		_, err = ca.db.Exec("DELETE FROM Users WHERE row=?", row)
		if err != nil {
			Error.Println(err)
		}
	}

	return err
}

func (ca *CA) readToken(id string) *sql.Row {
	Trace.Println("Reading token for "+id+".")

	return ca.db.QueryRow("SELECT token, state, key FROM Users WHERE id=?", id)
}

func (ca *CA) readRole(id string) int {
	Trace.Println("Reading role for "+id+".")

	var role int
	err := ca.db.QueryRow("SELECT role FROM Users WHERE id=?", id).Scan(&role)
	if err != nil {
		Panic.Panicln(err)
	}

	return role
}
<|endoftext|>


<commit_before>// Package runner provides a Chrome process runner.
package runner

import (
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"runtime"
	"sync"
	"syscall"

	"github.com/knq/chromedp/client"
)

const (
	// DefaultUserDataDirPrefix is the default user data directory prefix.
	DefaultUserDataDirPrefix = "chromedp-runner.%d."
)

// Runner holds information about a running Chrome process.
type Runner struct {
	opts    map[string]interface{}
	cmd     *exec.Cmd
	waiting bool
	rw      sync.RWMutex
}

// New creates a new Chrome process using the supplied command line options.
func New(opts ...CommandLineOption) (*Runner, error) {
	var err error

	cliOpts := map[string]interface{}{}

	// apply opts
	for _, o := range opts {
		err = o(cliOpts)
		if err != nil {
			return nil, err
		}
	}

	// set default Chrome options if exec-path not provided
	if _, ok := cliOpts["exec-path"]; !ok {
		cliOpts["exec-path"] = findChromePath()
		for k, v := range map[string]interface{}{
			"no-first-run":             true,
			"no-default-browser-check": true,
			"remote-debugging-port":    9222,
		} {
			if _, ok := cliOpts[k]; !ok {
				cliOpts[k] = v
			}
		}
	}

	// add KillProcessGroup and ForceKill if no other cmd opts provided
	if _, ok := cliOpts["cmd-opts"]; !ok {
		for _, o := range []CommandLineOption{KillProcessGroup, ForceKill} {
			err = o(cliOpts)
			if err != nil {
				return nil, err
			}
		}
	}

	return &Runner{
		opts: cliOpts,
	}, nil
}

// cliOptRE is a regular expression to validate a chrome cli option.
var cliOptRE = regexp.MustCompile(`^[a-z0-9\-]+$`)

// buildOpts generates the command line options for Chrome.
func (r *Runner) buildOpts() []string {
	var opts []string

	// process options
	var urlstr string
	for k, v := range r.opts {
		if !cliOptRE.MatchString(k) || v == nil {
			continue
		}

		switch k {
		case "exec-path", "cmd-opts", "process-opts":
			continue

		case "start-url":
			urlstr = v.(string)

		default:
			switch z := v.(type) {
			case bool:
				if z {
					opts = append(opts, "--"+k)
				}

			case string:
				opts = append(opts, "--"+k+"="+z)

			default:
				opts = append(opts, "--"+k+"="+fmt.Sprintf("%v", v))
			}
		}
	}

	if urlstr == "" {
		urlstr = "about:blank"
	}

	return append(opts, urlstr)
}

// Start starts a Chrome process using the specified context. The Chrome
// process can be terminated by closing the passed context.
func (r *Runner) Start(ctxt context.Context) error {
	var err error
	var ok bool

	r.rw.RLock()
	cmd := r.cmd
	r.rw.RUnlock()

	if cmd != nil {
		return errors.New("already started")
	}

	// setup context
	if ctxt == nil {
		ctxt = context.Background()
	}

	// set user data dir, if not provided
	_, ok = r.opts["user-data-dir"]
	if !ok {
		r.opts["user-data-dir"], err = ioutil.TempDir(
			DefaultUserDataTmpDir, fmt.Sprintf(DefaultUserDataDirPrefix, r.Port()),
		)
		if err != nil {
			return err
		}
	}

	// ensure exec-path set
	execPath, ok := r.opts["exec-path"]
	if !ok {
		return errors.New("exec-path command line option not set, or chrome executable not found in $PATH")
	}

	// create cmd
	r.cmd = exec.CommandContext(ctxt, execPath.(string), r.buildOpts()...)

	// apply cmd opts
	if cmdOpts, ok := r.opts["cmd-opts"]; ok {
		for _, co := range cmdOpts.([]func(*exec.Cmd) error) {
			err = co(r.cmd)
			if err != nil {
				return err
			}
		}
	}

	// start process
	err = r.cmd.Start()
	if err != nil {
		return err
	}

	// apply process opts
	if processOpts, ok := r.opts["process-opts"]; ok {
		for _, po := range processOpts.([]func(*os.Process) error) {
			err = po(r.cmd.Process)
			if err != nil {
				// TODO: do something better here, as we want to kill
				// the child process, do cleanup, etc.
				panic(err)
				//return err
			}
		}
	}

	return nil
}

// Shutdown shuts down and all the Chrome tabs.
func (r *Runner) Shutdown(ctxt context.Context, opts ...client.Option) error {
	var err error

	cl := r.Client(opts...)

	targets, err := cl.ListPageTargets(ctxt)
	if err != nil {
		return err
	}

	var wg sync.WaitGroup
	errs := make([]error, len(targets))
	for i, t := range targets {
		wg.Add(1)
		go func(wg *sync.WaitGroup, i int, t client.Target) {
			defer wg.Done()
			errs[i] = cl.CloseTarget(ctxt, t)
		}(&wg, i, t)
	}
	wg.Wait()

	for _, e := range errs {
		if e != nil {
			return e
		}
	}

	// osx applications do not automatically exit when all windows (ie, tabs)
	// closed, so send SIGTERM.
	//
	// TODO: add other behavior here for more process options on shutdown?
	if runtime.GOOS == "darwin" {
		return r.cmd.Process.Signal(syscall.SIGTERM)
	}

	return nil
}

// Wait waits for the previously started Chrome process to terminate, returning
// any encountered error.
func (r *Runner) Wait() error {
	r.rw.RLock()
	waiting := r.waiting
	r.rw.RUnlock()

	if waiting {
		return errors.New("already waiting")
	}

	r.rw.Lock()
	r.waiting = true
	r.rw.Unlock()

	defer func() {
		r.rw.Lock()
		r.waiting = false
		r.rw.Unlock()
	}()

	return r.cmd.Wait()
}

// Port returns the port the process was launched with.
func (r *Runner) Port() int {
	var port interface{}
	var ok bool
	port, ok = r.opts["remote-debugging-port"]
	if !ok {
		port, ok = r.opts["port"]
	}
	if !ok {
		panic("expected either remote-debugging-port or port to be specified in command line options")
	}

	var p int
	p, ok = port.(int)
	if !ok {
		panic("expected port to be type int")
	}

	return p
}

// Client returns a Chrome Debugging Protocol client for the running Chrome
// process.
func (r *Runner) Client(opts ...client.Option) *client.Client {
	return client.New(append(opts,
		client.URL(fmt.Sprintf("http://localhost:%d/json", r.Port())),
	)...)
}

// WatchPageTargets returns a channel that will receive new page targets as
// they are created.
func (r *Runner) WatchPageTargets(ctxt context.Context, opts ...client.Option) <-chan client.Target {
	return r.Client(opts...).WatchPageTargets(ctxt)
}

// Run starts a new Chrome process runner, using the provided context and
// command line options.
func Run(ctxt context.Context, opts ...CommandLineOption) (*Runner, error) {
	var err error

	// create
	r, err := New(opts...)
	if err != nil {
		return nil, err
	}

	// start
	err = r.Start(ctxt)
	if err != nil {
		return nil, err
	}

	return r, nil
}

// CommandLineOption is a Chrome command line option.
//
// see: http://peter.sh/experiments/chromium-command-line-switches/
type CommandLineOption func(map[string]interface{}) error

// Flag is a generic Chrome command line option to pass a name=value flag to
// Chrome.
func Flag(name string, value interface{}) CommandLineOption {
	return func(m map[string]interface{}) error {
		m[name] = value
		return nil
	}
}

// ExecPath sets the Chrome executable path.
func ExecPath(path string) CommandLineOption {
	return Flag("exec-path", path)
}

// Port is the Chrome command line option to set the remote debugging port.
func Port(port int) CommandLineOption {
	return Flag("remote-debugging-port", port)
}

// UserDataDir is the Chrome command line option to set the user data dir.
func UserDataDir(dir string) CommandLineOption {
	return Flag("user-data-dir", dir)
}

// StartURL is the Chrome command line option to set the initial URL.
func StartURL(urlstr string) CommandLineOption {
	return Flag("start-url", urlstr)
}

// Proxy is the Chrome command line option to set the proxy used.
func Proxy(proxy string) CommandLineOption {
	return Flag("proxy-server", proxy)
}

// WindowSize is the Chrome command line option to set the initial window size.
func WindowSize(width, height int) CommandLineOption {
	return Flag("window-size", fmt.Sprintf("%d,%d", width, height))
}

// UserAgent is the Chrome command line option to set the default User-Agent
// header.
func UserAgent(userAgent string) CommandLineOption {
	return Flag("user-agent", userAgent)
}

// CmdOpt is a Chrome command line option to modify the underlying exec.Cmd
// prior to invocation.
func CmdOpt(o func(*exec.Cmd) error) CommandLineOption {
	return func(m map[string]interface{}) error {
		var opts []func(*exec.Cmd) error

		if e, ok := m["cmd-opts"]; ok {
			opts, ok = e.([]func(*exec.Cmd) error)
			if !ok {
				return errors.New("cmd-opts is in invalid state")
			}
		}

		m["cmd-opts"] = append(opts, o)

		return nil
	}
}

// ProcessOpt is a Chrome command line option to modify the child os.Process
// after started exec.Cmd.Start.
func ProcessOpt(o func(*os.Process) error) CommandLineOption {
	return func(m map[string]interface{}) error {
		var opts []func(*os.Process) error

		if e, ok := m["process-opts"]; ok {
			opts, ok = e.([]func(*os.Process) error)
			if !ok {
				return errors.New("process-opts is in invalid state")
			}
		}

		m["process-opts"] = append(opts, o)

		return nil
	}
}
<commit_msg>Cleaning up chromedp/runner options<commit_after>// Package runner provides a Chrome process runner.
package runner

import (
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"runtime"
	"sync"
	"syscall"

	"github.com/knq/chromedp/client"
)

const (
	// DefaultUserDataDirPrefix is the default user data directory prefix.
	DefaultUserDataDirPrefix = "chromedp-runner.%d."
)

// Runner holds information about a running Chrome process.
type Runner struct {
	opts    map[string]interface{}
	cmd     *exec.Cmd
	waiting bool
	rw      sync.RWMutex
}

// New creates a new Chrome process using the supplied command line options.
func New(opts ...CommandLineOption) (*Runner, error) {
	var err error

	cliOpts := map[string]interface{}{}

	// apply opts
	for _, o := range opts {
		err = o(cliOpts)
		if err != nil {
			return nil, err
		}
	}

	// set default Chrome options if exec-path not provided
	if _, ok := cliOpts["exec-path"]; !ok {
		cliOpts["exec-path"] = findChromePath()
		for k, v := range map[string]interface{}{
			"no-first-run":             true,
			"no-default-browser-check": true,
			"remote-debugging-port":    9222,
		} {
			if _, ok := cliOpts[k]; !ok {
				cliOpts[k] = v
			}
		}
	}

	// add KillProcessGroup and ForceKill if no other cmd opts provided
	if _, ok := cliOpts["cmd-opts"]; !ok {
		for _, o := range []CommandLineOption{KillProcessGroup, ForceKill} {
			err = o(cliOpts)
			if err != nil {
				return nil, err
			}
		}
	}

	return &Runner{
		opts: cliOpts,
	}, nil
}

// cliOptRE is a regular expression to validate a chrome cli option.
var cliOptRE = regexp.MustCompile(`^[a-z0-9\-]+$`)

// buildOpts generates the command line options for Chrome.
func (r *Runner) buildOpts() []string {
	var opts []string

	// process options
	var urlstr string
	for k, v := range r.opts {
		if !cliOptRE.MatchString(k) || v == nil {
			continue
		}

		switch k {
		case "exec-path", "cmd-opts", "process-opts":
			continue

		case "start-url":
			urlstr = v.(string)

		default:
			switch z := v.(type) {
			case bool:
				if z {
					opts = append(opts, "--"+k)
				}

			case string:
				opts = append(opts, "--"+k+"="+z)

			default:
				opts = append(opts, "--"+k+"="+fmt.Sprintf("%v", v))
			}
		}
	}

	if urlstr == "" {
		urlstr = "about:blank"
	}

	return append(opts, urlstr)
}

// Start starts a Chrome process using the specified context. The Chrome
// process can be terminated by closing the passed context.
func (r *Runner) Start(ctxt context.Context) error {
	var err error
	var ok bool

	r.rw.RLock()
	cmd := r.cmd
	r.rw.RUnlock()

	if cmd != nil {
		return errors.New("already started")
	}

	// setup context
	if ctxt == nil {
		ctxt = context.Background()
	}

	// set user data dir, if not provided
	_, ok = r.opts["user-data-dir"]
	if !ok {
		r.opts["user-data-dir"], err = ioutil.TempDir(
			DefaultUserDataTmpDir, fmt.Sprintf(DefaultUserDataDirPrefix, r.Port()),
		)
		if err != nil {
			return err
		}
	}

	// ensure exec-path set
	execPath, ok := r.opts["exec-path"]
	if !ok {
		return errors.New("exec-path command line option not set, or chrome executable not found in $PATH")
	}

	// create cmd
	r.cmd = exec.CommandContext(ctxt, execPath.(string), r.buildOpts()...)

	// apply cmd opts
	if cmdOpts, ok := r.opts["cmd-opts"]; ok {
		for _, co := range cmdOpts.([]func(*exec.Cmd) error) {
			err = co(r.cmd)
			if err != nil {
				return err
			}
		}
	}

	// start process
	err = r.cmd.Start()
	if err != nil {
		return err
	}

	// apply process opts
	if processOpts, ok := r.opts["process-opts"]; ok {
		for _, po := range processOpts.([]func(*os.Process) error) {
			err = po(r.cmd.Process)
			if err != nil {
				// TODO: do something better here, as we want to kill
				// the child process, do cleanup, etc.
				panic(err)
				//return err
			}
		}
	}

	return nil
}

// Shutdown shuts down and all the Chrome tabs.
func (r *Runner) Shutdown(ctxt context.Context, opts ...client.Option) error {
	var err error

	cl := r.Client(opts...)

	targets, err := cl.ListPageTargets(ctxt)
	if err != nil {
		return err
	}

	var wg sync.WaitGroup
	errs := make([]error, len(targets))
	for i, t := range targets {
		wg.Add(1)
		go func(wg *sync.WaitGroup, i int, t client.Target) {
			defer wg.Done()
			errs[i] = cl.CloseTarget(ctxt, t)
		}(&wg, i, t)
	}
	wg.Wait()

	for _, e := range errs {
		if e != nil {
			return e
		}
	}

	// osx applications do not automatically exit when all windows (ie, tabs)
	// closed, so send SIGTERM.
	//
	// TODO: add other behavior here for more process options on shutdown?
	if runtime.GOOS == "darwin" && r.cmd != nil && r.cmd.Process != nil {
		return r.cmd.Process.Signal(syscall.SIGTERM)
	}

	return nil
}

// Wait waits for the previously started Chrome process to terminate, returning
// any encountered error.
func (r *Runner) Wait() error {
	r.rw.RLock()
	waiting := r.waiting
	r.rw.RUnlock()

	if waiting {
		return errors.New("already waiting")
	}

	r.rw.Lock()
	r.waiting = true
	r.rw.Unlock()

	defer func() {
		r.rw.Lock()
		r.waiting = false
		r.rw.Unlock()
	}()

	return r.cmd.Wait()
}

// Port returns the port the process was launched with.
func (r *Runner) Port() int {
	var port interface{}
	var ok bool
	port, ok = r.opts["remote-debugging-port"]
	if !ok {
		port, ok = r.opts["port"]
	}
	if !ok {
		panic("expected either remote-debugging-port or port to be specified in command line options")
	}

	var p int
	p, ok = port.(int)
	if !ok {
		panic("expected port to be type int")
	}

	return p
}

// Client returns a Chrome Debugging Protocol client for the running Chrome
// process.
func (r *Runner) Client(opts ...client.Option) *client.Client {
	return client.New(append(opts,
		client.URL(fmt.Sprintf("http://localhost:%d/json", r.Port())),
	)...)
}

// WatchPageTargets returns a channel that will receive new page targets as
// they are created.
func (r *Runner) WatchPageTargets(ctxt context.Context, opts ...client.Option) <-chan client.Target {
	return r.Client(opts...).WatchPageTargets(ctxt)
}

// Run starts a new Chrome process runner, using the provided context and
// command line options.
func Run(ctxt context.Context, opts ...CommandLineOption) (*Runner, error) {
	var err error

	// create
	r, err := New(opts...)
	if err != nil {
		return nil, err
	}

	// start
	err = r.Start(ctxt)
	if err != nil {
		return nil, err
	}

	return r, nil
}

// CommandLineOption is a Chrome command line option.
//
// see: http://peter.sh/experiments/chromium-command-line-switches/
type CommandLineOption func(map[string]interface{}) error

// Flag is a generic Chrome command line option to pass a name=value flag to
// Chrome.
func Flag(name string, value interface{}) CommandLineOption {
	return func(m map[string]interface{}) error {
		m[name] = value
		return nil
	}
}

// Path sets the path to the Chrome executable and sets default run options for
// Chrome. This will also set the remote debugging port to 9222, and disable
// the first run / default browser check.
//
// Note: use ExecPath if you do not want to set other options.
func Path(path string) CommandLineOption {
	return func(m map[string]interface{}) error {
		m["exec-path"] = path
		m["no-first-run"] = true
		m["no-default-browser-check"] = true
		m["remote-debugging-port"] = 9222
		return nil
	}
}

// ExecPath is a Chrome command line option to set the exec path.
func ExecPath(path string) CommandLineOption {
	return Flag("exec-path", path)
}

// Port is the Chrome command line option to set the remote debugging port.
func Port(port int) CommandLineOption {
	return Flag("remote-debugging-port", port)
}

// UserDataDir is the Chrome command line option to set the user data dir.
func UserDataDir(dir string) CommandLineOption {
	return Flag("user-data-dir", dir)
}

// StartURL is the Chrome command line option to set the initial URL.
func StartURL(urlstr string) CommandLineOption {
	return Flag("start-url", urlstr)
}

// Proxy is the Chrome command line option to set the outbound proxy.
func Proxy(proxy string) CommandLineOption {
	return Flag("proxy-server", proxy)
}

// WindowSize is the Chrome command line option to set the initial window size.
func WindowSize(width, height int) CommandLineOption {
	return Flag("window-size", fmt.Sprintf("%d,%d", width, height))
}

// UserAgent is the Chrome command line option to set the default User-Agent
// header.
func UserAgent(userAgent string) CommandLineOption {
	return Flag("user-agent", userAgent)
}

// CmdOpt is a Chrome command line option to modify the underlying exec.Cmd
// prior to invocation.
func CmdOpt(o func(*exec.Cmd) error) CommandLineOption {
	return func(m map[string]interface{}) error {
		var opts []func(*exec.Cmd) error

		if e, ok := m["cmd-opts"]; ok {
			opts, ok = e.([]func(*exec.Cmd) error)
			if !ok {
				return errors.New("cmd-opts is in invalid state")
			}
		}

		m["cmd-opts"] = append(opts, o)

		return nil
	}
}

// ProcessOpt is a Chrome command line option to modify the child os.Process
// after started exec.Cmd.Start.
func ProcessOpt(o func(*os.Process) error) CommandLineOption {
	return func(m map[string]interface{}) error {
		var opts []func(*os.Process) error

		if e, ok := m["process-opts"]; ok {
			opts, ok = e.([]func(*os.Process) error)
			if !ok {
				return errors.New("process-opts is in invalid state")
			}
		}

		m["process-opts"] = append(opts, o)

		return nil
	}
}
<|endoftext|>