<commit_before>package shared

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/gorilla/websocket"

	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/logger"
)

func RFC3493Dialer(network, address string) (net.Conn, error) {
	host, port, err := net.SplitHostPort(address)
	if err != nil {
		return nil, err
	}

	addrs, err := net.LookupHost(host)
	if err != nil {
		return nil, err
	}
	for _, a := range addrs {
		c, err := net.DialTimeout(network, net.JoinHostPort(a, port), 10*time.Second)
		if err != nil {
			continue
		}
		if tc, ok := c.(*net.TCPConn); ok {
			tc.SetKeepAlive(true)
			tc.SetKeepAlivePeriod(3 * time.Second)
		}
		return c, err
	}
	return nil, fmt.Errorf("Unable to connect to: " + address)
}

// InitTLSConfig returns a tls.Config populated with default encryption
// parameters. This is used as baseline config for both client and server
// certificates used by LXD.
func InitTLSConfig() *tls.Config {
	return &tls.Config{
		MinVersion: tls.VersionTLS12,
		MaxVersion: tls.VersionTLS12,
		CipherSuites: []uint16{
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA,
			tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA,
			tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA,
		},
		PreferServerCipherSuites: true,
	}
}

func finalizeTLSConfig(tlsConfig *tls.Config, tlsRemoteCert *x509.Certificate) {
	// Trusted certificates
	if tlsRemoteCert != nil {
		caCertPool := tlsConfig.RootCAs
		if caCertPool == nil {
			caCertPool, _ = x509.SystemCertPool()
			if caCertPool == nil {
				caCertPool = x509.NewCertPool()
			}
		}

		// Make it a valid RootCA
		tlsRemoteCert.IsCA = true
		tlsRemoteCert.KeyUsage = x509.KeyUsageCertSign

		// Setup the pool
		caCertPool.AddCert(tlsRemoteCert)
		tlsConfig.RootCAs = caCertPool

		// Set the ServerName
		if tlsRemoteCert.DNSNames != nil {
			tlsConfig.ServerName = tlsRemoteCert.DNSNames[0]
		}
	}

	tlsConfig.BuildNameToCertificate()
}

func GetTLSConfig(tlsClientCertFile string, tlsClientKeyFile string, tlsClientCAFile string, tlsRemoteCert *x509.Certificate) (*tls.Config, error) {
	tlsConfig := InitTLSConfig()

	// Client authentication
	if tlsClientCertFile != "" && tlsClientKeyFile != "" {
		cert, err := tls.LoadX509KeyPair(tlsClientCertFile, tlsClientKeyFile)
		if err != nil {
			return nil, err
		}

		tlsConfig.Certificates = []tls.Certificate{cert}
	}

	if tlsClientCAFile != "" {
		caCertificates, err := ioutil.ReadFile(tlsClientCAFile)
		if err != nil {
			return nil, err
		}

		caPool := x509.NewCertPool()
		caPool.AppendCertsFromPEM(caCertificates)

		tlsConfig.RootCAs = caPool
	}

	finalizeTLSConfig(tlsConfig, tlsRemoteCert)
	return tlsConfig, nil
}

func GetTLSConfigMem(tlsClientCert string, tlsClientKey string, tlsClientCA string, tlsRemoteCertPEM string, insecureSkipVerify bool) (*tls.Config, error) {
	tlsConfig := InitTLSConfig()
	tlsConfig.InsecureSkipVerify = insecureSkipVerify
	// Client authentication
	if tlsClientCert != "" && tlsClientKey != "" {
		cert, err := tls.X509KeyPair([]byte(tlsClientCert), []byte(tlsClientKey))
		if err != nil {
			return nil, err
		}

		tlsConfig.Certificates = []tls.Certificate{cert}
	}

	var tlsRemoteCert *x509.Certificate
	if tlsRemoteCertPEM != "" {
		// Ignore any content outside of the PEM bytes we care about
		certBlock, _ := pem.Decode([]byte(tlsRemoteCertPEM))
		if certBlock == nil {
			return nil, fmt.Errorf("Invalid remote certificate")
		}

		var err error
		tlsRemoteCert, err = x509.ParseCertificate(certBlock.Bytes)
		if err != nil {
			return nil, err
		}
	}

	if tlsClientCA != "" {
		caPool := x509.NewCertPool()
		caPool.AppendCertsFromPEM([]byte(tlsClientCA))

		tlsConfig.RootCAs = caPool
	}

	finalizeTLSConfig(tlsConfig, tlsRemoteCert)

	return tlsConfig, nil
}

func IsLoopback(iface *net.Interface) bool {
	return int(iface.Flags&net.FlagLoopback) > 0
}

func WebsocketSendStream(conn *websocket.Conn, r io.Reader, bufferSize int) chan bool {
	ch := make(chan bool)

	if r == nil {
		close(ch)
		return ch
	}

	go func(conn *websocket.Conn, r io.Reader) {
		in := ReaderToChannel(r, bufferSize)
		for {
			buf, ok := <-in
			if !ok {
				break
			}

			w, err := conn.NextWriter(websocket.BinaryMessage)
			if err != nil {
				logger.Debugf("Got error getting next writer %s", err)
				break
			}

			_, err = w.Write(buf)
			w.Close()
			if err != nil {
				logger.Debugf("Got err writing %s", err)
				break
			}
		}
		conn.WriteMessage(websocket.TextMessage, []byte{})
		ch <- true
	}(conn, r)

	return ch
}

func WebsocketRecvStream(w io.Writer, conn *websocket.Conn) chan bool {
	ch := make(chan bool)

	go func(w io.Writer, conn *websocket.Conn) {
		for {
			mt, r, err := conn.NextReader()
			if mt == websocket.CloseMessage {
				logger.Debugf("Got close message for reader")
				break
			}

			if mt == websocket.TextMessage {
				logger.Debugf("got message barrier")
				break
			}

			if err != nil {
				logger.Debugf("Got error getting next reader %s, %s", err, w)
				break
			}

			buf, err := ioutil.ReadAll(r)
			if err != nil {
				logger.Debugf("Got error writing to writer %s", err)
				break
			}

			if w == nil {
				continue
			}

			i, err := w.Write(buf)
			if i != len(buf) {
				logger.Debugf("Didn't write all of buf")
				break
			}
			if err != nil {
				logger.Debugf("Error writing buf %s", err)
				break
			}
		}
		ch <- true
	}(w, conn)

	return ch
}

func WebsocketProxy(source *websocket.Conn, target *websocket.Conn) chan bool {
	forward := func(in *websocket.Conn, out *websocket.Conn, ch chan bool) {
		for {
			mt, r, err := in.NextReader()
			if err != nil {
				break
			}

			w, err := out.NextWriter(mt)
			if err != nil {
				break
			}

			_, err = io.Copy(w, r)
			w.Close()
			if err != nil {
				break
			}
		}

		ch <- true
	}

	chSend := make(chan bool)
	go forward(source, target, chSend)

	chRecv := make(chan bool)
	go forward(target, source, chRecv)

	ch := make(chan bool)
	go func() {
		select {
		case <-chSend:
		case <-chRecv:
		}

		source.Close()
		target.Close()

		ch <- true
	}()

	return ch
}

func defaultReader(conn *websocket.Conn, r io.ReadCloser, readDone chan<- bool) {
	/* For now, we don't need to adjust buffer sizes in
	* WebsocketMirror, since it's used for interactive things like
	* exec.
	 */
	in := ReaderToChannel(r, -1)
	for {
		buf, ok := <-in
		if !ok {
			r.Close()
			logger.Debugf("sending write barrier")
			conn.WriteMessage(websocket.TextMessage, []byte{})
			readDone <- true
			return
		}
		w, err := conn.NextWriter(websocket.BinaryMessage)
		if err != nil {
			logger.Debugf("Got error getting next writer %s", err)
			break
		}

		_, err = w.Write(buf)
		w.Close()
		if err != nil {
			logger.Debugf("Got err writing %s", err)
			break
		}
	}
	closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
	conn.WriteMessage(websocket.CloseMessage, closeMsg)
	readDone <- true
	r.Close()
}

func defaultWriter(conn *websocket.Conn, w io.WriteCloser, writeDone chan<- bool) {
	for {
		mt, r, err := conn.NextReader()
		if err != nil {
			logger.Debugf("Got error getting next reader %s", err)
			break
		}

		if mt == websocket.CloseMessage {
			logger.Debugf("Got close message for reader")
			break
		}

		if mt == websocket.TextMessage {
			logger.Debugf("Got message barrier, resetting stream")
			break
		}

		buf, err := ioutil.ReadAll(r)
		if err != nil {
			logger.Debugf("Got error writing to writer %s", err)
			break
		}
		i, err := w.Write(buf)
		if i != len(buf) {
			logger.Debugf("Didn't write all of buf")
			break
		}
		if err != nil {
			logger.Debugf("Error writing buf %s", err)
			break
		}
	}
	writeDone <- true
	w.Close()
}

// WebsocketMirror allows mirroring a reader to a websocket and taking the
// result and writing it to a writer. This function allows for multiple
// mirrorings and correctly negotiates stream endings. However, it means any
// websocket.Conns passed to it are live when it returns, and must be closed
// explicitly.
type WebSocketMirrorReader func(conn *websocket.Conn, r io.ReadCloser, readDone chan<- bool)
type WebSocketMirrorWriter func(conn *websocket.Conn, w io.WriteCloser, writeDone chan<- bool)

func WebsocketMirror(conn *websocket.Conn, w io.WriteCloser, r io.ReadCloser, Reader WebSocketMirrorReader, Writer WebSocketMirrorWriter) (chan bool, chan bool) {
	readDone := make(chan bool, 1)
	writeDone := make(chan bool, 1)

	ReadFunc := Reader
	if ReadFunc == nil {
		ReadFunc = defaultReader
	}

	WriteFunc := Writer
	if WriteFunc == nil {
		WriteFunc = defaultWriter
	}

	go ReadFunc(conn, r, readDone)
	go WriteFunc(conn, w, writeDone)

	return readDone, writeDone
}

func WebsocketConsoleMirror(conn *websocket.Conn, w io.WriteCloser, r io.ReadCloser) (chan bool, chan bool) {
	readDone := make(chan bool, 1)
	writeDone := make(chan bool, 1)

	go defaultWriter(conn, w, writeDone)

	go func(conn *websocket.Conn, r io.ReadCloser) {
		in := ReaderToChannel(r, -1)
		for {
			buf, ok := <-in
			if !ok {
				r.Close()
				logger.Debugf("sending write barrier")
				conn.WriteMessage(websocket.TextMessage, []byte{})
				readDone <- true
				return
			}
			w, err := conn.NextWriter(websocket.BinaryMessage)
			if err != nil {
				logger.Debugf("Got error getting next writer %s", err)
				break
			}

			_, err = w.Write(buf)
			w.Close()
			if err != nil {
				logger.Debugf("Got err writing %s", err)
				break
			}
		}
		closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
		conn.WriteMessage(websocket.CloseMessage, closeMsg)
		readDone <- true
		r.Close()
	}(conn, r)

	return readDone, writeDone
}

var WebsocketUpgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

// AllocatePort asks the kernel for a free open port that is ready to use
func AllocatePort() (int, error) {
	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
	if err != nil {
		return -1, err
	}

	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return -1, err
	}
	defer l.Close()
	return l.Addr().(*net.TCPAddr).Port, nil
}

func NetworkGetCounters(ifName string) api.NetworkStateCounters {
	counters := api.NetworkStateCounters{}
	// Get counters
	content, err := ioutil.ReadFile("/proc/net/dev")
	if err == nil {
		for _, line := range strings.Split(string(content), "\n") {
			fields := strings.Fields(line)

			if len(fields) != 17 {
				continue
			}

			intName := strings.TrimSuffix(fields[0], ":")
			if intName != ifName {
				continue
			}

			rxBytes, err := strconv.ParseInt(fields[1], 10, 64)
			if err != nil {
				continue
			}

			rxPackets, err := strconv.ParseInt(fields[2], 10, 64)
			if err != nil {
				continue
			}

			txBytes, err := strconv.ParseInt(fields[9], 10, 64)
			if err != nil {
				continue
			}

			txPackets, err := strconv.ParseInt(fields[10], 10, 64)
			if err != nil {
				continue
			}

			counters.BytesSent = txBytes
			counters.BytesReceived = rxBytes
			counters.PacketsSent = txPackets
			counters.PacketsReceived = rxPackets
			break
		}
	}

	return counters
}
<commit_msg>shared/network: Allow TLS1.3<commit_after>package shared

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/gorilla/websocket"

	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/logger"
)

func RFC3493Dialer(network, address string) (net.Conn, error) {
	host, port, err := net.SplitHostPort(address)
	if err != nil {
		return nil, err
	}

	addrs, err := net.LookupHost(host)
	if err != nil {
		return nil, err
	}
	for _, a := range addrs {
		c, err := net.DialTimeout(network, net.JoinHostPort(a, port), 10*time.Second)
		if err != nil {
			continue
		}
		if tc, ok := c.(*net.TCPConn); ok {
			tc.SetKeepAlive(true)
			tc.SetKeepAlivePeriod(3 * time.Second)
		}
		return c, err
	}
	return nil, fmt.Errorf("Unable to connect to: " + address)
}

// InitTLSConfig returns a tls.Config populated with default encryption
// parameters. This is used as baseline config for both client and server
// certificates used by LXD.
func InitTLSConfig() *tls.Config {
	return &tls.Config{
		MinVersion: tls.VersionTLS12,
		CipherSuites: []uint16{
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA,
			tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA,
			tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA,
		},
		PreferServerCipherSuites: true,
	}
}

func finalizeTLSConfig(tlsConfig *tls.Config, tlsRemoteCert *x509.Certificate) {
	// Trusted certificates
	if tlsRemoteCert != nil {
		caCertPool := tlsConfig.RootCAs
		if caCertPool == nil {
			caCertPool, _ = x509.SystemCertPool()
			if caCertPool == nil {
				caCertPool = x509.NewCertPool()
			}
		}

		// Make it a valid RootCA
		tlsRemoteCert.IsCA = true
		tlsRemoteCert.KeyUsage = x509.KeyUsageCertSign

		// Setup the pool
		caCertPool.AddCert(tlsRemoteCert)
		tlsConfig.RootCAs = caCertPool

		// Set the ServerName
		if tlsRemoteCert.DNSNames != nil {
			tlsConfig.ServerName = tlsRemoteCert.DNSNames[0]
		}
	}

	tlsConfig.BuildNameToCertificate()
}

func GetTLSConfig(tlsClientCertFile string, tlsClientKeyFile string, tlsClientCAFile string, tlsRemoteCert *x509.Certificate) (*tls.Config, error) {
	tlsConfig := InitTLSConfig()

	// Client authentication
	if tlsClientCertFile != "" && tlsClientKeyFile != "" {
		cert, err := tls.LoadX509KeyPair(tlsClientCertFile, tlsClientKeyFile)
		if err != nil {
			return nil, err
		}

		tlsConfig.Certificates = []tls.Certificate{cert}
	}

	if tlsClientCAFile != "" {
		caCertificates, err := ioutil.ReadFile(tlsClientCAFile)
		if err != nil {
			return nil, err
		}

		caPool := x509.NewCertPool()
		caPool.AppendCertsFromPEM(caCertificates)

		tlsConfig.RootCAs = caPool
	}

	finalizeTLSConfig(tlsConfig, tlsRemoteCert)
	return tlsConfig, nil
}

func GetTLSConfigMem(tlsClientCert string, tlsClientKey string, tlsClientCA string, tlsRemoteCertPEM string, insecureSkipVerify bool) (*tls.Config, error) {
	tlsConfig := InitTLSConfig()
	tlsConfig.InsecureSkipVerify = insecureSkipVerify
	// Client authentication
	if tlsClientCert != "" && tlsClientKey != "" {
		cert, err := tls.X509KeyPair([]byte(tlsClientCert), []byte(tlsClientKey))
		if err != nil {
			return nil, err
		}

		tlsConfig.Certificates = []tls.Certificate{cert}
	}

	var tlsRemoteCert *x509.Certificate
	if tlsRemoteCertPEM != "" {
		// Ignore any content outside of the PEM bytes we care about
		certBlock, _ := pem.Decode([]byte(tlsRemoteCertPEM))
		if certBlock == nil {
			return nil, fmt.Errorf("Invalid remote certificate")
		}

		var err error
		tlsRemoteCert, err = x509.ParseCertificate(certBlock.Bytes)
		if err != nil {
			return nil, err
		}
	}

	if tlsClientCA != "" {
		caPool := x509.NewCertPool()
		caPool.AppendCertsFromPEM([]byte(tlsClientCA))

		tlsConfig.RootCAs = caPool
	}

	finalizeTLSConfig(tlsConfig, tlsRemoteCert)

	return tlsConfig, nil
}

func IsLoopback(iface *net.Interface) bool {
	return int(iface.Flags&net.FlagLoopback) > 0
}

func WebsocketSendStream(conn *websocket.Conn, r io.Reader, bufferSize int) chan bool {
	ch := make(chan bool)

	if r == nil {
		close(ch)
		return ch
	}

	go func(conn *websocket.Conn, r io.Reader) {
		in := ReaderToChannel(r, bufferSize)
		for {
			buf, ok := <-in
			if !ok {
				break
			}

			w, err := conn.NextWriter(websocket.BinaryMessage)
			if err != nil {
				logger.Debugf("Got error getting next writer %s", err)
				break
			}

			_, err = w.Write(buf)
			w.Close()
			if err != nil {
				logger.Debugf("Got err writing %s", err)
				break
			}
		}
		conn.WriteMessage(websocket.TextMessage, []byte{})
		ch <- true
	}(conn, r)

	return ch
}

func WebsocketRecvStream(w io.Writer, conn *websocket.Conn) chan bool {
	ch := make(chan bool)

	go func(w io.Writer, conn *websocket.Conn) {
		for {
			mt, r, err := conn.NextReader()
			if mt == websocket.CloseMessage {
				logger.Debugf("Got close message for reader")
				break
			}

			if mt == websocket.TextMessage {
				logger.Debugf("got message barrier")
				break
			}

			if err != nil {
				logger.Debugf("Got error getting next reader %s, %s", err, w)
				break
			}

			buf, err := ioutil.ReadAll(r)
			if err != nil {
				logger.Debugf("Got error writing to writer %s", err)
				break
			}

			if w == nil {
				continue
			}

			i, err := w.Write(buf)
			if i != len(buf) {
				logger.Debugf("Didn't write all of buf")
				break
			}
			if err != nil {
				logger.Debugf("Error writing buf %s", err)
				break
			}
		}
		ch <- true
	}(w, conn)

	return ch
}

func WebsocketProxy(source *websocket.Conn, target *websocket.Conn) chan bool {
	forward := func(in *websocket.Conn, out *websocket.Conn, ch chan bool) {
		for {
			mt, r, err := in.NextReader()
			if err != nil {
				break
			}

			w, err := out.NextWriter(mt)
			if err != nil {
				break
			}

			_, err = io.Copy(w, r)
			w.Close()
			if err != nil {
				break
			}
		}

		ch <- true
	}

	chSend := make(chan bool)
	go forward(source, target, chSend)

	chRecv := make(chan bool)
	go forward(target, source, chRecv)

	ch := make(chan bool)
	go func() {
		select {
		case <-chSend:
		case <-chRecv:
		}

		source.Close()
		target.Close()

		ch <- true
	}()

	return ch
}

func defaultReader(conn *websocket.Conn, r io.ReadCloser, readDone chan<- bool) {
	/* For now, we don't need to adjust buffer sizes in
	* WebsocketMirror, since it's used for interactive things like
	* exec.
	 */
	in := ReaderToChannel(r, -1)
	for {
		buf, ok := <-in
		if !ok {
			r.Close()
			logger.Debugf("sending write barrier")
			conn.WriteMessage(websocket.TextMessage, []byte{})
			readDone <- true
			return
		}
		w, err := conn.NextWriter(websocket.BinaryMessage)
		if err != nil {
			logger.Debugf("Got error getting next writer %s", err)
			break
		}

		_, err = w.Write(buf)
		w.Close()
		if err != nil {
			logger.Debugf("Got err writing %s", err)
			break
		}
	}
	closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
	conn.WriteMessage(websocket.CloseMessage, closeMsg)
	readDone <- true
	r.Close()
}

func defaultWriter(conn *websocket.Conn, w io.WriteCloser, writeDone chan<- bool) {
	for {
		mt, r, err := conn.NextReader()
		if err != nil {
			logger.Debugf("Got error getting next reader %s", err)
			break
		}

		if mt == websocket.CloseMessage {
			logger.Debugf("Got close message for reader")
			break
		}

		if mt == websocket.TextMessage {
			logger.Debugf("Got message barrier, resetting stream")
			break
		}

		buf, err := ioutil.ReadAll(r)
		if err != nil {
			logger.Debugf("Got error writing to writer %s", err)
			break
		}
		i, err := w.Write(buf)
		if i != len(buf) {
			logger.Debugf("Didn't write all of buf")
			break
		}
		if err != nil {
			logger.Debugf("Error writing buf %s", err)
			break
		}
	}
	writeDone <- true
	w.Close()
}

// WebsocketMirror allows mirroring a reader to a websocket and taking the
// result and writing it to a writer. This function allows for multiple
// mirrorings and correctly negotiates stream endings. However, it means any
// websocket.Conns passed to it are live when it returns, and must be closed
// explicitly.
type WebSocketMirrorReader func(conn *websocket.Conn, r io.ReadCloser, readDone chan<- bool)
type WebSocketMirrorWriter func(conn *websocket.Conn, w io.WriteCloser, writeDone chan<- bool)

func WebsocketMirror(conn *websocket.Conn, w io.WriteCloser, r io.ReadCloser, Reader WebSocketMirrorReader, Writer WebSocketMirrorWriter) (chan bool, chan bool) {
	readDone := make(chan bool, 1)
	writeDone := make(chan bool, 1)

	ReadFunc := Reader
	if ReadFunc == nil {
		ReadFunc = defaultReader
	}

	WriteFunc := Writer
	if WriteFunc == nil {
		WriteFunc = defaultWriter
	}

	go ReadFunc(conn, r, readDone)
	go WriteFunc(conn, w, writeDone)

	return readDone, writeDone
}

func WebsocketConsoleMirror(conn *websocket.Conn, w io.WriteCloser, r io.ReadCloser) (chan bool, chan bool) {
	readDone := make(chan bool, 1)
	writeDone := make(chan bool, 1)

	go defaultWriter(conn, w, writeDone)

	go func(conn *websocket.Conn, r io.ReadCloser) {
		in := ReaderToChannel(r, -1)
		for {
			buf, ok := <-in
			if !ok {
				r.Close()
				logger.Debugf("sending write barrier")
				conn.WriteMessage(websocket.TextMessage, []byte{})
				readDone <- true
				return
			}
			w, err := conn.NextWriter(websocket.BinaryMessage)
			if err != nil {
				logger.Debugf("Got error getting next writer %s", err)
				break
			}

			_, err = w.Write(buf)
			w.Close()
			if err != nil {
				logger.Debugf("Got err writing %s", err)
				break
			}
		}
		closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
		conn.WriteMessage(websocket.CloseMessage, closeMsg)
		readDone <- true
		r.Close()
	}(conn, r)

	return readDone, writeDone
}

var WebsocketUpgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

// AllocatePort asks the kernel for a free open port that is ready to use
func AllocatePort() (int, error) {
	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
	if err != nil {
		return -1, err
	}

	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return -1, err
	}
	defer l.Close()
	return l.Addr().(*net.TCPAddr).Port, nil
}

func NetworkGetCounters(ifName string) api.NetworkStateCounters {
	counters := api.NetworkStateCounters{}
	// Get counters
	content, err := ioutil.ReadFile("/proc/net/dev")
	if err == nil {
		for _, line := range strings.Split(string(content), "\n") {
			fields := strings.Fields(line)

			if len(fields) != 17 {
				continue
			}

			intName := strings.TrimSuffix(fields[0], ":")
			if intName != ifName {
				continue
			}

			rxBytes, err := strconv.ParseInt(fields[1], 10, 64)
			if err != nil {
				continue
			}

			rxPackets, err := strconv.ParseInt(fields[2], 10, 64)
			if err != nil {
				continue
			}

			txBytes, err := strconv.ParseInt(fields[9], 10, 64)
			if err != nil {
				continue
			}

			txPackets, err := strconv.ParseInt(fields[10], 10, 64)
			if err != nil {
				continue
			}

			counters.BytesSent = txBytes
			counters.BytesReceived = rxBytes
			counters.PacketsSent = txPackets
			counters.PacketsReceived = rxPackets
			break
		}
	}

	return counters
}
<|endoftext|>


<commit_before>package lib

import (
	"encoding/json"
	"fmt"
	"io"

	"github.com/streadway/amqp"
	"golang.org/x/net/context"
)

// A JobQueue allows getting Jobs out of an AMQP queue.
type JobQueue struct {
	conn  *amqp.Connection
	queue string
}

type amqpJob struct {
	conn     *amqp.Connection
	delivery amqp.Delivery
	payload  JobPayload
}

func (j amqpJob) Payload() JobPayload {
	return j.payload
}

func (j amqpJob) Error(err error) error {
	fmt.Printf("amqpJob.Error(%v)\n", err)

	return nil
}

func (j amqpJob) Requeue() error {
	fmt.Printf("amqpJob.Requeue()\n")

	return nil
}

func (j amqpJob) Finish(state FinishState) error {
	fmt.Printf("amqpJob.Finish(%v)\n", state)

	return nil
}

func (j amqpJob) LogWriter(ctx context.Context) (io.WriteCloser, error) {
	return NewLogWriter(ctx, j.conn, j.payload.Job.ID)
}

// NewJobQueue creates a JobQueue backed by the given AMQP connections and
// connects to the AMQP queue with the given name. The queue will be declared
// in AMQP when this function is called, so an error could be raised if the
// queue already exists, but with different attributes than we expect.
func NewJobQueue(conn *amqp.Connection, queue string) (*JobQueue, error) {
	channel, err := conn.Channel()
	if err != nil {
		return nil, err
	}

	_, err = channel.QueueDeclare(queue, true, false, false, false, nil)
	if err != nil {
		return nil, err
	}

	err = channel.Close()
	if err != nil {
		return nil, err
	}

	return &JobQueue{
		conn:  conn,
		queue: queue,
	}, nil
}

// Jobs creates a new consumer on the queue, and returns three channels. The
// first channel gets sent every BuildJob that we receive from AMQP. The
// stopChan is a channel that can be closed in order to stop the consumer.
func (q *JobQueue) Jobs() (outChan <-chan Job, err error) {
	channel, err := q.conn.Channel()
	if err != nil {
		return
	}

	err = channel.Qos(1, 0, false)
	if err != nil {
		return
	}

	deliveries, err := channel.Consume(q.queue, "build-job-consumer", false, false, false, false, nil)
	if err != nil {
		return
	}

	buildJobChan := make(chan Job)
	outChan = buildJobChan

	go func() {
		for delivery := range deliveries {
			var buildJob amqpJob
			err = json.Unmarshal(delivery.Body, &buildJob.payload)
			if err != nil {
				fmt.Printf("JSON parse error: %v\n", err)
			}
			buildJob.conn = q.conn
			buildJob.delivery = delivery

			buildJobChan <- buildJob
		}

		channel.Close()
	}()

	return
}
<commit_msg>Fix a data race<commit_after>package lib

import (
	"encoding/json"
	"fmt"
	"io"

	"github.com/streadway/amqp"
	"golang.org/x/net/context"
)

// A JobQueue allows getting Jobs out of an AMQP queue.
type JobQueue struct {
	conn  *amqp.Connection
	queue string
}

type amqpJob struct {
	conn     *amqp.Connection
	delivery amqp.Delivery
	payload  JobPayload
}

func (j amqpJob) Payload() JobPayload {
	return j.payload
}

func (j amqpJob) Error(err error) error {
	fmt.Printf("amqpJob.Error(%v)\n", err)

	return nil
}

func (j amqpJob) Requeue() error {
	fmt.Printf("amqpJob.Requeue()\n")

	return nil
}

func (j amqpJob) Finish(state FinishState) error {
	fmt.Printf("amqpJob.Finish(%v)\n", state)

	return nil
}

func (j amqpJob) LogWriter(ctx context.Context) (io.WriteCloser, error) {
	return NewLogWriter(ctx, j.conn, j.payload.Job.ID)
}

// NewJobQueue creates a JobQueue backed by the given AMQP connections and
// connects to the AMQP queue with the given name. The queue will be declared
// in AMQP when this function is called, so an error could be raised if the
// queue already exists, but with different attributes than we expect.
func NewJobQueue(conn *amqp.Connection, queue string) (*JobQueue, error) {
	channel, err := conn.Channel()
	if err != nil {
		return nil, err
	}

	_, err = channel.QueueDeclare(queue, true, false, false, false, nil)
	if err != nil {
		return nil, err
	}

	err = channel.Close()
	if err != nil {
		return nil, err
	}

	return &JobQueue{
		conn:  conn,
		queue: queue,
	}, nil
}

// Jobs creates a new consumer on the queue, and returns three channels. The
// first channel gets sent every BuildJob that we receive from AMQP. The
// stopChan is a channel that can be closed in order to stop the consumer.
func (q *JobQueue) Jobs() (outChan <-chan Job, err error) {
	channel, err := q.conn.Channel()
	if err != nil {
		return
	}

	err = channel.Qos(1, 0, false)
	if err != nil {
		return
	}

	deliveries, err := channel.Consume(q.queue, "build-job-consumer", false, false, false, false, nil)
	if err != nil {
		return
	}

	buildJobChan := make(chan Job)
	outChan = buildJobChan

	go func() {
		for delivery := range deliveries {
			var buildJob amqpJob
			err := json.Unmarshal(delivery.Body, &buildJob.payload)
			if err != nil {
				fmt.Printf("JSON parse error: %v\n", err)
			}
			buildJob.conn = q.conn
			buildJob.delivery = delivery

			buildJobChan <- buildJob
		}

		channel.Close()
	}()

	return
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"encoding/base64"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"regexp"

	"github.com/Syfaro/telegram-bot-api"
)

// version
const VERSION = "1.0"

// Command - one command type
type Commands map[string]string

// Config - config struct
type Config struct {
	token   string // bot token
	addExit bool   // add /exit command
}

// Users in chat
type User struct {
	UserName     string
	AuthCode     string
	IsAuthorized bool
}

type Users map[int]*User

// ------------------------------------------------------------------
// get config
func getConfig() (commands Commands, app_config Config, err error) {
	flag.StringVar(&app_config.token, "tb-token", "", "set bot token (or set TB_TOKEN variable)")
	flag.BoolVar(&app_config.addExit, "add-exit", false, "add /exit command for terminate bot")

	flag.Usage = func() {
		fmt.Printf("usage: %s [options] /chat_command \"shell command\" /chat_command2 \"shell command2\"\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(0)
	}
	version := flag.Bool("version", false, "get version")
	flag.Parse()
	if *version {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	commands = Commands{}
	// need >= 2 arguments and count of it must be even
	args := flag.Args()
	if len(args) < 2 || len(args)%2 == 1 {
		return commands, app_config, fmt.Errorf("error: need pairs of chat-command and shell-command")
	}

	for i := 0; i < len(args); i += 2 {
		path, cmd := args[i], args[i+1]
		if path[0] != '/' {
			return commands, app_config, fmt.Errorf("error: path %s dont starts with /", path)
		}
		commands[path] = cmd
	}

	if app_config.token == "" {
		if app_config.token = os.Getenv("TB_TOKEN"); app_config.token == "" {
			return commands, app_config, fmt.Errorf("TB_TOKEN env var not found. See https://core.telegram.org/bots#botfather for more information\n")
		}
	}

	return commands, app_config, nil
}

// ------------------------------------------------------------------
func getRandomCode() string {
	buffer := make([]byte, 15)
	_, err := rand.Read(buffer)
	if err != nil {
		log.Fatal("Get code error:", err)
	}

	return base64.URLEncoding.EncodeToString(buffer)
}

// ------------------------------------------------------------------
func main() {
	commands, app_config, err := getConfig()
	if err != nil {
		log.Fatal(err)
	}

	bot, err := tgbotapi.NewBotAPI(app_config.token)
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("Authorized on account %s", bot.Self.UserName)

	var ucfg tgbotapi.UpdateConfig = tgbotapi.NewUpdate(0)
	ucfg.Timeout = 60
	err = bot.UpdatesChan(ucfg)

	go_exit := false
	users := make(Users)

LOOP:
	for {
		select {
		case telegram_update := <-bot.Updates:

			chat_id := telegram_update.Message.Chat.ID

			parts := regexp.MustCompile(`\s+`).Split(telegram_update.Message.Text, 2)
			replay_msg := ""

			if len(parts) > 0 && parts[0][0] == '/' {

				user_from := telegram_update.Message.From
				allowExec := false
				if _, ok := users[user_from.ID]; ok {
					allowExec = users[user_from.ID].IsAuthorized
				}

				if parts[0] == "/auth" {

					if _, ok := users[user_from.ID]; !ok {
						users[user_from.ID] = &User{
							UserName:     user_from.UserName,
							AuthCode:     "",
							IsAuthorized: false,
						}
					}

					if len(parts) == 1 || parts[1] == "" {
						replay_msg = "See code in terminal with shell2telegram and type:\n/auth code"
						users[user_from.ID].IsAuthorized = false
						users[user_from.ID].AuthCode = getRandomCode()
						fmt.Printf("Code (for %s): %s\n", user_from.UserName, users[user_from.ID].AuthCode)
					} else if len(parts) > 1 && parts[1] != "" && parts[1] == users[user_from.ID].AuthCode {
						users[user_from.ID].IsAuthorized = true
						replay_msg = fmt.Sprintf("You (%s) authorized.", user_from.UserName)
					} else if len(parts) > 1 && parts[1] != "" && parts[1] != users[user_from.ID].AuthCode {
						replay_msg = fmt.Sprintf("Code is not valid.")
					}

				} else if parts[0] == "/help" {

					if allowExec {
						for cmd, shell_cmd := range commands {
							replay_msg += fmt.Sprintf("%s - %s\n", cmd, shell_cmd)
						}
						if app_config.addExit {
							replay_msg += fmt.Sprintf("%s - %s\n", "/exit", "terminate bot")
						}
					}
					replay_msg += fmt.Sprintf("%s - %s\n", "/auth [code]", "authorize user")

				} else if allowExec && app_config.addExit && parts[0] == "/exit" {
					replay_msg = "bye..."
					go_exit = true
				} else if cmd, found := commands[parts[0]]; allowExec && found {

					shell, params := "sh", []string{"-c", cmd}
					if len(parts) > 1 {
						params = append(params, parts[1])
					}

					os_exec_command := exec.Command(shell, params...)
					os_exec_command.Stderr = os.Stderr
					shell_out, err := os_exec_command.Output()
					if err != nil {
						log.Println("exec error: ", err)
						replay_msg = fmt.Sprintf("exec error: %s", err)
					} else {
						replay_msg = string(shell_out)
					}
				}

				if replay_msg != "" {
					bot.SendMessage(tgbotapi.NewMessage(chat_id, replay_msg))
					if go_exit {
						break LOOP
					}
				}
			}
		}
	}
}
<commit_msg>Fixed failure on image in message<commit_after>package main

import (
	"crypto/rand"
	"encoding/base64"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"regexp"

	"github.com/Syfaro/telegram-bot-api"
)

// version
const VERSION = "1.0"

// Command - one command type
type Commands map[string]string

// Config - config struct
type Config struct {
	token   string // bot token
	addExit bool   // add /exit command
}

// Users in chat
type User struct {
	UserName     string
	AuthCode     string
	IsAuthorized bool
}

type Users map[int]*User

// ------------------------------------------------------------------
// get config
func getConfig() (commands Commands, app_config Config, err error) {
	flag.StringVar(&app_config.token, "tb-token", "", "set bot token (or set TB_TOKEN variable)")
	flag.BoolVar(&app_config.addExit, "add-exit", false, "add /exit command for terminate bot")

	flag.Usage = func() {
		fmt.Printf("usage: %s [options] /chat_command \"shell command\" /chat_command2 \"shell command2\"\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(0)
	}
	version := flag.Bool("version", false, "get version")
	flag.Parse()
	if *version {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	commands = Commands{}
	// need >= 2 arguments and count of it must be even
	args := flag.Args()
	if len(args) < 2 || len(args)%2 == 1 {
		return commands, app_config, fmt.Errorf("error: need pairs of chat-command and shell-command")
	}

	for i := 0; i < len(args); i += 2 {
		path, cmd := args[i], args[i+1]
		if path[0] != '/' {
			return commands, app_config, fmt.Errorf("error: path %s dont starts with /", path)
		}
		commands[path] = cmd
	}

	if app_config.token == "" {
		if app_config.token = os.Getenv("TB_TOKEN"); app_config.token == "" {
			return commands, app_config, fmt.Errorf("TB_TOKEN env var not found. See https://core.telegram.org/bots#botfather for more information\n")
		}
	}

	return commands, app_config, nil
}

// ------------------------------------------------------------------
func getRandomCode() string {
	buffer := make([]byte, 15)
	_, err := rand.Read(buffer)
	if err != nil {
		log.Fatal("Get code error:", err)
	}

	return base64.URLEncoding.EncodeToString(buffer)
}

// ------------------------------------------------------------------
func main() {
	commands, app_config, err := getConfig()
	if err != nil {
		log.Fatal(err)
	}

	bot, err := tgbotapi.NewBotAPI(app_config.token)
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("Authorized on account %s", bot.Self.UserName)

	var ucfg tgbotapi.UpdateConfig = tgbotapi.NewUpdate(0)
	ucfg.Timeout = 60
	err = bot.UpdatesChan(ucfg)

	go_exit := false
	users := make(Users)

LOOP:
	for {
		select {
		case telegram_update := <-bot.Updates:

			chat_id := telegram_update.Message.Chat.ID

			parts := regexp.MustCompile(`\s+`).Split(telegram_update.Message.Text, 2)
			replay_msg := ""

			if len(parts) > 0 && len(parts[0]) > 0 && parts[0][0] == '/' {

				user_from := telegram_update.Message.From
				allowExec := false
				if _, ok := users[user_from.ID]; ok {
					allowExec = users[user_from.ID].IsAuthorized
				}

				if parts[0] == "/auth" {

					if _, ok := users[user_from.ID]; !ok {
						users[user_from.ID] = &User{
							UserName:     user_from.UserName,
							AuthCode:     "",
							IsAuthorized: false,
						}
					}

					if len(parts) == 1 || parts[1] == "" {
						replay_msg = "See code in terminal with shell2telegram and type:\n/auth code"
						users[user_from.ID].IsAuthorized = false
						users[user_from.ID].AuthCode = getRandomCode()
						fmt.Printf("Code (for %s): %s\n", user_from.UserName, users[user_from.ID].AuthCode)
					} else if len(parts) > 1 && parts[1] != "" && parts[1] == users[user_from.ID].AuthCode {
						users[user_from.ID].IsAuthorized = true
						replay_msg = fmt.Sprintf("You (%s) authorized.", user_from.UserName)
					} else if len(parts) > 1 && parts[1] != "" && parts[1] != users[user_from.ID].AuthCode {
						replay_msg = fmt.Sprintf("Code is not valid.")
					}

				} else if parts[0] == "/help" {

					if allowExec {
						for cmd, shell_cmd := range commands {
							replay_msg += fmt.Sprintf("%s - %s\n", cmd, shell_cmd)
						}
						if app_config.addExit {
							replay_msg += fmt.Sprintf("%s - %s\n", "/exit", "terminate bot")
						}
					}
					replay_msg += fmt.Sprintf("%s - %s\n", "/auth [code]", "authorize user")

				} else if allowExec && app_config.addExit && parts[0] == "/exit" {
					replay_msg = "bye..."
					go_exit = true
				} else if cmd, found := commands[parts[0]]; allowExec && found {

					shell, params := "sh", []string{"-c", cmd}
					if len(parts) > 1 {
						params = append(params, parts[1])
					}

					os_exec_command := exec.Command(shell, params...)
					os_exec_command.Stderr = os.Stderr
					shell_out, err := os_exec_command.Output()
					if err != nil {
						log.Println("exec error: ", err)
						replay_msg = fmt.Sprintf("exec error: %s", err)
					} else {
						replay_msg = string(shell_out)
					}
				}

				if replay_msg != "" {
					bot.SendMessage(tgbotapi.NewMessage(chat_id, replay_msg))
					if go_exit {
						break LOOP
					}
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"regexp"

	"github.com/Syfaro/telegram-bot-api"
)

// version
const VERSION = "1.0"

func main() {
	flag.Usage = func() {
		fmt.Printf("usage: %s [options] /chat_command \"shell command\" /chat_command2 \"shell command2\"\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(0)
	}
	version := flag.Bool("version", false, "get version")
	flag.Parse()
	if *version {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	// need >= 2 arguments and count of it must be even
	args := flag.Args()
	if len(args) < 2 || len(args)%2 == 1 {
		log.Fatal("error: need pairs of path and shell command")
	}

	cmd_handlers := map[string]string{}

	for i := 0; i < len(args); i += 2 {
		path, cmd := args[i], args[i+1]
		if path[0] != '/' {
			log.Fatalf("error: path %s dont starts with /", path)
		}
		cmd_handlers[path] = cmd
	}

	tb_token := os.Getenv("TB_TOKEN")
	if tb_token == "" {
		log.Fatal("TB_TOKEN env var not found. See https://core.telegram.org/bots#botfather for more information\n")
	}

	bot, err := tgbotapi.NewBotAPI(tb_token)
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("Authorized on account %s", bot.Self.UserName)

	var ucfg tgbotapi.UpdateConfig = tgbotapi.NewUpdate(0)
	ucfg.Timeout = 60
	err = bot.UpdatesChan(ucfg)

	for {
		select {
		case update := <-bot.Updates:

			chat_id := update.Message.Chat.ID

			parts := regexp.MustCompile(`\s+`).Split(update.Message.Text, 2)
			if len(parts) > 0 {
				if parts[0] == "/help" {
					reply := ""
					for cmd, shell_cmd := range cmd_handlers {
						reply += fmt.Sprintf("%s\n %s\n", cmd, shell_cmd)
					}
					msg := tgbotapi.NewMessage(chat_id, reply)
					bot.SendMessage(msg)
				}

				if cmd, found := cmd_handlers[parts[0]]; found {
					shell, params := "sh", []string{"-c", cmd}
					if len(parts) > 1 {
						params = append(params, parts[1])
					}

					os_exec_command := exec.Command(shell, params...)
					os_exec_command.Stderr = os.Stderr
					shell_out, err := os_exec_command.Output()
					if err != nil {
						log.Println("exec error: ", err)
						reply := fmt.Sprintf("exec error: %s", err)
						msg := tgbotapi.NewMessage(chat_id, reply)
						bot.SendMessage(msg)
					} else {
						reply := string(shell_out)
						msg := tgbotapi.NewMessage(chat_id, reply)
						bot.SendMessage(msg)
					}
				}
			}
		}
	}
}
<commit_msg>Refactoring<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"regexp"

	"github.com/Syfaro/telegram-bot-api"
)

// version
const VERSION = "1.0"

func main() {
	flag.Usage = func() {
		fmt.Printf("usage: %s [options] /chat_command \"shell command\" /chat_command2 \"shell command2\"\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(0)
	}
	version := flag.Bool("version", false, "get version")
	flag.Parse()
	if *version {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	// need >= 2 arguments and count of it must be even
	args := flag.Args()
	if len(args) < 2 || len(args)%2 == 1 {
		log.Fatal("error: need pairs of path and shell command")
	}

	cmd_handlers := map[string]string{}

	for i := 0; i < len(args); i += 2 {
		path, cmd := args[i], args[i+1]
		if path[0] != '/' {
			log.Fatalf("error: path %s dont starts with /", path)
		}
		cmd_handlers[path] = cmd
	}

	tb_token := os.Getenv("TB_TOKEN")
	if tb_token == "" {
		log.Fatal("TB_TOKEN env var not found. See https://core.telegram.org/bots#botfather for more information\n")
	}

	bot, err := tgbotapi.NewBotAPI(tb_token)
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("Authorized on account %s", bot.Self.UserName)

	var ucfg tgbotapi.UpdateConfig = tgbotapi.NewUpdate(0)
	ucfg.Timeout = 60
	err = bot.UpdatesChan(ucfg)

	for {
		select {
		case update := <-bot.Updates:

			chat_id := update.Message.Chat.ID

			parts := regexp.MustCompile(`\s+`).Split(update.Message.Text, 2)
			replay_msg := ""

			if len(parts) > 0 {
				if parts[0] == "/help" {
					for cmd, shell_cmd := range cmd_handlers {
						replay_msg += fmt.Sprintf("%s\n %s\n", cmd, shell_cmd)
					}

				} else if cmd, found := cmd_handlers[parts[0]]; found {

					shell, params := "sh", []string{"-c", cmd}
					if len(parts) > 1 {
						params = append(params, parts[1])
					}

					os_exec_command := exec.Command(shell, params...)
					os_exec_command.Stderr = os.Stderr
					shell_out, err := os_exec_command.Output()
					if err != nil {
						log.Println("exec error: ", err)
						replay_msg = fmt.Sprintf("exec error: %s", err)
					} else {
						replay_msg = string(shell_out)
					}
				}
			}

			if replay_msg != "" {
				bot.SendMessage(tgbotapi.NewMessage(chat_id, replay_msg))
			}
		}
	}
}
<|endoftext|>


<commit_before>package lnwallet

import (
	"bytes"
	"reflect"
	"testing"

	"github.com/roasbeef/btcd/btcec"
	"github.com/roasbeef/btcd/txscript"
	"github.com/roasbeef/btcd/wire"
)

func TestSignDescriptorSerialization(t *testing.T) {
	keys := [][]byte{
		{0x04, 0x11, 0xdb, 0x93, 0xe1, 0xdc, 0xdb, 0x8a,
			0x01, 0x6b, 0x49, 0x84, 0x0f, 0x8c, 0x53, 0xbc, 0x1e,
			0xb6, 0x8a, 0x38, 0x2e, 0x97, 0xb1, 0x48, 0x2e, 0xca,
			0xd7, 0xb1, 0x48, 0xa6, 0x90, 0x9a, 0x5c, 0xb2, 0xe0,
			0xea, 0xdd, 0xfb, 0x84, 0xcc, 0xf9, 0x74, 0x44, 0x64,
			0xf8, 0x2e, 0x16, 0x0b, 0xfa, 0x9b, 0x8b, 0x64, 0xf9,
			0xd4, 0xc0, 0x3f, 0x99, 0x9b, 0x86, 0x43, 0xf6, 0x56,
			0xb4, 0x12, 0xa3,
		},
		{0x04, 0x11, 0xdb, 0x93, 0xe1, 0xdc, 0xdb, 0x8a,
			0x01, 0x6b, 0x49, 0x84, 0x0f, 0x8c, 0x53, 0xbc, 0x1e,
			0xb6, 0x8a, 0x38, 0x2e, 0x97, 0xb1, 0x48, 0x2e, 0xca,
			0xd7, 0xb1, 0x48, 0xa6, 0x90, 0x9a, 0x5c, 0xb2, 0xe0,
			0xea, 0xdd, 0xfb, 0x84, 0xcc, 0xf9, 0x74, 0x44, 0x64,
			0xf8, 0x2e, 0x16, 0x0b, 0xfa, 0x9b, 0x8b, 0x64, 0xf9,
			0xd4, 0xc0, 0x3f, 0x99, 0x9b, 0x86, 0x43, 0xf6, 0x56,
			0xb4, 0x12, 0xa3,
		},
	}

	signDescriptors := []SignDescriptor{
		{
			SingleTweak: []byte{
				0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02,
				0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02,
				0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02,
				0x02, 0x02, 0x02, 0x02, 0x02,
			},
			WitnessScript: []byte{
				0x00, 0x14, 0xee, 0x91, 0x41, 0x7e, 0x85, 0x6c, 0xde,
				0x10, 0xa2, 0x91, 0x1e, 0xdc, 0xbd, 0xbd, 0x69, 0xe2,
				0xef, 0xb5, 0x71, 0x48,
			},
			Output: &wire.TxOut{
				Value: 5000000000,
				PkScript: []byte{
					0x41, // OP_DATA_65
					0x04, 0xd6, 0x4b, 0xdf, 0xd0, 0x9e, 0xb1, 0xc5,
					0xfe, 0x29, 0x5a, 0xbd, 0xeb, 0x1d, 0xca, 0x42,
					0x81, 0xbe, 0x98, 0x8e, 0x2d, 0xa0, 0xb6, 0xc1,
					0xc6, 0xa5, 0x9d, 0xc2, 0x26, 0xc2, 0x86, 0x24,
					0xe1, 0x81, 0x75, 0xe8, 0x51, 0xc9, 0x6b, 0x97,
					0x3d, 0x81, 0xb0, 0x1c, 0xc3, 0x1f, 0x04, 0x78,
					0x34, 0xbc, 0x06, 0xd6, 0xd6, 0xed, 0xf6, 0x20,
					0xd1, 0x84, 0x24, 0x1a, 0x6a, 0xed, 0x8b, 0x63,
					0xa6, // 65-byte signature
					0xac, // OP_CHECKSIG
				},
			},
			HashType: txscript.SigHashAll,
		},

		// Test serializing a SignDescriptor with a nil-valued PrivateTweak
		{
			SingleTweak: nil,
			WitnessScript: []byte{
				0x00, 0x14, 0xee, 0x91, 0x41, 0x7e, 0x85, 0x6c, 0xde,
				0x10, 0xa2, 0x91, 0x1e, 0xdc, 0xbd, 0xbd, 0x69, 0xe2,
				0xef, 0xb5, 0x71, 0x48,
			},
			Output: &wire.TxOut{
				Value: 5000000000,
				PkScript: []byte{
					0x41, // OP_DATA_65
					0x04, 0xd6, 0x4b, 0xdf, 0xd0, 0x9e, 0xb1, 0xc5,
					0xfe, 0x29, 0x5a, 0xbd, 0xeb, 0x1d, 0xca, 0x42,
					0x81, 0xbe, 0x98, 0x8e, 0x2d, 0xa0, 0xb6, 0xc1,
					0xc6, 0xa5, 0x9d, 0xc2, 0x26, 0xc2, 0x86, 0x24,
					0xe1, 0x81, 0x75, 0xe8, 0x51, 0xc9, 0x6b, 0x97,
					0x3d, 0x81, 0xb0, 0x1c, 0xc3, 0x1f, 0x04, 0x78,
					0x34, 0xbc, 0x06, 0xd6, 0xd6, 0xed, 0xf6, 0x20,
					0xd1, 0x84, 0x24, 0x1a, 0x6a, 0xed, 0x8b, 0x63,
					0xa6, // 65-byte signature
					0xac, // OP_CHECKSIG
				},
			},
			HashType: txscript.SigHashAll,
		},
	}

	for i := 0; i < len(signDescriptors); i++ {
		// Parse pubkeys for each sign descriptor.
		sd := &signDescriptors[i]
		pubkey, err := btcec.ParsePubKey(keys[i], btcec.S256())
		if err != nil {
			t.Fatalf("unable to parse pubkey: %v", err)
		}
		sd.PubKey = pubkey

		// Test that serialize -> deserialize yields same result as original.
		var buf bytes.Buffer
		if err := WriteSignDescriptor(&buf, sd); err != nil {
			t.Fatalf("unable to serialize sign descriptor[%v]: %v", i, sd)
		}

		desSd := &SignDescriptor{}
		if err := ReadSignDescriptor(&buf, desSd); err != nil {
			t.Fatalf("unable to deserialize sign descriptor[%v]: %v", i, sd)
		}

		if !reflect.DeepEqual(sd, desSd) {
			t.Fatalf("original and deserialized sign descriptors not equal:\n"+
				"original     : %+v\n"+
				"deserialized : %+v\n",
				sd, desSd)
		}
	}
}
<commit_msg>lnwallet: update signdescriptor_test.go due to latest API changes<commit_after>package lnwallet

import (
	"bytes"
	"reflect"
	"testing"

	"github.com/lightningnetwork/lnd/keychain"
	"github.com/roasbeef/btcd/btcec"
	"github.com/roasbeef/btcd/txscript"
	"github.com/roasbeef/btcd/wire"
)

func TestSignDescriptorSerialization(t *testing.T) {
	keys := [][]byte{
		{0x04, 0x11, 0xdb, 0x93, 0xe1, 0xdc, 0xdb, 0x8a,
			0x01, 0x6b, 0x49, 0x84, 0x0f, 0x8c, 0x53, 0xbc, 0x1e,
			0xb6, 0x8a, 0x38, 0x2e, 0x97, 0xb1, 0x48, 0x2e, 0xca,
			0xd7, 0xb1, 0x48, 0xa6, 0x90, 0x9a, 0x5c, 0xb2, 0xe0,
			0xea, 0xdd, 0xfb, 0x84, 0xcc, 0xf9, 0x74, 0x44, 0x64,
			0xf8, 0x2e, 0x16, 0x0b, 0xfa, 0x9b, 0x8b, 0x64, 0xf9,
			0xd4, 0xc0, 0x3f, 0x99, 0x9b, 0x86, 0x43, 0xf6, 0x56,
			0xb4, 0x12, 0xa3,
		},
		{0x04, 0x11, 0xdb, 0x93, 0xe1, 0xdc, 0xdb, 0x8a,
			0x01, 0x6b, 0x49, 0x84, 0x0f, 0x8c, 0x53, 0xbc, 0x1e,
			0xb6, 0x8a, 0x38, 0x2e, 0x97, 0xb1, 0x48, 0x2e, 0xca,
			0xd7, 0xb1, 0x48, 0xa6, 0x90, 0x9a, 0x5c, 0xb2, 0xe0,
			0xea, 0xdd, 0xfb, 0x84, 0xcc, 0xf9, 0x74, 0x44, 0x64,
			0xf8, 0x2e, 0x16, 0x0b, 0xfa, 0x9b, 0x8b, 0x64, 0xf9,
			0xd4, 0xc0, 0x3f, 0x99, 0x9b, 0x86, 0x43, 0xf6, 0x56,
			0xb4, 0x12, 0xa3,
		},
	}

	signDescriptors := []SignDescriptor{
		{
			SingleTweak: []byte{
				0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02,
				0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02,
				0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02,
				0x02, 0x02, 0x02, 0x02, 0x02,
			},
			WitnessScript: []byte{
				0x00, 0x14, 0xee, 0x91, 0x41, 0x7e, 0x85, 0x6c, 0xde,
				0x10, 0xa2, 0x91, 0x1e, 0xdc, 0xbd, 0xbd, 0x69, 0xe2,
				0xef, 0xb5, 0x71, 0x48,
			},
			Output: &wire.TxOut{
				Value: 5000000000,
				PkScript: []byte{
					0x41, // OP_DATA_65
					0x04, 0xd6, 0x4b, 0xdf, 0xd0, 0x9e, 0xb1, 0xc5,
					0xfe, 0x29, 0x5a, 0xbd, 0xeb, 0x1d, 0xca, 0x42,
					0x81, 0xbe, 0x98, 0x8e, 0x2d, 0xa0, 0xb6, 0xc1,
					0xc6, 0xa5, 0x9d, 0xc2, 0x26, 0xc2, 0x86, 0x24,
					0xe1, 0x81, 0x75, 0xe8, 0x51, 0xc9, 0x6b, 0x97,
					0x3d, 0x81, 0xb0, 0x1c, 0xc3, 0x1f, 0x04, 0x78,
					0x34, 0xbc, 0x06, 0xd6, 0xd6, 0xed, 0xf6, 0x20,
					0xd1, 0x84, 0x24, 0x1a, 0x6a, 0xed, 0x8b, 0x63,
					0xa6, // 65-byte signature
					0xac, // OP_CHECKSIG
				},
			},
			HashType: txscript.SigHashAll,
		},

		// Test serializing a SignDescriptor with a nil-valued PrivateTweak
		{
			SingleTweak: nil,
			WitnessScript: []byte{
				0x00, 0x14, 0xee, 0x91, 0x41, 0x7e, 0x85, 0x6c, 0xde,
				0x10, 0xa2, 0x91, 0x1e, 0xdc, 0xbd, 0xbd, 0x69, 0xe2,
				0xef, 0xb5, 0x71, 0x48,
			},
			Output: &wire.TxOut{
				Value: 5000000000,
				PkScript: []byte{
					0x41, // OP_DATA_65
					0x04, 0xd6, 0x4b, 0xdf, 0xd0, 0x9e, 0xb1, 0xc5,
					0xfe, 0x29, 0x5a, 0xbd, 0xeb, 0x1d, 0xca, 0x42,
					0x81, 0xbe, 0x98, 0x8e, 0x2d, 0xa0, 0xb6, 0xc1,
					0xc6, 0xa5, 0x9d, 0xc2, 0x26, 0xc2, 0x86, 0x24,
					0xe1, 0x81, 0x75, 0xe8, 0x51, 0xc9, 0x6b, 0x97,
					0x3d, 0x81, 0xb0, 0x1c, 0xc3, 0x1f, 0x04, 0x78,
					0x34, 0xbc, 0x06, 0xd6, 0xd6, 0xed, 0xf6, 0x20,
					0xd1, 0x84, 0x24, 0x1a, 0x6a, 0xed, 0x8b, 0x63,
					0xa6, // 65-byte signature
					0xac, // OP_CHECKSIG
				},
			},
			HashType: txscript.SigHashAll,
		},
	}

	for i := 0; i < len(signDescriptors); i++ {
		// Parse pubkeys for each sign descriptor.
		sd := &signDescriptors[i]
		pubkey, err := btcec.ParsePubKey(keys[i], btcec.S256())
		if err != nil {
			t.Fatalf("unable to parse pubkey: %v", err)
		}
		sd.KeyDesc = keychain.KeyDescriptor{
			KeyLocator: keychain.KeyLocator{
				Family: 50,
				Index:  99,
			},
			PubKey: pubkey,
		}

		// Test that serialize -> deserialize yields same result as original.
		var buf bytes.Buffer
		if err := WriteSignDescriptor(&buf, sd); err != nil {
			t.Fatalf("unable to serialize sign descriptor[%v]: %v", i, sd)
		}

		desSd := &SignDescriptor{}
		if err := ReadSignDescriptor(&buf, desSd); err != nil {
			t.Fatalf("unable to deserialize sign descriptor[%v]: %v", i, sd)
		}

		if !reflect.DeepEqual(sd, desSd) {
			t.Fatalf("original and deserialized sign descriptors not equal:\n"+
				"original     : %+v\n"+
				"deserialized : %+v\n",
				sd, desSd)
		}
	}
}
<|endoftext|>


<commit_before>package manta

import (
	"testing"

	"github.com/dotabuff/manta/dota"
	"github.com/golang/protobuf/proto"
	"github.com/stretchr/testify/assert"
)

func TestFieldpath(t *testing.T) {
	assert := assert.New(t)

	// roughly the same format used in property_test.go
	scenarios := []struct {
		tableName   string // the name of the table, must have a sendtable fixture.
		run         bool   // whether or not we run the test.
		debug       bool   // whether or not we print debugging output.
		expectCount int    // how many result entries we expect.
	}{
		{
			tableName:   "CRagdollManager",
			run:         true,
			debug:       true,
			expectCount: 1,
		},
		{
			tableName:   "CDOTATeam",
			run:         true,
			debug:       true,
			expectCount: 15,
		},
	}

	// Load our send tables
	m := &dota.CDemoSendTables{}
	if err := proto.Unmarshal(_read_fixture("send_tables/1560315800.pbmsg"), m); err != nil {
		panic(err)
	}

	// Retrieve the flattened field serializer
	fs, err := parseSendTablesNew(m)
	assert.Nil(err)

	// Build the huffman tree
	huf := fieldpath_huffman()

	// Iterate over the different scenarios
	// -! Create a new FieldPath for each scenario
	for _, s := range scenarios {
		// Load up a fixture
		buf := _read_fixture(_sprintf("instancebaseline/1560315800_%s.rawbuf", s.tableName))

		// Get the serializer
		// We don't really know which version is used to generate the baseline
		// 0 seems resonable
		serializer := fs.Serializers[s.tableName][0]
		assert.NotNil(serializer)

		// Optionally skip
		if !s.run {
			continue
		}

		// Set debug status
		debugMode = s.debug

		// Initialize a field path and walk it
		fieldPath := fieldpath_init(serializer, &huf)
		fieldPath.fieldpath_walk(newReader(buf))

		// Verify field count
		assert.Equal(len(fieldPath.fields), s.expectCount)
	}
}
<commit_msg>Add CWorld to fieldpath_test (WIP)<commit_after>package manta

import (
	"testing"

	"github.com/dotabuff/manta/dota"
	"github.com/golang/protobuf/proto"
	"github.com/stretchr/testify/assert"
)

func TestFieldpath(t *testing.T) {
	assert := assert.New(t)

	// roughly the same format used in property_test.go
	scenarios := []struct {
		tableName   string // the name of the table, must have a sendtable fixture.
		run         bool   // whether or not we run the test.
		debug       bool   // whether or not we print debugging output.
		expectCount int    // how many result entries we expect.
	}{
		{
			tableName:   "CRagdollManager",
			run:         true,
			debug:       false,
			expectCount: 1,
		},
		{
			tableName:   "CDOTATeam",
			run:         true,
			debug:       false,
			expectCount: 15,
		},
		{
			tableName:   "CWorld",
			run:         true,
			debug:       true,
			expectCount: 99, // tbd
		},
	}

	// Load our send tables
	m := &dota.CDemoSendTables{}
	if err := proto.Unmarshal(_read_fixture("send_tables/1560315800.pbmsg"), m); err != nil {
		panic(err)
	}

	// Retrieve the flattened field serializer
	fs, err := parseSendTablesNew(m)
	assert.Nil(err)

	// Build the huffman tree
	huf := fieldpath_huffman()

	// Iterate over the different scenarios
	// -! Create a new FieldPath for each scenario
	for _, s := range scenarios {
		// Load up a fixture
		buf := _read_fixture(_sprintf("instancebaseline/1560315800_%s.rawbuf", s.tableName))

		// Get the serializer
		// We don't really know which version is used to generate the baseline
		// 0 seems resonable
		serializer := fs.Serializers[s.tableName][0]
		assert.NotNil(serializer)

		// Optionally skip
		if !s.run {
			continue
		}

		// Set debug status
		debugMode = s.debug

		// Initialize a field path and walk it
		fieldPath := fieldpath_init(serializer, &huf)
		fieldPath.fieldpath_walk(newReader(buf))

		// Verify field count
		assert.Equal(len(fieldPath.fields), s.expectCount)
	}
}
<|endoftext|>


<commit_before>package translators

import (
	"fmt"
	"strings"

	"github.com/pkg/errors"

	"github.com/markbates/pop/fizz"
)

type SQLite struct {
	Schema Schema
}

func NewSQLite(url string) *SQLite {
	return &SQLite{
		Schema: &sqliteSchema{
			URL:    url,
			schema: map[string]*fizz.Table{},
		},
	}
}

func (p *SQLite) CreateTable(t fizz.Table) (string, error) {
	sql := []string{}
	cols := []string{}
	var s string
	for _, c := range t.Columns {
		if c.Primary {
			switch strings.ToLower(c.ColType) {
			case "integer":
				s = fmt.Sprintf("\"%s\" INTEGER PRIMARY KEY AUTOINCREMENT", c.Name)
			case "uuid", "string":
				s = fmt.Sprintf("\"%s\" TEXT PRIMARY KEY", c.Name)
			}
		} else {
			s = p.buildColumn(c)
		}
		cols = append(cols, s)
	}
	s = fmt.Sprintf("CREATE TABLE \"%s\" (\n%s\n);", t.Name, strings.Join(cols, ",\n"))
	sql = append(sql, s)

	for _, i := range t.Indexes {
		s, err := p.AddIndex(fizz.Table{
			Name:    t.Name,
			Indexes: []fizz.Index{i},
		})
		if err != nil {
			return "", err
		}
		sql = append(sql, s)
	}

	return strings.Join(sql, "\n"), nil
}

func (p *SQLite) DropTable(t fizz.Table) (string, error) {
	p.Schema.Delete(t.Name)
	s := fmt.Sprintf("DROP TABLE \"%s\";", t.Name)
	return s, nil
}

func (p *SQLite) RenameTable(t []fizz.Table) (string, error) {
	if len(t) < 2 {
		return "", errors.New("Not enough table names supplied!")
	}
	oldName := t[0].Name
	newName := t[1].Name
	tableInfo, err := p.Schema.TableInfo(oldName)
	if err != nil {
		return "", err
	}
	tableInfo.Name = newName
	s := fmt.Sprintf("ALTER TABLE \"%s\" RENAME TO \"%s\";", oldName, newName)
	return s, nil
}

func (p *SQLite) ChangeColumn(t fizz.Table) (string, error) {
	tableInfo, err := p.Schema.TableInfo(t.Name)

	if err != nil {
		return "", err
	}

	for i := range tableInfo.Columns {
		if tableInfo.Columns[i].Name == t.Columns[0].Name {
			tableInfo.Columns[i] = t.Columns[0]
			break
		}
	}

	sql := []string{}
	s, err := p.withTempTable(t.Name, func(tempTable fizz.Table) (string, error) {
		createTableSQL, err := p.CreateTable(*tableInfo)
		fmt.Printf("table: %v, sql: %v\n", tableInfo, createTableSQL)
		if err != nil {
			return "", err
		}

		ins := fmt.Sprintf("INSERT INTO \"%s\" (%s) SELECT %s FROM \"%s\";", t.Name, strings.Join(tableInfo.ColumnNames(), ", "), strings.Join(tableInfo.ColumnNames(), ", "), tempTable.Name)
		return strings.Join([]string{createTableSQL, ins}, "\n"), nil
	})

	if err != nil {
		return "", err
	}

	sql = append(sql, s)

	return strings.Join(sql, "\n"), nil
}

func (p *SQLite) AddColumn(t fizz.Table) (string, error) {
	if len(t.Columns) == 0 {
		return "", errors.New("Not enough columns supplied!")
	}
	c := t.Columns[0]

	tableInfo, err := p.Schema.TableInfo(t.Name)
	if err != nil {
		return "", err
	}

	tableInfo.Columns = append(tableInfo.Columns, c)

	s := fmt.Sprintf("ALTER TABLE \"%s\" ADD COLUMN %s;", t.Name, p.buildColumn(c))
	return s, nil
}

func (p *SQLite) DropColumn(t fizz.Table) (string, error) {
	if len(t.Columns) < 1 {
		return "", errors.New("Not enough columns supplied!")
	}

	tableInfo, err := p.Schema.TableInfo(t.Name)
	if err != nil {
		return "", err
	}

	sql := []string{}
	droppedColumn := t.Columns[0]

	newColumns := []fizz.Column{}
	for _, c := range tableInfo.Columns {
		if c.Name != droppedColumn.Name {
			newColumns = append(newColumns, c)
		}
	}
	tableInfo.Columns = newColumns

	newIndexes := []fizz.Index{}
	for _, i := range tableInfo.Indexes {
		s, err := p.DropIndex(fizz.Table{
			Name:    tableInfo.Name,
			Indexes: []fizz.Index{i},
		})
		if err != nil {
			return "", err
		}
		sql = append(sql, s)
		if tableInfo.HasColumns(i.Columns...) {
			newIndexes = append(newIndexes, i)
		}
	}
	tableInfo.Indexes = newIndexes

	s, err := p.withTempTable(t.Name, func(tempTable fizz.Table) (string, error) {
		createTableSQL, err := p.CreateTable(*tableInfo)
		if err != nil {
			return "", err
		}

		s := fmt.Sprintf("INSERT INTO \"%s\" (%s) SELECT %s FROM \"%s\";", tableInfo.Name, strings.Join(tableInfo.ColumnNames(), ", "), strings.Join(tableInfo.ColumnNames(), ", "), tempTable.Name)

		return strings.Join([]string{createTableSQL, s}, "\n"), nil
	})

	if err != nil {
		return "", err
	}
	sql = append(sql, s)

	return strings.Join(sql, "\n"), nil
}

func (p *SQLite) RenameColumn(t fizz.Table) (string, error) {
	if len(t.Columns) < 2 {
		return "", errors.New("Not enough columns supplied!")
	}

	tableInfo, err := p.Schema.TableInfo(t.Name)
	if err != nil {
		return "", err
	}

	oldColumn := t.Columns[0]
	newColumn := t.Columns[1]

	sql := []string{}

	oldColumns := tableInfo.ColumnNames()
	for ic, c := range tableInfo.Columns {
		if c.Name == oldColumn.Name {
			tableInfo.Columns[ic].Name = newColumn.Name
		}
	}

	for _, i := range tableInfo.Indexes {
		s, err := p.DropIndex(fizz.Table{
			Name:    tableInfo.Name,
			Indexes: []fizz.Index{i},
		})
		if err != nil {
			return "", err
		}
		sql = append(sql, s)
		for ic, c := range i.Columns {
			if c == oldColumn.Name {
				i.Columns[ic] = newColumn.Name
			}
		}
	}

	s, err := p.withTempTable(t.Name, func(tempTable fizz.Table) (string, error) {
		createTableSQL, err := p.CreateTable(*tableInfo)
		if err != nil {
			return "", err
		}

		ins := fmt.Sprintf("INSERT INTO \"%s\" (%s) SELECT %s FROM \"%s\";", t.Name, strings.Join(tableInfo.ColumnNames(), ", "), strings.Join(oldColumns, ", "), tempTable.Name)
		return strings.Join([]string{createTableSQL, ins}, "\n"), nil
	})

	if err != nil {
		return "", err
	}

	sql = append(sql, s)

	return strings.Join(sql, "\n"), nil
}

func (p *SQLite) AddIndex(t fizz.Table) (string, error) {
	if len(t.Indexes) == 0 {
		return "", errors.New("Not enough indexes supplied!")
	}
	i := t.Indexes[0]
	s := fmt.Sprintf("CREATE INDEX \"%s\" ON \"%s\" (%s);", i.Name, t.Name, strings.Join(i.Columns, ", "))
	if i.Unique {
		s = strings.Replace(s, "CREATE", "CREATE UNIQUE", 1)
	}
	return s, nil
}

func (p *SQLite) DropIndex(t fizz.Table) (string, error) {
	if len(t.Indexes) == 0 {
		return "", errors.New("Not enough indexes supplied!")
	}
	i := t.Indexes[0]
	s := fmt.Sprintf("DROP INDEX IF EXISTS \"%s\";", i.Name)
	return s, nil
}

func (p *SQLite) RenameIndex(t fizz.Table) (string, error) {
	if len(t.Indexes) < 2 {
		return "", errors.New("Not enough indexes supplied!")
	}

	tableInfo, err := p.Schema.TableInfo(t.Name)
	if err != nil {
		return "", err
	}

	sql := []string{}

	oldIndex := t.Indexes[0]
	newIndex := t.Indexes[1]

	for _, ti := range tableInfo.Indexes {
		if ti.Name == oldIndex.Name {
			ti.Name = newIndex.Name
			newIndex = ti
			break
		}
	}

	s, err := p.DropIndex(fizz.Table{
		Name:    tableInfo.Name,
		Indexes: []fizz.Index{oldIndex},
	})

	if err != nil {
		return "", err
	}

	sql = append(sql, s)

	s, err = p.AddIndex(fizz.Table{
		Name:    t.Name,
		Indexes: []fizz.Index{newIndex},
	})

	if err != nil {
		return "", err
	}

	sql = append(sql, s)

	return strings.Join(sql, "\n"), nil
}

func (p *SQLite) withTempTable(table string, fn func(fizz.Table) (string, error)) (string, error) {
	tempTable := fizz.Table{Name: fmt.Sprintf("_%s_tmp", table)}

	sql := []string{fmt.Sprintf("ALTER TABLE \"%s\" RENAME TO \"%s\";", table, tempTable.Name)}
	s, err := fn(tempTable)
	if err != nil {
		return "", err
	}
	sql = append(sql, s, fmt.Sprintf("DROP TABLE \"%s\";", tempTable.Name))

	return strings.Join(sql, "\n"), nil
}

func (p *SQLite) buildColumn(c fizz.Column) string {
	s := fmt.Sprintf("\"%s\" %s", c.Name, p.colType(c))
	if c.Options["null"] == nil {
		s = fmt.Sprintf("%s NOT NULL", s)
	}
	if c.Options["default"] != nil {
		s = fmt.Sprintf("%s DEFAULT '%v'", s, c.Options["default"])
	}
	if c.Options["default_raw"] != nil {
		s = fmt.Sprintf("%s DEFAULT %s", s, c.Options["default_raw"])
	}
	return s
}

func (p *SQLite) colType(c fizz.Column) string {
	switch strings.ToLower(c.ColType) {
	case "uuid":
		return "char(36)"
	case "timestamp", "time", "datetime":
		return "DATETIME"
	case "boolean", "date":
		return "NUMERIC"
	case "string":
		return "TEXT"
	default:
		return c.ColType
	}
}
<commit_msg>removed a print statement<commit_after>package translators

import (
	"fmt"
	"strings"

	"github.com/pkg/errors"

	"github.com/markbates/pop/fizz"
)

type SQLite struct {
	Schema Schema
}

func NewSQLite(url string) *SQLite {
	return &SQLite{
		Schema: &sqliteSchema{
			URL:    url,
			schema: map[string]*fizz.Table{},
		},
	}
}

func (p *SQLite) CreateTable(t fizz.Table) (string, error) {
	sql := []string{}
	cols := []string{}
	var s string
	for _, c := range t.Columns {
		if c.Primary {
			switch strings.ToLower(c.ColType) {
			case "integer":
				s = fmt.Sprintf("\"%s\" INTEGER PRIMARY KEY AUTOINCREMENT", c.Name)
			case "uuid", "string":
				s = fmt.Sprintf("\"%s\" TEXT PRIMARY KEY", c.Name)
			}
		} else {
			s = p.buildColumn(c)
		}
		cols = append(cols, s)
	}
	s = fmt.Sprintf("CREATE TABLE \"%s\" (\n%s\n);", t.Name, strings.Join(cols, ",\n"))
	sql = append(sql, s)

	for _, i := range t.Indexes {
		s, err := p.AddIndex(fizz.Table{
			Name:    t.Name,
			Indexes: []fizz.Index{i},
		})
		if err != nil {
			return "", err
		}
		sql = append(sql, s)
	}

	return strings.Join(sql, "\n"), nil
}

func (p *SQLite) DropTable(t fizz.Table) (string, error) {
	p.Schema.Delete(t.Name)
	s := fmt.Sprintf("DROP TABLE \"%s\";", t.Name)
	return s, nil
}

func (p *SQLite) RenameTable(t []fizz.Table) (string, error) {
	if len(t) < 2 {
		return "", errors.New("Not enough table names supplied!")
	}
	oldName := t[0].Name
	newName := t[1].Name
	tableInfo, err := p.Schema.TableInfo(oldName)
	if err != nil {
		return "", err
	}
	tableInfo.Name = newName
	s := fmt.Sprintf("ALTER TABLE \"%s\" RENAME TO \"%s\";", oldName, newName)
	return s, nil
}

func (p *SQLite) ChangeColumn(t fizz.Table) (string, error) {
	tableInfo, err := p.Schema.TableInfo(t.Name)

	if err != nil {
		return "", err
	}

	for i := range tableInfo.Columns {
		if tableInfo.Columns[i].Name == t.Columns[0].Name {
			tableInfo.Columns[i] = t.Columns[0]
			break
		}
	}

	sql := []string{}
	s, err := p.withTempTable(t.Name, func(tempTable fizz.Table) (string, error) {
		createTableSQL, err := p.CreateTable(*tableInfo)
		if err != nil {
			return "", err
		}

		ins := fmt.Sprintf("INSERT INTO \"%s\" (%s) SELECT %s FROM \"%s\";", t.Name, strings.Join(tableInfo.ColumnNames(), ", "), strings.Join(tableInfo.ColumnNames(), ", "), tempTable.Name)
		return strings.Join([]string{createTableSQL, ins}, "\n"), nil
	})

	if err != nil {
		return "", err
	}

	sql = append(sql, s)

	return strings.Join(sql, "\n"), nil
}

func (p *SQLite) AddColumn(t fizz.Table) (string, error) {
	if len(t.Columns) == 0 {
		return "", errors.New("Not enough columns supplied!")
	}
	c := t.Columns[0]

	tableInfo, err := p.Schema.TableInfo(t.Name)
	if err != nil {
		return "", err
	}

	tableInfo.Columns = append(tableInfo.Columns, c)

	s := fmt.Sprintf("ALTER TABLE \"%s\" ADD COLUMN %s;", t.Name, p.buildColumn(c))
	return s, nil
}

func (p *SQLite) DropColumn(t fizz.Table) (string, error) {
	if len(t.Columns) < 1 {
		return "", errors.New("Not enough columns supplied!")
	}

	tableInfo, err := p.Schema.TableInfo(t.Name)
	if err != nil {
		return "", err
	}

	sql := []string{}
	droppedColumn := t.Columns[0]

	newColumns := []fizz.Column{}
	for _, c := range tableInfo.Columns {
		if c.Name != droppedColumn.Name {
			newColumns = append(newColumns, c)
		}
	}
	tableInfo.Columns = newColumns

	newIndexes := []fizz.Index{}
	for _, i := range tableInfo.Indexes {
		s, err := p.DropIndex(fizz.Table{
			Name:    tableInfo.Name,
			Indexes: []fizz.Index{i},
		})
		if err != nil {
			return "", err
		}
		sql = append(sql, s)
		if tableInfo.HasColumns(i.Columns...) {
			newIndexes = append(newIndexes, i)
		}
	}
	tableInfo.Indexes = newIndexes

	s, err := p.withTempTable(t.Name, func(tempTable fizz.Table) (string, error) {
		createTableSQL, err := p.CreateTable(*tableInfo)
		if err != nil {
			return "", err
		}

		s := fmt.Sprintf("INSERT INTO \"%s\" (%s) SELECT %s FROM \"%s\";", tableInfo.Name, strings.Join(tableInfo.ColumnNames(), ", "), strings.Join(tableInfo.ColumnNames(), ", "), tempTable.Name)

		return strings.Join([]string{createTableSQL, s}, "\n"), nil
	})

	if err != nil {
		return "", err
	}
	sql = append(sql, s)

	return strings.Join(sql, "\n"), nil
}

func (p *SQLite) RenameColumn(t fizz.Table) (string, error) {
	if len(t.Columns) < 2 {
		return "", errors.New("Not enough columns supplied!")
	}

	tableInfo, err := p.Schema.TableInfo(t.Name)
	if err != nil {
		return "", err
	}

	oldColumn := t.Columns[0]
	newColumn := t.Columns[1]

	sql := []string{}

	oldColumns := tableInfo.ColumnNames()
	for ic, c := range tableInfo.Columns {
		if c.Name == oldColumn.Name {
			tableInfo.Columns[ic].Name = newColumn.Name
		}
	}

	for _, i := range tableInfo.Indexes {
		s, err := p.DropIndex(fizz.Table{
			Name:    tableInfo.Name,
			Indexes: []fizz.Index{i},
		})
		if err != nil {
			return "", err
		}
		sql = append(sql, s)
		for ic, c := range i.Columns {
			if c == oldColumn.Name {
				i.Columns[ic] = newColumn.Name
			}
		}
	}

	s, err := p.withTempTable(t.Name, func(tempTable fizz.Table) (string, error) {
		createTableSQL, err := p.CreateTable(*tableInfo)
		if err != nil {
			return "", err
		}

		ins := fmt.Sprintf("INSERT INTO \"%s\" (%s) SELECT %s FROM \"%s\";", t.Name, strings.Join(tableInfo.ColumnNames(), ", "), strings.Join(oldColumns, ", "), tempTable.Name)
		return strings.Join([]string{createTableSQL, ins}, "\n"), nil
	})

	if err != nil {
		return "", err
	}

	sql = append(sql, s)

	return strings.Join(sql, "\n"), nil
}

func (p *SQLite) AddIndex(t fizz.Table) (string, error) {
	if len(t.Indexes) == 0 {
		return "", errors.New("Not enough indexes supplied!")
	}
	i := t.Indexes[0]
	s := fmt.Sprintf("CREATE INDEX \"%s\" ON \"%s\" (%s);", i.Name, t.Name, strings.Join(i.Columns, ", "))
	if i.Unique {
		s = strings.Replace(s, "CREATE", "CREATE UNIQUE", 1)
	}
	return s, nil
}

func (p *SQLite) DropIndex(t fizz.Table) (string, error) {
	if len(t.Indexes) == 0 {
		return "", errors.New("Not enough indexes supplied!")
	}
	i := t.Indexes[0]
	s := fmt.Sprintf("DROP INDEX IF EXISTS \"%s\";", i.Name)
	return s, nil
}

func (p *SQLite) RenameIndex(t fizz.Table) (string, error) {
	if len(t.Indexes) < 2 {
		return "", errors.New("Not enough indexes supplied!")
	}

	tableInfo, err := p.Schema.TableInfo(t.Name)
	if err != nil {
		return "", err
	}

	sql := []string{}

	oldIndex := t.Indexes[0]
	newIndex := t.Indexes[1]

	for _, ti := range tableInfo.Indexes {
		if ti.Name == oldIndex.Name {
			ti.Name = newIndex.Name
			newIndex = ti
			break
		}
	}

	s, err := p.DropIndex(fizz.Table{
		Name:    tableInfo.Name,
		Indexes: []fizz.Index{oldIndex},
	})

	if err != nil {
		return "", err
	}

	sql = append(sql, s)

	s, err = p.AddIndex(fizz.Table{
		Name:    t.Name,
		Indexes: []fizz.Index{newIndex},
	})

	if err != nil {
		return "", err
	}

	sql = append(sql, s)

	return strings.Join(sql, "\n"), nil
}

func (p *SQLite) withTempTable(table string, fn func(fizz.Table) (string, error)) (string, error) {
	tempTable := fizz.Table{Name: fmt.Sprintf("_%s_tmp", table)}

	sql := []string{fmt.Sprintf("ALTER TABLE \"%s\" RENAME TO \"%s\";", table, tempTable.Name)}
	s, err := fn(tempTable)
	if err != nil {
		return "", err
	}
	sql = append(sql, s, fmt.Sprintf("DROP TABLE \"%s\";", tempTable.Name))

	return strings.Join(sql, "\n"), nil
}

func (p *SQLite) buildColumn(c fizz.Column) string {
	s := fmt.Sprintf("\"%s\" %s", c.Name, p.colType(c))
	if c.Options["null"] == nil {
		s = fmt.Sprintf("%s NOT NULL", s)
	}
	if c.Options["default"] != nil {
		s = fmt.Sprintf("%s DEFAULT '%v'", s, c.Options["default"])
	}
	if c.Options["default_raw"] != nil {
		s = fmt.Sprintf("%s DEFAULT %s", s, c.Options["default_raw"])
	}
	return s
}

func (p *SQLite) colType(c fizz.Column) string {
	switch strings.ToLower(c.ColType) {
	case "uuid":
		return "char(36)"
	case "timestamp", "time", "datetime":
		return "DATETIME"
	case "boolean", "date":
		return "NUMERIC"
	case "string":
		return "TEXT"
	default:
		return c.ColType
	}
}
<|endoftext|>


<commit_before>package files

import (
	"os"
	"time"
	"path/filepath"
	"strings"
	"fmt"
	"github.com/Unknwon/com"
)

const (
	IsDirectory = iota
	IsRegular
	IsSymlink
)

type sysFile struct {
	fType  int
	fName  string
	fLink  string
	fSize  int64
	fMtime time.Time
	fPerm  os.FileMode
}

type F struct {
	files []*sysFile
}

//  _______ _________ _        _        _______
// (  ____ \\__   __/( (    /|( \      (  ___  )|\     /|
// | (    \/   ) (   |  \  ( || (      | (   ) || )   ( |
// | (_____    | |   |   \ | || |      | |   | || |   | |
// (_____  )   | |   | (\ \) || |      | |   | |( (   ) )
//       ) |   | |   | | \   || |      | |   | | \ \_/ /
// /\____) |___) (___| )  \  || (____/\| (___) |  \   /
// \_______)\_______/|/    )_)(_______/(_______)   \_/
//
func (f *F) visit(path string, file os.FileInfo, err error) error {
	if file == nil {
		return err
	}
	var tp int
	if file.IsDir() {
		tp = IsDirectory
	} else if (file.Mode() & os.ModeSymlink) > 0 {
		tp = IsSymlink
	} else {
		tp = IsRegular
	}
	inoFile := &sysFile{
		fName:  path,
		fType:  tp,
		fPerm:  file.Mode(),
		fMtime: file.ModTime(),
		fSize:  file.Size(),
	}
	f.files = append(f.files, inoFile)
	return nil
}

func (f *F) MergeFolderPath(sourceDir string, targetDir string) error {

	source := F{
		files: make([]*sysFile, 0),
	}
	err := filepath.Walk(sourceDir,
		func(path string, f os.FileInfo, err error) error {
			return source.visit(path, f, err)
		},
	)
	if err != nil {
		return err
	}
	target := F{
		files: make([]*sysFile, 0),
	}
	err = filepath.Walk(targetDir, func(path string, f os.FileInfo, err error) error {
		return target.visit(path, f, err)
	})
	if err != nil {
		return err
	}

	for _, v := range source.files {

		if com.IsFile(v.fName) == true {

			tmp1 := strings.Split(v.fName, "\\")
			sourceName := tmp1[len(tmp1)-1]
			for _, r := range target.files {
				if com.IsFile(r.fName) == true {
					tmp2 := strings.Split(r.fName, "\\")
					targetName := tmp2[len(tmp2)-1]
					if sourceName == targetName {
						fmt.Printf("the same file: %s\n", sourceName)
						com.Copy(v.fName, r.fName)
					}
				}
			}

		}
	}

	return nil
}
<commit_msg>refactor fileCopy<commit_after>package files

import (
	"os"
	"time"
	"path/filepath"
	"strings"
	"fmt"
	"github.com/Unknwon/com"
)

const (
	IsDirectory = iota
	IsRegular
	IsSymlink
)

type sysFile struct {
	fType  int
	fName  string
	fLink  string
	fSize  int64
	fMTime time.Time
	fPerm  os.FileMode
}

type F struct {
	files []*sysFile
}

func (f *F) visit(path string, file os.FileInfo, err error) error {
	if file == nil {
		return err
	}
	var tp int
	if file.IsDir() {
		tp = IsDirectory
	} else if (file.Mode() & os.ModeSymlink) > 0 {
		tp = IsSymlink
	} else {
		tp = IsRegular
	}
	inoFile := &sysFile{
		fName:  path,
		fType:  tp,
		fPerm:  file.Mode(),
		fMTime: file.ModTime(),
		fSize:  file.Size(),
	}
	f.files = append(f.files, inoFile)
	return nil
}

func (f *F) MergeFolderPath(sourceDir string, targetDir string) error {

	source := F{
		files: make([]*sysFile, 0),
	}
	err := filepath.Walk(sourceDir,
		func(path string, f os.FileInfo, err error) error {
			return source.visit(path, f, err)
		},
	)
	if err != nil {
		return err
	}
	target := F{
		files: make([]*sysFile, 0),
	}
	err = filepath.Walk(targetDir, func(path string, f os.FileInfo, err error) error {
		return target.visit(path, f, err)
	})
	if err != nil {
		return err
	}

	for _, v := range source.files {

		if com.IsFile(v.fName) == true {

			tmp1 := strings.Split(v.fName, "\\")
			sourceName := tmp1[len(tmp1)-1]
			for _, r := range target.files {
				if com.IsFile(r.fName) == true {
					tmp2 := strings.Split(r.fName, "\\")
					targetName := tmp2[len(tmp2)-1]
					if sourceName == targetName {
						fmt.Printf("the same file: %s\n", sourceName)
						com.Copy(v.fName, r.fName)
					}
				}
			}

		}
	}

	return nil
}
<|endoftext|>


<commit_before><commit_msg>fixed log message<commit_after><|endoftext|>


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

package pmon

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"

	sys "golang.org/x/sys/unix"
)

type Process struct {
	W    io.Writer
	Freq time.Duration

	quit chan struct{}

	fc chan func() error
	ec chan error

	cmd  *exec.Cmd
	proc *os.Process

	start func() error
	stop  func() error
}

func New(cmd string, args ...string) *Process {
	c := exec.Command(cmd, args...)
	c.Stdin = os.Stdin
	c.Stdout = os.Stdout
	c.Stderr = os.Stderr

	c.SysProcAttr = &sys.SysProcAttr{
		Ptrace:  true,
		Setpgid: true,
	}

	proc := &Process{
		cmd:  c,
		Freq: 1 * time.Second,
		W:    ioutil.Discard,
		quit: make(chan struct{}),
		fc:   make(chan func() error),
		ec:   make(chan error),

		start: c.Start,
		stop: func() error {
			pgid, err := sys.Getpgid(c.Process.Pid)
			if err != nil {
				return fmt.Errorf("could not get process group of pid=%d: %w", c.Process.Pid, err)
			}
			err = sys.Kill(-pgid, sys.SIGKILL) // note the minus sign
			if err != nil {
				return fmt.Errorf("could not kill process group %d: %w", pgid, err)
			}

			return nil
		},
	}

	go ptraceRun(proc.fc, proc.ec)
	return proc
}

func (p *Process) Run() error {
	switch {
	case p.cmd != nil:
		return p.runCmd()
	default:
		return p.runPID()
	}
}

func (p *Process) runCmd() error {
	defer close(p.quit)
	defer func() {
		if w, ok := p.W.(interface{ Flush() error }); ok {
			_ = w.Flush()
		}
	}()

	err := p.ptraceRun(p.cmd.Start)
	if err != nil {
		return fmt.Errorf("could not start process: %w", err)
	}

	start := time.Now()

	pid := p.cmd.Process.Pid
	collector, err := newCollector(pid)
	if err != nil {
		return fmt.Errorf("could not create collector: %w", err)
	}
	defer collector.Close()

	_, err = fmt.Fprintf(p.W,
		"# pmon: %s\n# freq: %v\n# format: %#v\n# start: %v\n",
		strings.Join(p.cmd.Args, " "),
		p.Freq,
		Infos{},
		start,
	)
	if err != nil {
		return fmt.Errorf("error writing log-file header: %w", err)
	}

	defer func() {
		stop := time.Now()
		delta := time.Since(start)
		_, _ = fmt.Fprintf(p.W,
			"# elapsed: %v\n# stop: %v\n",
			delta,
			stop,
		)
	}()

	err = p.wait(pid, 0)
	if err != nil {
		return fmt.Errorf("waiting for target execve failed: %w", err)
	}

	err = p.ptraceDetach(pid)
	if err != nil {
		return fmt.Errorf("could not ptrace-detach pid=%d: %w", pid, err)
	}

	go p.monitor(collector)

	log.Printf(
		"monitoring... (pid=%d, freq=%v)\n",
		p.cmd.Process.Pid,
		p.Freq,
	)
	err = p.cmd.Wait()
	if err != nil {
		return fmt.Errorf("could not wait for pid=%d: %w", pid, err)
	}

	return nil
}

func (p *Process) runPID() error {
	start := time.Now()

	pid := p.proc.Pid
	collector, err := newCollector(pid)
	if err != nil {
		return fmt.Errorf("could not create collector: %w", err)
	}
	defer collector.Close()

	_, err = fmt.Fprintf(p.W,
		"# pmon: %s\n# freq: %v\n# format: %#v\n# start: %v\n",
		p.cmdline(pid),
		p.Freq,
		Infos{},
		start,
	)
	if err != nil {
		return fmt.Errorf("error writing log-file header: %w", err)
	}

	defer func() {
		stop := time.Now()
		delta := time.Since(start)
		_, _ = fmt.Fprintf(p.W,
			"# elapsed: %v\n# stop: %v\n",
			delta,
			stop,
		)
	}()

	go p.monitor(collector)

	log.Printf(
		"monitoring... (pid=%d, freq=%v)\n",
		p.proc.Pid,
		p.Freq,
	)
	<-p.quit
	if err != nil {
		return fmt.Errorf("could not wait for pid=%d: %w", pid, err)
	}

	return nil
}

func (p *Process) Kill() error {
	return p.stop()
}

func (p *Process) monitor(c *collector) {
	p.collect(c)
	tick := time.Tick(p.Freq)
	for {
		select {
		case <-tick:
			p.collect(c)
		case <-p.quit:
			return
		}
	}
}

func (p *Process) collect(c *collector) {

	if p.cmd != nil && p.cmd.ProcessState != nil {
		// process already stopped. nothing to collect.
		return
	}

	infos, err := c.collect()
	if err != nil {
		log.Printf("error collecting: %+v", err)
		return
	}

	fmt.Fprintf(
		p.W, "%f %f %f %d %d %d %d %d %d %d\n",
		milliseconds(infos.CPU), milliseconds(infos.UTime), milliseconds(infos.STime),
		infos.VMem, infos.RSS,
		infos.Threads,
		infos.Rchar, infos.Wchar,
		infos.Rdisk, infos.Wdisk,
	)
}

func milliseconds(t time.Duration) float64 {
	return t.Seconds() * 1e3
}
<commit_msg>pmon: use stable time.Layout for start/stop comments<commit_after>// Copyright 2015 The pmon Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package pmon

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"

	sys "golang.org/x/sys/unix"
)

type Process struct {
	W    io.Writer
	Freq time.Duration

	quit chan struct{}

	fc chan func() error
	ec chan error

	cmd  *exec.Cmd
	proc *os.Process

	start func() error
	stop  func() error
}

func New(cmd string, args ...string) *Process {
	c := exec.Command(cmd, args...)
	c.Stdin = os.Stdin
	c.Stdout = os.Stdout
	c.Stderr = os.Stderr

	c.SysProcAttr = &sys.SysProcAttr{
		Ptrace:  true,
		Setpgid: true,
	}

	proc := &Process{
		cmd:  c,
		Freq: 1 * time.Second,
		W:    ioutil.Discard,
		quit: make(chan struct{}),
		fc:   make(chan func() error),
		ec:   make(chan error),

		start: c.Start,
		stop: func() error {
			pgid, err := sys.Getpgid(c.Process.Pid)
			if err != nil {
				return fmt.Errorf("could not get process group of pid=%d: %w", c.Process.Pid, err)
			}
			err = sys.Kill(-pgid, sys.SIGKILL) // note the minus sign
			if err != nil {
				return fmt.Errorf("could not kill process group %d: %w", pgid, err)
			}

			return nil
		},
	}

	go ptraceRun(proc.fc, proc.ec)
	return proc
}

func (p *Process) Run() error {
	switch {
	case p.cmd != nil:
		return p.runCmd()
	default:
		return p.runPID()
	}
}

func (p *Process) runCmd() error {
	defer close(p.quit)
	defer func() {
		if w, ok := p.W.(interface{ Flush() error }); ok {
			_ = w.Flush()
		}
	}()

	err := p.ptraceRun(p.cmd.Start)
	if err != nil {
		return fmt.Errorf("could not start process: %w", err)
	}

	start := time.Now()

	pid := p.cmd.Process.Pid
	collector, err := newCollector(pid)
	if err != nil {
		return fmt.Errorf("could not create collector: %w", err)
	}
	defer collector.Close()

	_, err = fmt.Fprintf(p.W,
		"# pmon: %s\n# freq: %v\n# format: %#v\n# start: %v\n",
		strings.Join(p.cmd.Args, " "),
		p.Freq,
		Infos{},
		start.Format(time.RFC3339Nano),
	)
	if err != nil {
		return fmt.Errorf("error writing log-file header: %w", err)
	}

	defer func() {
		stop := time.Now()
		delta := time.Since(start)
		_, _ = fmt.Fprintf(p.W,
			"# elapsed: %v\n# stop: %v\n",
			delta,
			stop.Format(time.RFC3339Nano),
		)
	}()

	err = p.wait(pid, 0)
	if err != nil {
		return fmt.Errorf("waiting for target execve failed: %w", err)
	}

	err = p.ptraceDetach(pid)
	if err != nil {
		return fmt.Errorf("could not ptrace-detach pid=%d: %w", pid, err)
	}

	go p.monitor(collector)

	log.Printf(
		"monitoring... (pid=%d, freq=%v)\n",
		p.cmd.Process.Pid,
		p.Freq,
	)
	err = p.cmd.Wait()
	if err != nil {
		return fmt.Errorf("could not wait for pid=%d: %w", pid, err)
	}

	return nil
}

func (p *Process) runPID() error {
	start := time.Now()

	pid := p.proc.Pid
	collector, err := newCollector(pid)
	if err != nil {
		return fmt.Errorf("could not create collector: %w", err)
	}
	defer collector.Close()

	_, err = fmt.Fprintf(p.W,
		"# pmon: %s\n# freq: %v\n# format: %#v\n# start: %v\n",
		p.cmdline(pid),
		p.Freq,
		Infos{},
		start.Format(time.RFC3339Nano),
	)
	if err != nil {
		return fmt.Errorf("error writing log-file header: %w", err)
	}

	defer func() {
		stop := time.Now()
		delta := time.Since(start)
		_, _ = fmt.Fprintf(p.W,
			"# elapsed: %v\n# stop: %v\n",
			delta,
			stop.Format(time.RFC3339Nano),
		)
	}()

	go p.monitor(collector)

	log.Printf(
		"monitoring... (pid=%d, freq=%v)\n",
		p.proc.Pid,
		p.Freq,
	)
	<-p.quit
	if err != nil {
		return fmt.Errorf("could not wait for pid=%d: %w", pid, err)
	}

	return nil
}

func (p *Process) Kill() error {
	return p.stop()
}

func (p *Process) monitor(c *collector) {
	p.collect(c)
	tick := time.Tick(p.Freq)
	for {
		select {
		case <-tick:
			p.collect(c)
		case <-p.quit:
			return
		}
	}
}

func (p *Process) collect(c *collector) {

	if p.cmd != nil && p.cmd.ProcessState != nil {
		// process already stopped. nothing to collect.
		return
	}

	infos, err := c.collect()
	if err != nil {
		log.Printf("error collecting: %+v", err)
		return
	}

	fmt.Fprintf(
		p.W, "%f %f %f %d %d %d %d %d %d %d\n",
		milliseconds(infos.CPU), milliseconds(infos.UTime), milliseconds(infos.STime),
		infos.VMem, infos.RSS,
		infos.Threads,
		infos.Rchar, infos.Wchar,
		infos.Rdisk, infos.Wdisk,
	)
}

func milliseconds(t time.Duration) float64 {
	return t.Seconds() * 1e3
}
<|endoftext|>


<commit_before>package main

import "fmt"
import _ "github.com/jaybennett89/thorium-go/process"

func main(){
	fmt.Println("hello world")
}
<commit_msg>cleanup<commit_after><|endoftext|>