<commit_before>package sia

import (
	"errors"
	"net"
	"strconv"
	"time"
)

const (
	timeout   = time.Second * 5
	maxMsgLen = 1 << 16
)

// A NetAddress contains the information needed to contact a peer over TCP.
type NetAddress struct {
	Host string
	Port uint16
}

// String returns the NetAddress as a string, concatentating the hostname and
// port number.
func (na *NetAddress) String() string {
	return net.JoinHostPort(na.Host, strconv.Itoa(int(na.Port)))
}

// call establishes a TCP connection to the NetAddress, calls the provided
// function on it, and closes the connection.
func (na *NetAddress) Call(fn func(net.Conn) error) error {
	conn, err := net.Dial("tcp", na.String())
	if err != nil {
		return err
	}
	defer conn.Close()
	return fn(conn)
}

// TBD
var BootstrapPeers = []NetAddress{}

// A TCPServer sends and receives messages. It also maintains an address book
// of peers to broadcast to and make requests of.
type TCPServer struct {
	net.Listener
	myAddr      NetAddress
	addressbook map[NetAddress]struct{}
	handlerMap  map[byte]func(net.Conn, []byte) error
}

func NewTCPServer(port uint16) (tcps *TCPServer, err error) {
	tcpServ, err := net.Listen("tcp", ":"+strconv.Itoa(int(port)))
	if err != nil {
		return
	}
	tcps = &TCPServer{
		Listener:    tcpServ,
		addressbook: make(map[NetAddress]struct{}),
		handlerMap:  make(map[byte]func(net.Conn, []byte) error),
	}
	// default handlers
	tcps.handlerMap['H'] = tcps.sendHostname
	tcps.handlerMap['P'] = tcps.sendHostname

	// spawn listener
	go tcps.listen()
	return
}

// Register registers a message type with a message handler. The existing
// handler for that type will be overwritten.
func (tcps *TCPServer) Register(t byte, fn func(net.Conn, []byte) error) {
	tcps.handlerMap[t] = fn
}

// listen runs in the background, accepting incoming connections and serving
// them. listen will return after TCPServer.Close() is called, because the
// Accept() call will fail.
func (tcps *TCPServer) listen() {
	for {
		conn, err := tcps.Accept()
		if err != nil {
			return
		}
		// it is the handler's responsibility to close the connection
		go tcps.handleConn(conn)
	}
}

// handleConn reads header data from a connection, unmarshals the data
// structures it contains, and routes the data to other functions for
// processing.
func (tcps *TCPServer) handleConn(conn net.Conn) {
	defer conn.Close()
	var (
		msgType   []byte = make([]byte, 1)
		msgLenBuf []byte = make([]byte, 4)
		msgData   []byte // length determined by msgLen
	)
	// TODO: make this DRYer?
	if n, err := conn.Read(msgType); err != nil || n != 1 {
		// TODO: log error
		return
	}
	if n, err := conn.Read(msgLenBuf); err != nil || n != 4 {
		// TODO: log error
		return
	}
	msgLen := DecUint64(msgLenBuf)
	if msgLen > maxMsgLen {
		// TODO: log error
		return
	}
	msgData = make([]byte, msgLen)
	if n, err := conn.Read(msgData); err != nil || uint64(n) != msgLen {
		// TODO: log error
		return
	}

	// call registered handler for this message type
	if fn, ok := tcps.handlerMap[msgType[0]]; ok {
		fn(conn, msgData)
		// TODO: log error
	}
	return
}

// sendHostname replies to the send with the sender's external IP.
func (tcps *TCPServer) sendHostname(conn net.Conn, _ []byte) error {
	_, err := conn.Write([]byte(conn.RemoteAddr().String()))
	return err
}

// sharePeers transmits at most 'num' peers over the connection.
// TODO: choose random peers?
func (tcps *TCPServer) sharePeers(conn net.Conn, msgData []byte) error {
	if len(msgData) != 1 {
		return errors.New("invalid number of peers")
	}
	num := msgData[0]
	var addrs []NetAddress
	for addr := range tcps.addressbook {
		if num == 0 {
			break
		}
		addrs = append(addrs, addr)
		num--
	}
	_, err := conn.Write(Marshal(addrs))
	return err
}

// Ping returns whether a NetAddress is reachable. It accomplishes this by
// initiating a TCP connection and immediately closes it. This is pretty
// unsophisticated. I'll add a Pong later.
func (tcps *TCPServer) Ping(addr NetAddress) bool {
	conn, err := net.Dial("tcp", addr.String())
	if err != nil {
		return false
	}
	conn.Close()
	return true
}

// learnHostname learns the external IP of the TCPServer.
func (tcps *TCPServer) learnHostname(conn net.Conn) (err error) {
	// send hostname request
	if _, err = conn.Write([]byte{'H', 0}); err != nil {
		return
	}
	// read response
	buf := make([]byte, 128)
	n, err := conn.Read(buf)
	if err != nil {
		return
	}
	// TODO: try to ping ourselves?
	host, portStr, err := net.SplitHostPort(string(buf[:n]))
	if err != nil {
		return
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return
	}
	tcps.myAddr = NetAddress{host, uint16(port)}
	return
}

// requestPeers queries a peer for additional peers, and adds any new peers to
// the address book.
func (tcps *TCPServer) requestPeers(conn net.Conn) (err error) {
	// request 10 peers
	if _, err = conn.Write([]byte{'P', 1, 10}); err != nil {
		return
	}
	// read response
	buf := make([]byte, 1024)
	n, err := conn.Read(buf)
	if err != nil {
		return
	}
	var addrs []NetAddress
	if err = Unmarshal(buf[:n], &addrs); err != nil {
		return
	}
	// add peers
	// TODO: make sure we don't add ourself
	for _, addr := range addrs {
		if tcps.Ping(addr) {
			tcps.addressbook[addr] = struct{}{}
		}
	}
	return
}

// Bootstrap discovers the external IP of the TCPServer, requests peers from
// the initial peer list, and announces itself to those peers.
func (tcps *TCPServer) Bootstrap() (err error) {
	// populate initial peer list
	for _, addr := range BootstrapPeers {
		if tcps.Ping(addr) {
			tcps.addressbook[addr] = struct{}{}
		}
	}

	// learn hostname
	for addr := range tcps.addressbook {
		if addr.Call(tcps.learnHostname) == nil {
			break
		}
	}
	// request peers
	// TODO: maybe iterate until we have enough new peers?
	for addr := range tcps.addressbook {
		addr.Call(tcps.requestPeers)
	}
	// TODO: announce ourselves to new peers
	return
}
<commit_msg>simplify handleConn a bit<commit_after>package sia

import (
	"errors"
	"net"
	"strconv"
	"time"
)

const (
	timeout   = time.Second * 5
	maxMsgLen = 1 << 16
)

// A NetAddress contains the information needed to contact a peer over TCP.
type NetAddress struct {
	Host string
	Port uint16
}

// String returns the NetAddress as a string, concatentating the hostname and
// port number.
func (na *NetAddress) String() string {
	return net.JoinHostPort(na.Host, strconv.Itoa(int(na.Port)))
}

// call establishes a TCP connection to the NetAddress, calls the provided
// function on it, and closes the connection.
func (na *NetAddress) Call(fn func(net.Conn) error) error {
	conn, err := net.Dial("tcp", na.String())
	if err != nil {
		return err
	}
	defer conn.Close()
	return fn(conn)
}

// TBD
var BootstrapPeers = []NetAddress{}

// A TCPServer sends and receives messages. It also maintains an address book
// of peers to broadcast to and make requests of.
type TCPServer struct {
	net.Listener
	myAddr      NetAddress
	addressbook map[NetAddress]struct{}
	handlerMap  map[byte]func(net.Conn, []byte) error
}

func NewTCPServer(port uint16) (tcps *TCPServer, err error) {
	tcpServ, err := net.Listen("tcp", ":"+strconv.Itoa(int(port)))
	if err != nil {
		return
	}
	tcps = &TCPServer{
		Listener:    tcpServ,
		addressbook: make(map[NetAddress]struct{}),
		handlerMap:  make(map[byte]func(net.Conn, []byte) error),
	}
	// default handlers
	tcps.handlerMap['H'] = tcps.sendHostname
	tcps.handlerMap['P'] = tcps.sendHostname

	// spawn listener
	go tcps.listen()
	return
}

// Register registers a message type with a message handler. The existing
// handler for that type will be overwritten.
func (tcps *TCPServer) Register(t byte, fn func(net.Conn, []byte) error) {
	tcps.handlerMap[t] = fn
}

// listen runs in the background, accepting incoming connections and serving
// them. listen will return after TCPServer.Close() is called, because the
// Accept() call will fail.
func (tcps *TCPServer) listen() {
	for {
		conn, err := tcps.Accept()
		if err != nil {
			return
		}
		// it is the handler's responsibility to close the connection
		go tcps.handleConn(conn)
	}
}

// handleConn reads header data from a connection, unmarshals the data
// structures it contains, and routes the data to other functions for
// processing.
func (tcps *TCPServer) handleConn(conn net.Conn) {
	defer conn.Close()
	var (
		msgHead []byte = make([]byte, 5)
		msgData []byte // length determined by msgHead
	)
	if n, err := conn.Read(msgHead); err != nil || n != 5 {
		// TODO: log error
		return
	}
	msgLen := DecUint64(msgHead[1:])
	if msgLen > maxMsgLen {
		// TODO: log error
		return
	}
	msgData = make([]byte, msgLen)
	if n, err := conn.Read(msgData); err != nil || uint64(n) != msgLen {
		// TODO: log error
		return
	}

	// call registered handler for this message type
	if fn, ok := tcps.handlerMap[msgHead[0]]; ok {
		fn(conn, msgData)
		// TODO: log error
	}
	return
}

// sendHostname replies to the send with the sender's external IP.
func (tcps *TCPServer) sendHostname(conn net.Conn, _ []byte) error {
	_, err := conn.Write([]byte(conn.RemoteAddr().String()))
	return err
}

// sharePeers transmits at most 'num' peers over the connection.
// TODO: choose random peers?
func (tcps *TCPServer) sharePeers(conn net.Conn, msgData []byte) error {
	if len(msgData) != 1 {
		return errors.New("invalid number of peers")
	}
	num := msgData[0]
	var addrs []NetAddress
	for addr := range tcps.addressbook {
		if num == 0 {
			break
		}
		addrs = append(addrs, addr)
		num--
	}
	_, err := conn.Write(Marshal(addrs))
	return err
}

// Ping returns whether a NetAddress is reachable. It accomplishes this by
// initiating a TCP connection and immediately closes it. This is pretty
// unsophisticated. I'll add a Pong later.
func (tcps *TCPServer) Ping(addr NetAddress) bool {
	conn, err := net.Dial("tcp", addr.String())
	if err != nil {
		return false
	}
	conn.Close()
	return true
}

// learnHostname learns the external IP of the TCPServer.
func (tcps *TCPServer) learnHostname(conn net.Conn) (err error) {
	// send hostname request
	if _, err = conn.Write([]byte{'H', 0}); err != nil {
		return
	}
	// read response
	buf := make([]byte, 128)
	n, err := conn.Read(buf)
	if err != nil {
		return
	}
	// TODO: try to ping ourselves?
	host, portStr, err := net.SplitHostPort(string(buf[:n]))
	if err != nil {
		return
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return
	}
	tcps.myAddr = NetAddress{host, uint16(port)}
	return
}

// requestPeers queries a peer for additional peers, and adds any new peers to
// the address book.
func (tcps *TCPServer) requestPeers(conn net.Conn) (err error) {
	// request 10 peers
	if _, err = conn.Write([]byte{'P', 1, 10}); err != nil {
		return
	}
	// read response
	buf := make([]byte, 1024)
	n, err := conn.Read(buf)
	if err != nil {
		return
	}
	var addrs []NetAddress
	if err = Unmarshal(buf[:n], &addrs); err != nil {
		return
	}
	// add peers
	// TODO: make sure we don't add ourself
	for _, addr := range addrs {
		if tcps.Ping(addr) {
			tcps.addressbook[addr] = struct{}{}
		}
	}
	return
}

// Bootstrap discovers the external IP of the TCPServer, requests peers from
// the initial peer list, and announces itself to those peers.
func (tcps *TCPServer) Bootstrap() (err error) {
	// populate initial peer list
	for _, addr := range BootstrapPeers {
		if tcps.Ping(addr) {
			tcps.addressbook[addr] = struct{}{}
		}
	}

	// learn hostname
	for addr := range tcps.addressbook {
		if addr.Call(tcps.learnHostname) == nil {
			break
		}
	}
	// request peers
	// TODO: maybe iterate until we have enough new peers?
	for addr := range tcps.addressbook {
		addr.Call(tcps.requestPeers)
	}
	// TODO: announce ourselves to new peers
	return
}
<|endoftext|>


<commit_before>package signal

import (
	"os"
	gosignal "os/signal"
	"sync/atomic"
	"syscall"

	log "github.com/Sirupsen/logrus"
)

// Trap sets up a simplified signal "trap", appropriate for common
// behavior expected from a vanilla unix command-line tool in general
// (and the Docker engine in particular).
//
// * If SIGINT or SIGTERM are received, `cleanup` is called, then the process is terminated.
// * If SIGINT or SIGTERM are repeated 3 times before cleanup is complete, then cleanup is
// skipped and the process terminated directly.
// * If "DEBUG" is set in the environment, SIGQUIT causes an exit without cleanup.
//
func Trap(cleanup func()) {
	c := make(chan os.Signal, 1)
	signals := []os.Signal{os.Interrupt, syscall.SIGTERM}
	if os.Getenv("DEBUG") == "" {
		signals = append(signals, syscall.SIGQUIT)
	}
	gosignal.Notify(c, signals...)
	go func() {
		interruptCount := uint32(0)
		for sig := range c {
			go func(sig os.Signal) {
				log.Infof("Received signal '%v', starting shutdown of docker...", sig)
				switch sig {
				case os.Interrupt, syscall.SIGTERM:
					// If the user really wants to interrupt, let him do so.
					if atomic.LoadUint32(&interruptCount) < 3 {
						// Initiate the cleanup only once
						if atomic.AddUint32(&interruptCount, 1) == 1 {
							// Call cleanup handler
							cleanup()
							os.Exit(0)
						} else {
							return
						}
					} else {
						log.Infof("Force shutdown of docker, interrupting cleanup")
					}
				case syscall.SIGQUIT:
				}
				os.Exit(128 + int(sig.(syscall.Signal)))
			}(sig)
		}
	}()
}
<commit_msg>Replace aliased imports of logrus, fixes #11762<commit_after>package signal

import (
	"os"
	gosignal "os/signal"
	"sync/atomic"
	"syscall"

	"github.com/Sirupsen/logrus"
)

// Trap sets up a simplified signal "trap", appropriate for common
// behavior expected from a vanilla unix command-line tool in general
// (and the Docker engine in particular).
//
// * If SIGINT or SIGTERM are received, `cleanup` is called, then the process is terminated.
// * If SIGINT or SIGTERM are repeated 3 times before cleanup is complete, then cleanup is
// skipped and the process terminated directly.
// * If "DEBUG" is set in the environment, SIGQUIT causes an exit without cleanup.
//
func Trap(cleanup func()) {
	c := make(chan os.Signal, 1)
	signals := []os.Signal{os.Interrupt, syscall.SIGTERM}
	if os.Getenv("DEBUG") == "" {
		signals = append(signals, syscall.SIGQUIT)
	}
	gosignal.Notify(c, signals...)
	go func() {
		interruptCount := uint32(0)
		for sig := range c {
			go func(sig os.Signal) {
				logrus.Infof("Received signal '%v', starting shutdown of docker...", sig)
				switch sig {
				case os.Interrupt, syscall.SIGTERM:
					// If the user really wants to interrupt, let him do so.
					if atomic.LoadUint32(&interruptCount) < 3 {
						// Initiate the cleanup only once
						if atomic.AddUint32(&interruptCount, 1) == 1 {
							// Call cleanup handler
							cleanup()
							os.Exit(0)
						} else {
							return
						}
					} else {
						logrus.Infof("Force shutdown of docker, interrupting cleanup")
					}
				case syscall.SIGQUIT:
				}
				os.Exit(128 + int(sig.(syscall.Signal)))
			}(sig)
		}
	}()
}
<|endoftext|>


<commit_before>// Package smaz is an implementation of the smaz library (https://github.com/antirez/smaz) for compressing
// small strings.
package smaz

import (
	"bytes"
	"errors"

	"github.com/cespare/go-smaz/trie"
)

var codeStrings = []string{" ",
	"the", "e", "t", "a", "of", "o", "and", "i", "n", "s", "e ", "r", " th",
	" t", "in", "he", "th", "h", "he ", "to", "\r\n", "l", "s ", "d", " a", "an",
	"er", "c", " o", "d ", "on", " of", "re", "of ", "t ", ", ", "is", "u", "at",
	"   ", "n ", "or", "which", "f", "m", "as", "it", "that", "\n", "was", "en",
	"  ", " w", "es", " an", " i", "\r", "f ", "g", "p", "nd", " s", "nd ", "ed ",
	"w", "ed", "http://", "for", "te", "ing", "y ", "The", " c", "ti", "r ", "his",
	"st", " in", "ar", "nt", ",", " to", "y", "ng", " h", "with", "le", "al", "to ",
	"b", "ou", "be", "were", " b", "se", "o ", "ent", "ha", "ng ", "their", "\"",
	"hi", "from", " f", "in ", "de", "ion", "me", "v", ".", "ve", "all", "re ",
	"ri", "ro", "is ", "co", "f t", "are", "ea", ". ", "her", " m", "er ", " p",
	"es ", "by", "they", "di", "ra", "ic", "not", "s, ", "d t", "at ", "ce", "la",
	"h ", "ne", "as ", "tio", "on ", "n t", "io", "we", " a ", "om", ", a", "s o",
	"ur", "li", "ll", "ch", "had", "this", "e t", "g ", "e\r\n", " wh", "ere",
	" co", "e o", "a ", "us", " d", "ss", "\n\r\n", "\r\n\r", "=\"", " be", " e",
	"s a", "ma", "one", "t t", "or ", "but", "el", "so", "l ", "e s", "s,", "no",
	"ter", " wa", "iv", "ho", "e a", " r", "hat", "s t", "ns", "ch ", "wh", "tr",
	"ut", "/", "have", "ly ", "ta", " ha", " on", "tha", "-", " l", "ati", "en ",
	"pe", " re", "there", "ass", "si", " fo", "wa", "ec", "our", "who", "its", "z",
	"fo", "rs", ">", "ot", "un", "<", "im", "th ", "nc", "ate", "><", "ver", "ad",
	" we", "ly", "ee", " n", "id", " cl", "ac", "il", "</", "rt", " wi", "div",
	"e, ", " it", "whi", " ma", "ge", "x", "e c", "men", ".com",
}

var codes = make([][]byte, len(codeStrings))
var codeTrie = trie.New()

func init() {
	for i, code := range codeStrings {
		codes[i] = []byte(code)
		codeTrie.Put([]byte(code), i)
	}
}

func flushVerb(outBuf, verbBuf *bytes.Buffer) {
	// We can write a max of 255 continuous verbatim characters, because the length of the continous verbatim
	// section is represented by a single byte.
	for verbBuf.Len() > 0 {
		chunk := verbBuf.Next(255)
		if len(chunk) == 1 {
			// 254 is code for a single verbatim byte
			outBuf.WriteByte(byte(254))
		} else {
			// 255 is code for a verbatim string. It is followed by a byte containing the length of the string.
			outBuf.WriteByte(byte(255))
			outBuf.WriteByte(byte(len(chunk)))
		}
		outBuf.Write(chunk)
	}
	verbBuf.Reset()
}

// Compress compresses a byte slice and returns the compressed data.
func Compress(input []byte) []byte {
	var outBuf bytes.Buffer
	var verbBuf bytes.Buffer
	root := codeTrie.Root()

	for len(input) > 0 {
		prefixLen := 0
		code := 0
		node := root
		for i, c := range input {
			next, ok := node.Walk(c)
			if !ok {
				break
			}
			node = next
			if node.Terminal() {
				prefixLen = i + 1
				code = node.Val()
			}
		}

		if prefixLen > 0 {
			input = input[prefixLen:]
			flushVerb(&outBuf, &verbBuf)
			outBuf.WriteByte(byte(code))
		} else {
			verbBuf.WriteByte(input[0])
			input = input[1:]
		}
	}
	flushVerb(&outBuf, &verbBuf)

	return outBuf.Bytes()
}

// DecompressionError is returned when decompressing invalid smaz-encoded data.
var DecompressionError = errors.New("Invalid or corrupted compressed data.")

// Decompress decompresses a smaz-compressed byte slice and return a new slice with the decompressed data. err
// is nil if and only if decompression fails for any reason (e.g., corrupted data).
func Decompress(compressed []byte) ([]byte, error) {
	decompressed := bytes.NewBuffer(make([]byte, 0, len(compressed))) // Estimate initial size

	for len(compressed) > 0 {
		switch compressed[0] {
		case 254: // Verbatim byte
			if len(compressed) < 2 {
				return nil, DecompressionError
			}
			decompressed.WriteByte(compressed[1])
			compressed = compressed[2:]
		case 255: // Verbatim string
			if len(compressed) < 2 {
				return nil, DecompressionError
			}
			n := int(compressed[1])
			if len(compressed) < n+2 {
				return nil, DecompressionError
			}
			decompressed.Write(compressed[2 : n+2])
			compressed = compressed[n+2:]
		default: // Look up encoded value
			decompressed.Write(codes[int(compressed[0])])
			compressed = compressed[1:]
		}
	}

	return decompressed.Bytes(), nil
}
<commit_msg>speed up smaz decompression 2x<commit_after>// Package smaz is an implementation of the smaz library (https://github.com/antirez/smaz) for compressing
// small strings.
package smaz

import (
	"bytes"
	"errors"

	"github.com/kjk/go-smaz/trie"
)

var codeStrings = []string{" ",
	"the", "e", "t", "a", "of", "o", "and", "i", "n", "s", "e ", "r", " th",
	" t", "in", "he", "th", "h", "he ", "to", "\r\n", "l", "s ", "d", " a", "an",
	"er", "c", " o", "d ", "on", " of", "re", "of ", "t ", ", ", "is", "u", "at",
	"   ", "n ", "or", "which", "f", "m", "as", "it", "that", "\n", "was", "en",
	"  ", " w", "es", " an", " i", "\r", "f ", "g", "p", "nd", " s", "nd ", "ed ",
	"w", "ed", "http://", "for", "te", "ing", "y ", "The", " c", "ti", "r ", "his",
	"st", " in", "ar", "nt", ",", " to", "y", "ng", " h", "with", "le", "al", "to ",
	"b", "ou", "be", "were", " b", "se", "o ", "ent", "ha", "ng ", "their", "\"",
	"hi", "from", " f", "in ", "de", "ion", "me", "v", ".", "ve", "all", "re ",
	"ri", "ro", "is ", "co", "f t", "are", "ea", ". ", "her", " m", "er ", " p",
	"es ", "by", "they", "di", "ra", "ic", "not", "s, ", "d t", "at ", "ce", "la",
	"h ", "ne", "as ", "tio", "on ", "n t", "io", "we", " a ", "om", ", a", "s o",
	"ur", "li", "ll", "ch", "had", "this", "e t", "g ", "e\r\n", " wh", "ere",
	" co", "e o", "a ", "us", " d", "ss", "\n\r\n", "\r\n\r", "=\"", " be", " e",
	"s a", "ma", "one", "t t", "or ", "but", "el", "so", "l ", "e s", "s,", "no",
	"ter", " wa", "iv", "ho", "e a", " r", "hat", "s t", "ns", "ch ", "wh", "tr",
	"ut", "/", "have", "ly ", "ta", " ha", " on", "tha", "-", " l", "ati", "en ",
	"pe", " re", "there", "ass", "si", " fo", "wa", "ec", "our", "who", "its", "z",
	"fo", "rs", ">", "ot", "un", "<", "im", "th ", "nc", "ate", "><", "ver", "ad",
	" we", "ly", "ee", " n", "id", " cl", "ac", "il", "</", "rt", " wi", "div",
	"e, ", " it", "whi", " ma", "ge", "x", "e c", "men", ".com",
}

var codes = make([][]byte, len(codeStrings))
var codeTrie = trie.New()

func init() {
	for i, code := range codeStrings {
		codes[i] = []byte(code)
		codeTrie.Put([]byte(code), i)
	}
}

func flushVerb(outBuf, verbBuf *bytes.Buffer) {
	// We can write a max of 255 continuous verbatim characters, because the length of the continous verbatim
	// section is represented by a single byte.
	for verbBuf.Len() > 0 {
		chunk := verbBuf.Next(255)
		if len(chunk) == 1 {
			// 254 is code for a single verbatim byte
			outBuf.WriteByte(byte(254))
		} else {
			// 255 is code for a verbatim string. It is followed by a byte containing the length of the string.
			outBuf.WriteByte(byte(255))
			outBuf.WriteByte(byte(len(chunk)))
		}
		outBuf.Write(chunk)
	}
	verbBuf.Reset()
}

// Compress compresses a byte slice and returns the compressed data.
func Compress(input []byte) []byte {
	var outBuf bytes.Buffer
	var verbBuf bytes.Buffer
	root := codeTrie.Root()

	for len(input) > 0 {
		prefixLen := 0
		code := 0
		node := root
		for i, c := range input {
			next, ok := node.Walk(c)
			if !ok {
				break
			}
			node = next
			if node.Terminal() {
				prefixLen = i + 1
				code = node.Val()
			}
		}

		if prefixLen > 0 {
			input = input[prefixLen:]
			flushVerb(&outBuf, &verbBuf)
			outBuf.WriteByte(byte(code))
		} else {
			verbBuf.WriteByte(input[0])
			input = input[1:]
		}
	}
	flushVerb(&outBuf, &verbBuf)

	return outBuf.Bytes()
}

// ErrDecompression is returned when decompressing invalid smaz-encoded data.
var ErrDecompression = errors.New("Invalid or corrupted compressed data.")

// Decompress decompresses a smaz-compressed byte slice and return a new slice with the decompressed data. err
// is nil if and only if decompression fails for any reason (e.g., corrupted data).
func Decompress(compressed []byte) ([]byte, error) {
	decompressed := make([]byte, 0, len(compressed)) // Estimate initial size
	for len(compressed) > 0 {
		n := int(compressed[0])
		switch n {
		case 254: // Verbatim byte
			if len(compressed) < 2 {
				return nil, ErrDecompression
			}
			decompressed = append(decompressed, compressed[1])
			compressed = compressed[2:]
		case 255: // Verbatim string
			if len(compressed) < 2 {
				return nil, ErrDecompression
			}
			n = int(compressed[1])
			if len(compressed) < n+2 {
				return nil, ErrDecompression
			}
			decompressed = append(decompressed, compressed[2:n+2]...)
			compressed = compressed[n+2:]
		default: // Look up encoded value
			d := codes[n]
			decompressed = append(decompressed, d...)
			compressed = compressed[1:]
		}
	}

	return decompressed, nil
}
<|endoftext|>


<commit_before>package smtp

import (
	"bytes"
	"io"
	"net"
	"strconv"
	"strings"
)

type Mail struct {
	From, To string
	Mail     string
}

type Handler func(*Mail)

const sizeLimit = 32 * 1024
const maxLineLength = sizeLimit

type conn struct {
	domain string

	conn   io.ReadWriteCloser
	reader *bufferedReader

	handler Handler

	state    state
	from, to string
}

func (c *conn) greeting() {
	c.conn.Write([]byte("220 " + c.domain + " keytree/smtp ready!\r\n"))
}

func (c *conn) ehlo() {
	c.conn.Write([]byte("250-" + c.domain + "\r\n250-PIPELINING\r\n250-8BITMIME\r\n250-SMTPUTF8\r\n250-CHUNKING\r\n250 SIZE " + strconv.Itoa(sizeLimit) + "\r\n"))
}

func (c *conn) helo() {
	c.conn.Write([]byte("250 " + c.domain + "\r\n"))
}

func (c *conn) syntaxError(message string) {
	c.conn.Write([]byte("500 " + message + "\r\n"))
}

func (c *conn) tooManyRecipients() {
	c.conn.Write([]byte("451 only one recipient per mail, please\r\n"))
}

func (c *conn) tooMuchMail() {
	c.conn.Write([]byte("552 too much data\r\n"))
}

func (c *conn) unexpectedCommand() {
	c.conn.Write([]byte("503 did not expect that command\r\n"))
}

func (c *conn) ok() {
	c.conn.Write([]byte("250 ok\r\n"))
}

func (c *conn) quitOk() {
	c.conn.Write([]byte("221 ok\r\n"))
}

func (c *conn) weDontVerify() {
	c.conn.Write([]byte("252 vrfy is so 90s\r\n"))
}

func (c *conn) startMail() {
	c.conn.Write([]byte("354 here we go\r\n"))
}

func (c *conn) readData() (string, bool) {
	c.startMail()

	var lines []string
	length := 0

	for {
		line, err := c.reader.ReadLine()
		if err != nil {
			return "", false
		}
		if line == "." {
			break
		}
		line = strings.TrimPrefix(line, ".")

		length += len(line) + 2
		if length > sizeLimit {
			c.tooMuchMail()
			return "", false
		}
		lines = append(lines, line)
	}

	lines = append(lines, "") // include final CRLF
	email := strings.Join(lines, "\r\n")
	c.ok()

	return email, true
}

func (c *conn) readNextBdat() (*bdatCmd, bool) {
	for {
		line, err := c.reader.ReadLine()
		if err != nil {
			return nil, false
		}
		cmd, err := parseCommand(line)
		if err != nil {
			c.syntaxError(err.Error())
			continue
		}
		switch cmd := cmd.(type) {
		case *bdatCmd:
			return cmd, true
		default:
			c.unexpectedCommand()
		}
	}
}

func (c *conn) readBdat(cmd *bdatCmd) (string, bool) {
	var data [][]byte
	length := 0

	for {
		length += cmd.length
		if length > sizeLimit {
			c.tooMuchMail()
			return "", false
		}

		slice := make([]byte, cmd.length)
		if _, err := io.ReadFull(c.reader, slice); err != nil {
			return "", false
		}
		data = append(data, slice)

		if cmd.last {
			break
		}

		var ok bool
		cmd, ok = c.readNextBdat()
		if !ok {
			return "", false
		}
	}

	c.ok()
	return string(bytes.Join(data, nil)), true
}

type state int

const (
	initial state = iota
	gotFrom
	gotTo
	gotData
)

func (c *conn) processCommand(cmd interface{}) bool {
	switch cmd := cmd.(type) {
	case *heloCmd:
		if c.state != initial {
			c.unexpectedCommand()
			return true
		}
		if cmd.isEhlo {
			c.ehlo()
		} else {
			c.helo()
		}
		return true

	case *mailFromCmd:
		if c.state != initial {
			c.unexpectedCommand()
			return true
		}
		c.state, c.from = gotFrom, cmd.from
		c.ok()
		return true

	case *rcptToCmd:
		if c.state != gotFrom {
			c.unexpectedCommand()
			return true
		}
		c.state, c.to = gotTo, cmd.to
		c.ok()
		return true

	case *bdatCmd:
		if c.state != gotTo {
			c.unexpectedCommand()
			return true
		} else {
			mail, ok := c.readBdat(cmd)
			if !ok {
				return false
			}
			c.handler(&Mail{From: c.from, To: c.to, Mail: mail})
			c.state, c.from, c.to = initial, "", ""
			return true
		}

	case *dataCmd:
		if c.state != gotTo {
			c.unexpectedCommand()
			return true
		} else {
			mail, ok := c.readData()
			if !ok {
				return false
			}
			c.handler(&Mail{From: c.from, To: c.to, Mail: mail})
			c.state, c.from, c.to = initial, "", ""
			return true
		}

	case *rsetCmd:
		c.state, c.from, c.to = initial, "", ""
		c.ok()
		return true

	case *noopCmd:
		c.ok()
		return true

	case *quitCmd:
		c.quitOk()
		return false

	case *vrfyCmd:
		c.weDontVerify()
		return true

	default:
		c.unexpectedCommand()
		return true
	}
}

func (c *conn) handle() {
	c.greeting()
	defer c.conn.Close()

	c.state = initial

	for {
		line, err := c.reader.ReadLine()
		if err != nil {
			break
		}
		cmd, err := parseCommand(line)
		if err != nil {
			c.syntaxError(err.Error())
			continue
		}

		if ok := c.processCommand(cmd); !ok {
			break
		}
	}
}

func ListenAndServe(domain string, address string, handler Handler) error {
	if address == "" {
		address = ":smtp"
	}

	l, err := net.Listen("tcp", address)
	if err != nil {
		return err
	}

	for {
		var c io.ReadWriteCloser
		c, err := l.Accept()
		if err != nil {
			return err
		}

		conn := &conn{
			domain:  domain,
			conn:    c,
			reader:  newBufferedReader(c, maxLineLength),
			handler: handler,
		}
		go conn.handle()
	}
}
<commit_msg>change API a little<commit_after>package smtp

import (
	"bytes"
	"io"
	"net"
	"strconv"
	"strings"
)

type Mail struct {
	From, To string
	Mail     string
}

type Handler func(*Mail)

const sizeLimit = 32 * 1024
const maxLineLength = sizeLimit

type conn struct {
	domain string

	conn   io.ReadWriteCloser
	reader *bufferedReader

	handler Handler

	state    state
	from, to string
}

func (c *conn) greeting() {
	c.conn.Write([]byte("220 " + c.domain + " jellevandenhooff/smtp ready!\r\n"))
}

func (c *conn) ehlo() {
	c.conn.Write([]byte("250-" + c.domain + "\r\n250-PIPELINING\r\n250-8BITMIME\r\n250-SMTPUTF8\r\n250-CHUNKING\r\n250 SIZE " + strconv.Itoa(sizeLimit) + "\r\n"))
}

func (c *conn) helo() {
	c.conn.Write([]byte("250 " + c.domain + "\r\n"))
}

func (c *conn) syntaxError(message string) {
	c.conn.Write([]byte("500 " + message + "\r\n"))
}

func (c *conn) tooManyRecipients() {
	c.conn.Write([]byte("451 only one recipient per mail, please\r\n"))
}

func (c *conn) tooMuchMail() {
	c.conn.Write([]byte("552 too much data\r\n"))
}

func (c *conn) unexpectedCommand() {
	c.conn.Write([]byte("503 did not expect that command\r\n"))
}

func (c *conn) ok() {
	c.conn.Write([]byte("250 ok\r\n"))
}

func (c *conn) quitOk() {
	c.conn.Write([]byte("221 ok\r\n"))
}

func (c *conn) weDontVerify() {
	c.conn.Write([]byte("252 vrfy is so 90s\r\n"))
}

func (c *conn) startMail() {
	c.conn.Write([]byte("354 here we go\r\n"))
}

func (c *conn) readData() (string, bool) {
	c.startMail()

	var lines []string
	length := 0

	for {
		line, err := c.reader.ReadLine()
		if err != nil {
			return "", false
		}
		if line == "." {
			break
		}
		line = strings.TrimPrefix(line, ".")

		length += len(line) + 2
		if length > sizeLimit {
			c.tooMuchMail()
			return "", false
		}
		lines = append(lines, line)
	}

	lines = append(lines, "") // include final CRLF
	email := strings.Join(lines, "\r\n")
	c.ok()

	return email, true
}

func (c *conn) readNextBdat() (*bdatCmd, bool) {
	for {
		line, err := c.reader.ReadLine()
		if err != nil {
			return nil, false
		}
		cmd, err := parseCommand(line)
		if err != nil {
			c.syntaxError(err.Error())
			continue
		}
		switch cmd := cmd.(type) {
		case *bdatCmd:
			return cmd, true
		default:
			c.unexpectedCommand()
		}
	}
}

func (c *conn) readBdat(cmd *bdatCmd) (string, bool) {
	var data [][]byte
	length := 0

	for {
		length += cmd.length
		if length > sizeLimit {
			c.tooMuchMail()
			return "", false
		}

		slice := make([]byte, cmd.length)
		if _, err := io.ReadFull(c.reader, slice); err != nil {
			return "", false
		}
		data = append(data, slice)

		if cmd.last {
			break
		}

		var ok bool
		cmd, ok = c.readNextBdat()
		if !ok {
			return "", false
		}
	}

	c.ok()
	return string(bytes.Join(data, nil)), true
}

type state int

const (
	initial state = iota
	gotFrom
	gotTo
	gotData
)

func (c *conn) processCommand(cmd interface{}) bool {
	switch cmd := cmd.(type) {
	case *heloCmd:
		if c.state != initial {
			c.unexpectedCommand()
			return true
		}
		if cmd.isEhlo {
			c.ehlo()
		} else {
			c.helo()
		}
		return true

	case *mailFromCmd:
		if c.state != initial {
			c.unexpectedCommand()
			return true
		}
		c.state, c.from = gotFrom, cmd.from
		c.ok()
		return true

	case *rcptToCmd:
		if c.state != gotFrom {
			c.unexpectedCommand()
			return true
		}
		c.state, c.to = gotTo, cmd.to
		c.ok()
		return true

	case *bdatCmd:
		if c.state != gotTo {
			c.unexpectedCommand()
			return true
		} else {
			mail, ok := c.readBdat(cmd)
			if !ok {
				return false
			}
			c.handler(&Mail{From: c.from, To: c.to, Mail: mail})
			c.state, c.from, c.to = initial, "", ""
			return true
		}

	case *dataCmd:
		if c.state != gotTo {
			c.unexpectedCommand()
			return true
		} else {
			mail, ok := c.readData()
			if !ok {
				return false
			}
			c.handler(&Mail{From: c.from, To: c.to, Mail: mail})
			c.state, c.from, c.to = initial, "", ""
			return true
		}

	case *rsetCmd:
		c.state, c.from, c.to = initial, "", ""
		c.ok()
		return true

	case *noopCmd:
		c.ok()
		return true

	case *quitCmd:
		c.quitOk()
		return false

	case *vrfyCmd:
		c.weDontVerify()
		return true

	default:
		c.unexpectedCommand()
		return true
	}
}

func (c *conn) handle() {
	c.greeting()
	defer c.conn.Close()

	c.state = initial

	for {
		line, err := c.reader.ReadLine()
		if err != nil {
			break
		}
		cmd, err := parseCommand(line)
		if err != nil {
			c.syntaxError(err.Error())
			continue
		}

		if ok := c.processCommand(cmd); !ok {
			break
		}
	}
}

func Serve(domain string, listener net.Listener, handler Handler) error {
	for {
		var c io.ReadWriteCloser
		c, err := listener.Accept()
		if err != nil {
			return err
		}

		conn := &conn{
			domain:  domain,
			conn:    c,
			reader:  newBufferedReader(c, maxLineLength),
			handler: handler,
		}
		go conn.handle()
	}
}
<|endoftext|>


<commit_before>package configgenerator

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/url"
	"path/filepath"

	ignconfig "github.com/coreos/ignition/config/v2_2"
	ignconfigtypes "github.com/coreos/ignition/config/v2_2/types"
	"github.com/openshift/installer/installer/pkg/config"
	"github.com/vincent-petithory/dataurl"
)

const (
	caPath = "generated/tls/root-ca.crt"
)

// GenerateIgnConfig generates Ignition configs for the workers and masters.
func (c *ConfigGenerator) GenerateIgnConfig(clusterDir string) error {
	var masters config.NodePool
	var workers config.NodePool
	for _, pool := range c.NodePools {
		switch pool.Name {
		case "master":
			masters = pool
		case "worker":
			workers = pool
		case "etcd": // FIXME: ignore these until openshift/release stops defining them
		default:
			return fmt.Errorf("unrecognized role: %s", pool.Name)
		}
	}

	ca, err := ioutil.ReadFile(filepath.Join(clusterDir, caPath))
	if err != nil {
		return err
	}

	workerCfg, err := parseIgnFile(workers.IgnitionFile)
	if err != nil {
		return fmt.Errorf("failed to parse Ignition config for workers: %v", err)
	}

	// XXX(crawford): The SSH key should only be added to the bootstrap
	//                node. After that, MCO should be responsible for
	//                distributing SSH keys.
	c.embedUserBlock(&workerCfg)
	c.appendCertificateAuthority(&workerCfg, ca)
	c.embedAppendBlock(&workerCfg, "worker", "")

	if err = ignCfgToFile(workerCfg, filepath.Join(clusterDir, config.IgnitionPathWorker)); err != nil {
		return err
	}

	masterCfg, err := parseIgnFile(masters.IgnitionFile)
	if err != nil {
		return fmt.Errorf("failed to parse Ignition config for masters: %v", err)
	}

	for i := 0; i < masters.Count; i++ {
		ignCfg := masterCfg

		// XXX(crawford): The SSH key should only be added to the bootstrap
		//                node. After that, MCO should be responsible for
		//                distributing SSH keys.
		c.embedUserBlock(&ignCfg)
		c.appendCertificateAuthority(&ignCfg, ca)
		c.embedAppendBlock(&ignCfg, "master", fmt.Sprintf("etcd_index=%d", i))

		if err = ignCfgToFile(ignCfg, filepath.Join(clusterDir, fmt.Sprintf(config.IgnitionPathMaster, i))); err != nil {
			return err
		}
	}

	return nil
}

func parseIgnFile(filePath string) (ignconfigtypes.Config, error) {
	if filePath == "" {
		return ignconfigtypes.Config{
			Ignition: ignconfigtypes.Ignition{
				Version: ignconfigtypes.MaxVersion.String(),
			},
		}, nil
	}

	data, err := ioutil.ReadFile(filePath)
	if err != nil {
		return ignconfigtypes.Config{}, err
	}

	cfg, rpt, _ := ignconfig.Parse(data)
	if len(rpt.Entries) > 0 {
		return ignconfigtypes.Config{}, fmt.Errorf("failed to parse ignition file %s: %s", filePath, rpt.String())
	}

	return cfg, nil
}

func (c *ConfigGenerator) embedAppendBlock(ignCfg *ignconfigtypes.Config, role string, query string) {
	appendBlock := ignconfigtypes.ConfigReference{
		Source:       c.getTNCURL(role, query),
		Verification: ignconfigtypes.Verification{Hash: nil},
	}
	ignCfg.Ignition.Config.Append = append(ignCfg.Ignition.Config.Append, appendBlock)
}

func (c *ConfigGenerator) appendCertificateAuthority(ignCfg *ignconfigtypes.Config, ca []byte) {
	ignCfg.Ignition.Security.TLS.CertificateAuthorities = append(ignCfg.Ignition.Security.TLS.CertificateAuthorities, ignconfigtypes.CaReference{
		Source: dataurl.EncodeBytes(ca),
	})
}

func (c *ConfigGenerator) embedUserBlock(ignCfg *ignconfigtypes.Config) {
	userBlock := ignconfigtypes.PasswdUser{
		Name: "core",
		SSHAuthorizedKeys: []ignconfigtypes.SSHAuthorizedKey{
			ignconfigtypes.SSHAuthorizedKey(c.SSHKey),
		},
	}

	ignCfg.Passwd.Users = append(ignCfg.Passwd.Users, userBlock)
}

func (c *ConfigGenerator) getTNCURL(role string, query string) string {
	var u string

	// cloud platforms put this behind a load balancer which remaps ports;
	// libvirt doesn't do that - use the tnc port directly
	port := 80
	if c.Platform == config.PlatformLibvirt {
		port = 49500
	}

	if role == "master" || role == "worker" {
		u = func() *url.URL {
			return &url.URL{
				Scheme:   "https",
				Host:     fmt.Sprintf("%s-tnc.%s:%d", c.Name, c.BaseDomain, port),
				Path:     fmt.Sprintf("/config/%s", role),
				RawQuery: query,
			}
		}().String()
	}
	return u
}

func ignCfgToFile(ignCfg ignconfigtypes.Config, filePath string) error {
	data, err := json.MarshalIndent(&ignCfg, "", "  ")
	if err != nil {
		return err
	}

	return ioutil.WriteFile(filePath, data, 0666)
}
<commit_msg>Revert "installer/pkg/config-generator/ignition: Ignore etcd node pools"<commit_after>package configgenerator

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/url"
	"path/filepath"

	ignconfig "github.com/coreos/ignition/config/v2_2"
	ignconfigtypes "github.com/coreos/ignition/config/v2_2/types"
	"github.com/openshift/installer/installer/pkg/config"
	"github.com/vincent-petithory/dataurl"
)

const (
	caPath = "generated/tls/root-ca.crt"
)

// GenerateIgnConfig generates Ignition configs for the workers and masters.
func (c *ConfigGenerator) GenerateIgnConfig(clusterDir string) error {
	var masters config.NodePool
	var workers config.NodePool
	for _, pool := range c.NodePools {
		switch pool.Name {
		case "master":
			masters = pool
		case "worker":
			workers = pool
		default:
			return fmt.Errorf("unrecognized role: %s", pool.Name)
		}
	}

	ca, err := ioutil.ReadFile(filepath.Join(clusterDir, caPath))
	if err != nil {
		return err
	}

	workerCfg, err := parseIgnFile(workers.IgnitionFile)
	if err != nil {
		return fmt.Errorf("failed to parse Ignition config for workers: %v", err)
	}

	// XXX(crawford): The SSH key should only be added to the bootstrap
	//                node. After that, MCO should be responsible for
	//                distributing SSH keys.
	c.embedUserBlock(&workerCfg)
	c.appendCertificateAuthority(&workerCfg, ca)
	c.embedAppendBlock(&workerCfg, "worker", "")

	if err = ignCfgToFile(workerCfg, filepath.Join(clusterDir, config.IgnitionPathWorker)); err != nil {
		return err
	}

	masterCfg, err := parseIgnFile(masters.IgnitionFile)
	if err != nil {
		return fmt.Errorf("failed to parse Ignition config for masters: %v", err)
	}

	for i := 0; i < masters.Count; i++ {
		ignCfg := masterCfg

		// XXX(crawford): The SSH key should only be added to the bootstrap
		//                node. After that, MCO should be responsible for
		//                distributing SSH keys.
		c.embedUserBlock(&ignCfg)
		c.appendCertificateAuthority(&ignCfg, ca)
		c.embedAppendBlock(&ignCfg, "master", fmt.Sprintf("etcd_index=%d", i))

		if err = ignCfgToFile(ignCfg, filepath.Join(clusterDir, fmt.Sprintf(config.IgnitionPathMaster, i))); err != nil {
			return err
		}
	}

	return nil
}

func parseIgnFile(filePath string) (ignconfigtypes.Config, error) {
	if filePath == "" {
		return ignconfigtypes.Config{
			Ignition: ignconfigtypes.Ignition{
				Version: ignconfigtypes.MaxVersion.String(),
			},
		}, nil
	}

	data, err := ioutil.ReadFile(filePath)
	if err != nil {
		return ignconfigtypes.Config{}, err
	}

	cfg, rpt, _ := ignconfig.Parse(data)
	if len(rpt.Entries) > 0 {
		return ignconfigtypes.Config{}, fmt.Errorf("failed to parse ignition file %s: %s", filePath, rpt.String())
	}

	return cfg, nil
}

func (c *ConfigGenerator) embedAppendBlock(ignCfg *ignconfigtypes.Config, role string, query string) {
	appendBlock := ignconfigtypes.ConfigReference{
		Source:       c.getTNCURL(role, query),
		Verification: ignconfigtypes.Verification{Hash: nil},
	}
	ignCfg.Ignition.Config.Append = append(ignCfg.Ignition.Config.Append, appendBlock)
}

func (c *ConfigGenerator) appendCertificateAuthority(ignCfg *ignconfigtypes.Config, ca []byte) {
	ignCfg.Ignition.Security.TLS.CertificateAuthorities = append(ignCfg.Ignition.Security.TLS.CertificateAuthorities, ignconfigtypes.CaReference{
		Source: dataurl.EncodeBytes(ca),
	})
}

func (c *ConfigGenerator) embedUserBlock(ignCfg *ignconfigtypes.Config) {
	userBlock := ignconfigtypes.PasswdUser{
		Name: "core",
		SSHAuthorizedKeys: []ignconfigtypes.SSHAuthorizedKey{
			ignconfigtypes.SSHAuthorizedKey(c.SSHKey),
		},
	}

	ignCfg.Passwd.Users = append(ignCfg.Passwd.Users, userBlock)
}

func (c *ConfigGenerator) getTNCURL(role string, query string) string {
	var u string

	// cloud platforms put this behind a load balancer which remaps ports;
	// libvirt doesn't do that - use the tnc port directly
	port := 80
	if c.Platform == config.PlatformLibvirt {
		port = 49500
	}

	if role == "master" || role == "worker" {
		u = func() *url.URL {
			return &url.URL{
				Scheme:   "https",
				Host:     fmt.Sprintf("%s-tnc.%s:%d", c.Name, c.BaseDomain, port),
				Path:     fmt.Sprintf("/config/%s", role),
				RawQuery: query,
			}
		}().String()
	}
	return u
}

func ignCfgToFile(ignCfg ignconfigtypes.Config, filePath string) error {
	data, err := json.MarshalIndent(&ignCfg, "", "  ")
	if err != nil {
		return err
	}

	return ioutil.WriteFile(filePath, data, 0666)
}
<|endoftext|>


<commit_before>package main

import (
	"os/exec"
	"strings"
	"testing"
	"time"
)

func TestRestartStoppedContainer(t *testing.T) {
	defer deleteAllContainers()

	runCmd := exec.Command(dockerBinary, "run", "-d", "busybox", "echo", "foobar")
	out, _, err := runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	cleanedContainerID := stripTrailingCharacters(out)

	runCmd = exec.Command(dockerBinary, "wait", cleanedContainerID)
	if out, _, err = runCommandWithOutput(runCmd); err != nil {
		t.Fatal(out, err)
	}

	runCmd = exec.Command(dockerBinary, "logs", cleanedContainerID)
	out, _, err = runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	if out != "foobar\n" {
		t.Errorf("container should've printed 'foobar'")
	}

	runCmd = exec.Command(dockerBinary, "restart", cleanedContainerID)
	if out, _, err = runCommandWithOutput(runCmd); err != nil {
		t.Fatal(out, err)
	}

	runCmd = exec.Command(dockerBinary, "logs", cleanedContainerID)
	out, _, err = runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	if out != "foobar\nfoobar\n" {
		t.Errorf("container should've printed 'foobar' twice")
	}

	logDone("restart - echo foobar for stopped container")
}

func TestRestartRunningContainer(t *testing.T) {
	defer deleteAllContainers()

	runCmd := exec.Command(dockerBinary, "run", "-d", "busybox", "sh", "-c", "echo foobar && sleep 30 && echo 'should not print this'")
	out, _, err := runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	cleanedContainerID := stripTrailingCharacters(out)

	time.Sleep(1 * time.Second)

	runCmd = exec.Command(dockerBinary, "logs", cleanedContainerID)
	out, _, err = runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	if out != "foobar\n" {
		t.Errorf("container should've printed 'foobar'")
	}

	runCmd = exec.Command(dockerBinary, "restart", "-t", "1", cleanedContainerID)
	if out, _, err = runCommandWithOutput(runCmd); err != nil {
		t.Fatal(out, err)
	}

	runCmd = exec.Command(dockerBinary, "logs", cleanedContainerID)
	out, _, err = runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	time.Sleep(1 * time.Second)

	if out != "foobar\nfoobar\n" {
		t.Errorf("container should've printed 'foobar' twice")
	}

	logDone("restart - echo foobar for running container")
}

// Test that restarting a container with a volume does not create a new volume on restart. Regression test for #819.
func TestRestartWithVolumes(t *testing.T) {
	defer deleteAllContainers()

	runCmd := exec.Command(dockerBinary, "run", "-d", "-v", "/test", "busybox", "top")
	out, _, err := runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	cleanedContainerID := stripTrailingCharacters(out)

	runCmd = exec.Command(dockerBinary, "inspect", "--format", "{{ len .Volumes }}", cleanedContainerID)
	out, _, err = runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	if out = strings.Trim(out, " \n\r"); out != "1" {
		t.Errorf("expect 1 volume received %s", out)
	}

	runCmd = exec.Command(dockerBinary, "inspect", "--format", "{{ .Volumes }}", cleanedContainerID)
	volumes, _, err := runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(volumes, err)
	}

	runCmd = exec.Command(dockerBinary, "restart", cleanedContainerID)
	if out, _, err = runCommandWithOutput(runCmd); err != nil {
		t.Fatal(out, err)
	}

	runCmd = exec.Command(dockerBinary, "inspect", "--format", "{{ len .Volumes }}", cleanedContainerID)
	out, _, err = runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	if out = strings.Trim(out, " \n\r"); out != "1" {
		t.Errorf("expect 1 volume after restart received %s", out)
	}

	runCmd = exec.Command(dockerBinary, "inspect", "--format", "{{ .Volumes }}", cleanedContainerID)
	volumesAfterRestart, _, err := runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(volumesAfterRestart, err)
	}

	if volumes != volumesAfterRestart {
		volumes = strings.Trim(volumes, " \n\r")
		volumesAfterRestart = strings.Trim(volumesAfterRestart, " \n\r")
		t.Errorf("expected volume path: %s Actual path: %s", volumes, volumesAfterRestart)
	}

	logDone("restart - does not create a new volume on restart")
}

func TestRestartPolicyNO(t *testing.T) {
	defer deleteAllContainers()

	cmd := exec.Command(dockerBinary, "run", "-d", "--restart=no", "busybox", "false")
	out, _, err := runCommandWithOutput(cmd)
	if err != nil {
		t.Fatal(err, out)
	}

	id := strings.TrimSpace(string(out))
	name, err := inspectField(id, "HostConfig.RestartPolicy.Name")
	if err != nil {
		t.Fatal(err, out)
	}
	if name != "no" {
		t.Fatalf("Container restart policy name is %s, expected %s", name, "no")
	}

	logDone("restart - recording restart policy name for --restart=no")
}

func TestRestartPolicyAlways(t *testing.T) {
	defer deleteAllContainers()

	cmd := exec.Command(dockerBinary, "run", "-d", "--restart=always", "busybox", "false")
	out, _, err := runCommandWithOutput(cmd)
	if err != nil {
		t.Fatal(err, out)
	}

	id := strings.TrimSpace(string(out))
	name, err := inspectField(id, "HostConfig.RestartPolicy.Name")
	if err != nil {
		t.Fatal(err, out)
	}
	if name != "always" {
		t.Fatalf("Container restart policy name is %s, expected %s", name, "always")
	}

	logDone("restart - recording restart policy name for --restart=always")
}

func TestRestartPolicyOnFailure(t *testing.T) {
	defer deleteAllContainers()

	cmd := exec.Command(dockerBinary, "run", "-d", "--restart=on-failure:1", "busybox", "false")
	out, _, err := runCommandWithOutput(cmd)
	if err != nil {
		t.Fatal(err, out)
	}

	id := strings.TrimSpace(string(out))
	name, err := inspectField(id, "HostConfig.RestartPolicy.Name")
	if err != nil {
		t.Fatal(err, out)
	}
	if name != "on-failure" {
		t.Fatalf("Container restart policy name is %s, expected %s", name, "on-failure")
	}

	logDone("restart - recording restart policy name for --restart=on-failure")
}
<commit_msg>add TestContainerRestartwithGoodContainer test case<commit_after>package main

import (
	"os/exec"
	"strings"
	"testing"
	"time"
)

func TestRestartStoppedContainer(t *testing.T) {
	defer deleteAllContainers()

	runCmd := exec.Command(dockerBinary, "run", "-d", "busybox", "echo", "foobar")
	out, _, err := runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	cleanedContainerID := stripTrailingCharacters(out)

	runCmd = exec.Command(dockerBinary, "wait", cleanedContainerID)
	if out, _, err = runCommandWithOutput(runCmd); err != nil {
		t.Fatal(out, err)
	}

	runCmd = exec.Command(dockerBinary, "logs", cleanedContainerID)
	out, _, err = runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	if out != "foobar\n" {
		t.Errorf("container should've printed 'foobar'")
	}

	runCmd = exec.Command(dockerBinary, "restart", cleanedContainerID)
	if out, _, err = runCommandWithOutput(runCmd); err != nil {
		t.Fatal(out, err)
	}

	runCmd = exec.Command(dockerBinary, "logs", cleanedContainerID)
	out, _, err = runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	if out != "foobar\nfoobar\n" {
		t.Errorf("container should've printed 'foobar' twice")
	}

	logDone("restart - echo foobar for stopped container")
}

func TestRestartRunningContainer(t *testing.T) {
	defer deleteAllContainers()

	runCmd := exec.Command(dockerBinary, "run", "-d", "busybox", "sh", "-c", "echo foobar && sleep 30 && echo 'should not print this'")
	out, _, err := runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	cleanedContainerID := stripTrailingCharacters(out)

	time.Sleep(1 * time.Second)

	runCmd = exec.Command(dockerBinary, "logs", cleanedContainerID)
	out, _, err = runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	if out != "foobar\n" {
		t.Errorf("container should've printed 'foobar'")
	}

	runCmd = exec.Command(dockerBinary, "restart", "-t", "1", cleanedContainerID)
	if out, _, err = runCommandWithOutput(runCmd); err != nil {
		t.Fatal(out, err)
	}

	runCmd = exec.Command(dockerBinary, "logs", cleanedContainerID)
	out, _, err = runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	time.Sleep(1 * time.Second)

	if out != "foobar\nfoobar\n" {
		t.Errorf("container should've printed 'foobar' twice")
	}

	logDone("restart - echo foobar for running container")
}

// Test that restarting a container with a volume does not create a new volume on restart. Regression test for #819.
func TestRestartWithVolumes(t *testing.T) {
	defer deleteAllContainers()

	runCmd := exec.Command(dockerBinary, "run", "-d", "-v", "/test", "busybox", "top")
	out, _, err := runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	cleanedContainerID := stripTrailingCharacters(out)

	runCmd = exec.Command(dockerBinary, "inspect", "--format", "{{ len .Volumes }}", cleanedContainerID)
	out, _, err = runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	if out = strings.Trim(out, " \n\r"); out != "1" {
		t.Errorf("expect 1 volume received %s", out)
	}

	runCmd = exec.Command(dockerBinary, "inspect", "--format", "{{ .Volumes }}", cleanedContainerID)
	volumes, _, err := runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(volumes, err)
	}

	runCmd = exec.Command(dockerBinary, "restart", cleanedContainerID)
	if out, _, err = runCommandWithOutput(runCmd); err != nil {
		t.Fatal(out, err)
	}

	runCmd = exec.Command(dockerBinary, "inspect", "--format", "{{ len .Volumes }}", cleanedContainerID)
	out, _, err = runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(out, err)
	}

	if out = strings.Trim(out, " \n\r"); out != "1" {
		t.Errorf("expect 1 volume after restart received %s", out)
	}

	runCmd = exec.Command(dockerBinary, "inspect", "--format", "{{ .Volumes }}", cleanedContainerID)
	volumesAfterRestart, _, err := runCommandWithOutput(runCmd)
	if err != nil {
		t.Fatal(volumesAfterRestart, err)
	}

	if volumes != volumesAfterRestart {
		volumes = strings.Trim(volumes, " \n\r")
		volumesAfterRestart = strings.Trim(volumesAfterRestart, " \n\r")
		t.Errorf("expected volume path: %s Actual path: %s", volumes, volumesAfterRestart)
	}

	logDone("restart - does not create a new volume on restart")
}

func TestRestartPolicyNO(t *testing.T) {
	defer deleteAllContainers()

	cmd := exec.Command(dockerBinary, "run", "-d", "--restart=no", "busybox", "false")
	out, _, err := runCommandWithOutput(cmd)
	if err != nil {
		t.Fatal(err, out)
	}

	id := strings.TrimSpace(string(out))
	name, err := inspectField(id, "HostConfig.RestartPolicy.Name")
	if err != nil {
		t.Fatal(err, out)
	}
	if name != "no" {
		t.Fatalf("Container restart policy name is %s, expected %s", name, "no")
	}

	logDone("restart - recording restart policy name for --restart=no")
}

func TestRestartPolicyAlways(t *testing.T) {
	defer deleteAllContainers()

	cmd := exec.Command(dockerBinary, "run", "-d", "--restart=always", "busybox", "false")
	out, _, err := runCommandWithOutput(cmd)
	if err != nil {
		t.Fatal(err, out)
	}

	id := strings.TrimSpace(string(out))
	name, err := inspectField(id, "HostConfig.RestartPolicy.Name")
	if err != nil {
		t.Fatal(err, out)
	}
	if name != "always" {
		t.Fatalf("Container restart policy name is %s, expected %s", name, "always")
	}

	logDone("restart - recording restart policy name for --restart=always")
}

func TestRestartPolicyOnFailure(t *testing.T) {
	defer deleteAllContainers()

	cmd := exec.Command(dockerBinary, "run", "-d", "--restart=on-failure:1", "busybox", "false")
	out, _, err := runCommandWithOutput(cmd)
	if err != nil {
		t.Fatal(err, out)
	}

	id := strings.TrimSpace(string(out))
	name, err := inspectField(id, "HostConfig.RestartPolicy.Name")
	if err != nil {
		t.Fatal(err, out)
	}
	if name != "on-failure" {
		t.Fatalf("Container restart policy name is %s, expected %s", name, "on-failure")
	}

	logDone("restart - recording restart policy name for --restart=on-failure")
}

// a good container with --restart=on-failure:3
// MaximumRetryCount!=0; RestartCount=0
func TestContainerRestartwithGoodContainer(t *testing.T) {
	defer deleteAllContainers()
	out, err := exec.Command(dockerBinary, "run", "-d", "--restart=on-failure:3", "busybox", "true").CombinedOutput()
	if err != nil {
		t.Fatal(string(out), err)
	}
	id := strings.TrimSpace(string(out))
	if err := waitInspect(id, "{{ .State.Restarting }} {{ .State.Running }}", "false false", 5); err != nil {
		t.Fatal(err)
	}
	count, err := inspectField(id, "RestartCount")
	if err != nil {
		t.Fatal(err)
	}
	if count != "0" {
		t.Fatalf("Container was restarted %s times, expected %d", count, 0)
	}
	MaximumRetryCount, err := inspectField(id, "HostConfig.RestartPolicy.MaximumRetryCount")
	if err != nil {
		t.Fatal(err)
	}
	if MaximumRetryCount != "3" {
		t.Fatalf("Container Maximum Retry Count is %s, expected %s", MaximumRetryCount, "3")
	}

	logDone("restart - for a good container with restart policy, MaximumRetryCount is not 0 and RestartCount is 0")
}
<|endoftext|>


<commit_before>package stager_runner

import (
	"os/exec"
	"strings"
	"time"

	"github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
)

type StagerRunner struct {
	stagerBin     string
	etcdCluster   []string
	natsAddresses []string

	session     *gexec.Session
	CompilerUrl string
}

func New(stagerBin string, etcdCluster []string, natsAddresses []string) *StagerRunner {
	return &StagerRunner{
		stagerBin:     stagerBin,
		etcdCluster:   etcdCluster,
		natsAddresses: natsAddresses,
	}
}

func (r *StagerRunner) Start(args ...string) {
	stagerSession, err := gexec.Start(
		exec.Command(
			r.stagerBin,
			append([]string{
				"-etcdCluster", strings.Join(r.etcdCluster, ","),
				"-natsAddresses", strings.Join(r.natsAddresses, ","),
			}, args...)...,
		),
		ginkgo.GinkgoWriter,
		ginkgo.GinkgoWriter,
	)

	Ω(err).ShouldNot(HaveOccurred())
	Eventually(stagerSession).Should(gbytes.Say("Listening for staging requests!"))

	r.session = stagerSession
}

func (r *StagerRunner) Stop() {
	if r.session != nil {
		r.session.Interrupt().Wait(5 * time.Second)
	}
}

func (r *StagerRunner) Session() *gexec.Session {
	return r.session
}
<commit_msg>Add KillWithFire to runner<commit_after>package stager_runner

import (
	"os/exec"
	"strings"
	"time"

	"github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
)

type StagerRunner struct {
	stagerBin     string
	etcdCluster   []string
	natsAddresses []string

	session     *gexec.Session
	CompilerUrl string
}

func New(stagerBin string, etcdCluster []string, natsAddresses []string) *StagerRunner {
	return &StagerRunner{
		stagerBin:     stagerBin,
		etcdCluster:   etcdCluster,
		natsAddresses: natsAddresses,
	}
}

func (r *StagerRunner) Start(args ...string) {
	if r.session != nil {
		panic("starting more than one stager runner!!!")
	}

	stagerSession, err := gexec.Start(
		exec.Command(
			r.stagerBin,
			append([]string{
				"-etcdCluster", strings.Join(r.etcdCluster, ","),
				"-natsAddresses", strings.Join(r.natsAddresses, ","),
			}, args...)...,
		),
		ginkgo.GinkgoWriter,
		ginkgo.GinkgoWriter,
	)

	Ω(err).ShouldNot(HaveOccurred())
	Eventually(stagerSession).Should(gbytes.Say("Listening for staging requests!"))

	r.session = stagerSession
}

func (r *StagerRunner) Stop() {
	if r.session != nil {
		r.session.Interrupt().Wait(5 * time.Second)
		r.session = nil
	}
}

func (r *StagerRunner) KillWithFire() {
	if r.session != nil {
		r.session.Kill().Wait(5 * time.Second)
		r.session = nil
	}
}

func (r *StagerRunner) Session() *gexec.Session {
	return r.session
}
<|endoftext|>


<commit_before>package slice

import "github.com/cheekybits/genny/generic"

type T generic.Type
type T1 generic.Type
type T2 generic.Type

type Slice_T []T
type Slice_T1 []T1
type Slice_T2 []T2

func (s Slice_T) Map(f func(e T) T) Slice_T {
	res := []T{}
	for _, e := range s {
		res = append(res, f(e))
	}
	return res
}

func (s Slice_T1) Map2(f func(e T1) T2) Slice_T2 {
	res := []T2{}
	for _, e := range s {
		res = append(res, f(e))
	}
	return res
}
<commit_msg>better version of 2 typed Map<commit_after>package slice

import "github.com/cheekybits/genny/generic"

type T generic.Type
type T1 generic.Type

type Slice_T []T
type Slice_T1 []T1

func (s Slice_T) Map(f func(e T) T) Slice_T {
	res := []T{}
	for _, e := range s {
		res = append(res, f(e))
	}
	return res
}

func (s Slice_T) Map_T1(f func(e T) T1) Slice_T1 {
	res := []T1{}
	for _, e := range s {
		res = append(res, f(e))
	}
	return res
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"encoding/json"
	"go-imap"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"time"
	"fmt"
)

type ServerConfig struct {
	Debug        bool   `json:"debug"`
	Hostname     string `json:"hostname"`
	Port         string `json:"port"`
	UseTLS       bool   `json:"useTLS"`
	CertFilename string `json:"certFilename"`
	KeyFilename  string `json:"keyFilename"`
	IDLETimeout  int    `json:"IDLETimeoutMinutes"`
}

var gServerConfig ServerConfig

func notifyNewMessageHandler(request *registrationRequest) {
	body := strings.NewReader("version=" + string(int32(time.Now().Unix())))
	r, err := http.NewRequest("PUT", request.OnNewMessageURL, body)
	if err != nil {
		log.Println(err)
		return
	}

	var client *http.Client

	if gServerConfig.UseTLS {
		client = &http.Client{}
	} else {
		config := &tls.Config{InsecureSkipVerify: true} // this line here
		tr := &http.Transport{TLSClientConfig: config}
		client = &http.Client{Transport: tr}
	}

	_, err = client.Do(r)
	if err != nil {
		log.Println(err)
		return
	}
}

func readConfig() {

	var data []byte
	var err error

	data, err = ioutil.ReadFile("config.json")
	if err != nil {
		log.Println("Not configured.  Could not find config.json")
		os.Exit(-1)
	}

	err = json.Unmarshal(data, &gServerConfig)
	if err != nil {
		log.Println("Could not unmarshal config.json", err)
		os.Exit(-1)
		return
	}
}

type registrationRequest struct {
	Username        string `json:"username"`
	Password        string `json:"password"`
	OnNewMessageURL string `json:"onNewMessageURL"`
	OnReconnectURL  string `json:"onReconnectURL"`
}

func registerHandler(w http.ResponseWriter, r *http.Request) {
	log.Println("Got notification from app server ", r.URL)
	if r.Method != "POST" {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("Not a post request"))
		return
	}
	// receive posted data
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("Could not access body of request"))
		return
	}

	request := new(registrationRequest)
	err = json.Unmarshal(body, request)

	if err != nil {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("Could not unmarshal body, "))
		log.Println(err)
		return
	}

	log.Println("got request: ", request)
	log.Println("got username: ", request.Username)
	log.Println("got password: ", request.Password)
	log.Println("got onNewMessageURL: ", request.OnNewMessageURL)
	log.Println("got onReconnectURL: ", request.OnReconnectURL)

	conn, err := tls.Dial("tcp", "imap.gmail.com:993", nil)
	var reader io.Reader = conn
	im := imap.New(reader, conn)
	im.Unsolicited = make(chan interface{}, 100)

	hello, err := im.Start()
	if err != nil {
		log.Println("imap didn't start...")
		return
	}
	log.Println("server hello: %s", hello)

	log.Println("logging in...")

	resp, caps, err := im.Auth(request.Username, request.Password)
	if err != nil {
		log.Println("server failed to auth: %s", err)
		return
	}

	log.Println("%s", resp)
	log.Println("server capabilities: %s", caps)

	respExamine, err := im.Examine("inbox")
	log.Println("EXISTS: ", respExamine.Exists)

	log.Println("Beginnig to IDLE")
	idleChan, err := im.Idle()

	if err != nil {
		log.Println("failed to send IDLE command")
	}

	go func() {
		for {
			select {
			case message := <-idleChan:
				switch message := message.(type) {
				case *imap.ResponseExists:
					log.Println("Got EXISTS ", message.Count)
					notifyNewMessageHandler(request)
				case *imap.ResponseStatus:
					if message.Status != imap.OK {
						panic(fmt.Sprintf("Non-OK response from IDLE: %+v", message))
					}

					log.Println("Restarting IDLE")
					idleChan, err = im.Idle()
				}
			case <-time.After(time.Duration(gServerConfig.IDLETimeout) * time.Minute):
				/* RFC 2177:
				 * (...) clients using IDLE are advised to terminate the IDLE and
				 * re-issue it at least every 29 minutes to avoid being logged
				 * off.
				 */
				 log.Println("Sending DONE")
				 im.Done()
			 }
		}
	}()

	w.WriteHeader(http.StatusOK)
	w.Write([]byte("OK\n"))
}

func main() {

	readConfig()

	http.HandleFunc("/register", registerHandler)

	log.Println("Listening on", gServerConfig.Hostname+":"+gServerConfig.Port)

	var err error
	if gServerConfig.UseTLS {
		err = http.ListenAndServeTLS(gServerConfig.Hostname+":"+gServerConfig.Port,
			gServerConfig.CertFilename,
			gServerConfig.KeyFilename,
			nil)
	} else {
		for i := 0; i < 5; i++ {
			log.Println("This is a really unsafe way to run the push server.  Really.  Don't do this in production.")
		}
		err = http.ListenAndServe(gServerConfig.Hostname+":"+gServerConfig.Port, nil)
	}

	log.Println("Exiting... ", err)
}
<commit_msg>Run gofmt.<commit_after>package main

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"go-imap"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"time"
)

type ServerConfig struct {
	Debug        bool   `json:"debug"`
	Hostname     string `json:"hostname"`
	Port         string `json:"port"`
	UseTLS       bool   `json:"useTLS"`
	CertFilename string `json:"certFilename"`
	KeyFilename  string `json:"keyFilename"`
	IDLETimeout  int    `json:"IDLETimeoutMinutes"`
}

var gServerConfig ServerConfig

func notifyNewMessageHandler(request *registrationRequest) {
	body := strings.NewReader("version=" + string(int32(time.Now().Unix())))
	r, err := http.NewRequest("PUT", request.OnNewMessageURL, body)
	if err != nil {
		log.Println(err)
		return
	}

	var client *http.Client

	if gServerConfig.UseTLS {
		client = &http.Client{}
	} else {
		config := &tls.Config{InsecureSkipVerify: true} // this line here
		tr := &http.Transport{TLSClientConfig: config}
		client = &http.Client{Transport: tr}
	}

	_, err = client.Do(r)
	if err != nil {
		log.Println(err)
		return
	}
}

func readConfig() {

	var data []byte
	var err error

	data, err = ioutil.ReadFile("config.json")
	if err != nil {
		log.Println("Not configured.  Could not find config.json")
		os.Exit(-1)
	}

	err = json.Unmarshal(data, &gServerConfig)
	if err != nil {
		log.Println("Could not unmarshal config.json", err)
		os.Exit(-1)
		return
	}
}

type registrationRequest struct {
	Username        string `json:"username"`
	Password        string `json:"password"`
	OnNewMessageURL string `json:"onNewMessageURL"`
	OnReconnectURL  string `json:"onReconnectURL"`
}

func registerHandler(w http.ResponseWriter, r *http.Request) {
	log.Println("Got notification from app server ", r.URL)
	if r.Method != "POST" {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("Not a post request"))
		return
	}
	// receive posted data
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("Could not access body of request"))
		return
	}

	request := new(registrationRequest)
	err = json.Unmarshal(body, request)

	if err != nil {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("Could not unmarshal body, "))
		log.Println(err)
		return
	}

	log.Println("got request: ", request)
	log.Println("got username: ", request.Username)
	log.Println("got password: ", request.Password)
	log.Println("got onNewMessageURL: ", request.OnNewMessageURL)
	log.Println("got onReconnectURL: ", request.OnReconnectURL)

	conn, err := tls.Dial("tcp", "imap.gmail.com:993", nil)
	var reader io.Reader = conn
	im := imap.New(reader, conn)
	im.Unsolicited = make(chan interface{}, 100)

	hello, err := im.Start()
	if err != nil {
		log.Println("imap didn't start...")
		return
	}
	log.Println("server hello: %s", hello)

	log.Println("logging in...")

	resp, caps, err := im.Auth(request.Username, request.Password)
	if err != nil {
		log.Println("server failed to auth: %s", err)
		return
	}

	log.Println("%s", resp)
	log.Println("server capabilities: %s", caps)

	respExamine, err := im.Examine("inbox")
	log.Println("EXISTS: ", respExamine.Exists)

	log.Println("Beginnig to IDLE")
	idleChan, err := im.Idle()

	if err != nil {
		log.Println("failed to send IDLE command")
	}

	go func() {
		for {
			select {
			case message := <-idleChan:
				switch message := message.(type) {
				case *imap.ResponseExists:
					log.Println("Got EXISTS ", message.Count)
					notifyNewMessageHandler(request)
				case *imap.ResponseStatus:
					if message.Status != imap.OK {
						panic(fmt.Sprintf("Non-OK response from IDLE: %+v", message))
					}

					log.Println("Restarting IDLE")
					idleChan, err = im.Idle()
				}
			case <-time.After(time.Duration(gServerConfig.IDLETimeout) * time.Minute):
				/* RFC 2177:
				 * (...) clients using IDLE are advised to terminate the IDLE and
				 * re-issue it at least every 29 minutes to avoid being logged
				 * off.
				 */
				log.Println("Sending DONE")
				im.Done()
			}
		}
	}()

	w.WriteHeader(http.StatusOK)
	w.Write([]byte("OK\n"))
}

func main() {

	readConfig()

	http.HandleFunc("/register", registerHandler)

	log.Println("Listening on", gServerConfig.Hostname+":"+gServerConfig.Port)

	var err error
	if gServerConfig.UseTLS {
		err = http.ListenAndServeTLS(gServerConfig.Hostname+":"+gServerConfig.Port,
			gServerConfig.CertFilename,
			gServerConfig.KeyFilename,
			nil)
	} else {
		for i := 0; i < 5; i++ {
			log.Println("This is a really unsafe way to run the push server.  Really.  Don't do this in production.")
		}
		err = http.ListenAndServe(gServerConfig.Hostname+":"+gServerConfig.Port, nil)
	}

	log.Println("Exiting... ", err)
}
<|endoftext|>


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

package main

import (
	"time"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/rand/mathrand"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/gcloud/gs"
	gcps "go.chromium.org/luci/common/gcloud/pubsub"
	log "go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/common/tsmon/distribution"
	"go.chromium.org/luci/common/tsmon/field"
	"go.chromium.org/luci/common/tsmon/metric"
	montypes "go.chromium.org/luci/common/tsmon/types"
	"go.chromium.org/luci/grpc/grpcutil"
	"go.chromium.org/luci/logdog/api/config/svcconfig"
	"go.chromium.org/luci/logdog/common/types"
	"go.chromium.org/luci/logdog/server/archivist"
	"go.chromium.org/luci/logdog/server/bundleServicesClient"
	"go.chromium.org/luci/logdog/server/service"

	"cloud.google.com/go/pubsub"
	"golang.org/x/net/context"

	"go.chromium.org/luci/hardcoded/chromeinfra"
)

var (
	errInvalidConfig = errors.New("invalid configuration")

	// tsTaskProcessingTime measures the amount of time spent processing a single
	// task.
	//
	// The "consumed" field is true if the underlying task was consumed and
	// false if it was not.
	tsTaskProcessingTime = metric.NewCumulativeDistribution("logdog/archivist/task_processing_time_ms",
		"The amount of time (in milliseconds) that a single task takes to process.",
		&montypes.MetricMetadata{Units: montypes.Milliseconds},
		distribution.DefaultBucketer,
		field.Bool("consumed"))
)

// application is the Archivist application state.
type application struct {
	service.Service
}

// run is the main execution function.
func (a *application) runArchivist(c context.Context) error {
	cfg := a.ServiceConfig()

	coordCfg, acfg := cfg.GetCoordinator(), cfg.GetArchivist()
	switch {
	case coordCfg == nil:
		fallthrough

	case acfg == nil:
		return errors.New("missing required config: archivist")
	case acfg.GsStagingBucket == "":
		return errors.New("missing required config: archivist.gs_staging_bucket")
	}

	// Initialize Pub/Sub client.
	//
	// We will initialize both an authenticated Client instance and an
	// authenticated Context, since we need the latter for raw ACK deadline
	// updates.
	taskSub := gcps.Subscription(acfg.Subscription)
	if err := taskSub.Validate(); err != nil {
		log.Fields{
			log.ErrorKey: err,
			"value":      taskSub,
		}.Errorf(c, "Task subscription did not validate.")
		return errors.New("invalid task subscription name")
	}
	psProject, psSubscriptionName := taskSub.Split()

	// New PubSub instance with the authenticated client.
	psClient, err := a.Service.PubSubSubscriberClient(c, psProject)
	if err != nil {
		log.WithError(err).Errorf(c, "Failed to create Pub/Sub client.")
		return err
	}
	sub := psClient.Subscription(psSubscriptionName)
	sub.ReceiveSettings = pubsub.ReceiveSettings{
		MaxExtension:           24 * time.Hour,
		MaxOutstandingMessages: int(acfg.Tasks), // If < 1, default.
		MaxOutstandingBytes:    0,               // Default.
	}

	// Initialize our Storage.
	//
	// NOTE: We're requesting read/write access even though we only need read-only
	// access because GKE doesn't understand the read-only scope:
	// https://www.googleapis.com/auth/bigtable.readonly
	st, err := a.IntermediateStorage(c, true)
	if err != nil {
		log.WithError(err).Errorf(c, "Failed to get storage instance.")
		return err
	}
	defer st.Close()

	// Initialize our Google Storage client.
	gsClient, err := a.GSClient(c)
	if err != nil {
		log.WithError(err).Errorf(c, "Failed to get Google Storage client.")
		return err
	}
	defer gsClient.Close()

	// Initialize a Coordinator client that bundles requests together.
	coordClient := &bundleServicesClient.Client{
		ServicesClient:       a.Coordinator(),
		DelayThreshold:       time.Second,
		BundleCountThreshold: 100,
	}
	defer coordClient.Flush()

	ar := archivist.Archivist{
		Service:        coordClient,
		SettingsLoader: a.GetSettingsLoader(acfg),
		Storage:        st,
		GSClient:       gsClient,
	}

	// Application shutdown will now operate by stopping the Iterator.
	c, cancelFunc := context.WithCancel(c)
	defer cancelFunc()

	// Application shutdown will now operate by cancelling the Collector's
	// shutdown Context.
	a.SetShutdownFunc(cancelFunc)

	// Execute our main subscription pull loop. It will run until the supplied
	// Context is cancelled.
	log.Fields{
		"subscription": taskSub,
	}.Infof(c, "Pulling tasks from Pub/Sub subscription.")

	retryForever := func() retry.Iterator {
		return &retry.ExponentialBackoff{
			Limited: retry.Limited{
				Delay:   200 * time.Millisecond,
				Retries: -1, // Unlimited.
			},
			MaxDelay:   10 * time.Second,
			Multiplier: 2,
		}
	}

	err = retry.Retry(c, transient.Only(retryForever), func() error {
		return grpcutil.WrapIfTransient(sub.Receive(c, func(c context.Context, msg *pubsub.Message) {
			c = log.SetFields(c, log.Fields{
				"messageID": msg.ID,
			})

			// ACK (or not) the message based on whether our task was consumed.
			deleteTask := false
			defer func() {
				// ACK the message if it is completed. If not, we do not NACK it, as we
				// want to let Pub/Sub redelivery delay occur as a form of backoff.
				if deleteTask {
					msg.Ack()
				}
			}()

			// Time how long task processing takes for metrics.
			startTime := clock.Now(c)
			defer func() {
				duration := clock.Now(c).Sub(startTime)

				if deleteTask {
					log.Fields{
						"duration": duration,
					}.Infof(c, "Task successfully processed; deleting.")
				} else {
					log.Fields{
						"duration": duration,
					}.Infof(c, "Task processing incomplete. Not deleting.")
				}

				// Add to our processing time metric.
				tsTaskProcessingTime.Add(c, duration.Seconds()*1000, deleteTask)
			}()

			task, err := makePubSubArchivistTask(c, psSubscriptionName, msg)
			c = log.SetFields(c, log.Fields{
				"consumed":         task.consumed,
				"subscriptionName": task.subscriptionName,
				"taskTimestamp":    task.timestamp.Format(time.RFC3339Nano),
				"archiveTask":      task.at,
			})
			if task.msg != nil {
				// Log all fields except data.
				c = log.SetFields(c, log.Fields{
					"message": map[string]interface{}{
						"id":          task.msg.ID,
						"attributes":  task.msg.Attributes,
						"publishTime": task.msg.PublishTime.Format(time.RFC3339Nano),
					},
				})
			}
			if err != nil {
				log.WithError(err).Errorf(c, "Failed to unmarshal archive task from message.")
				deleteTask = true
				return
			}

			ar.ArchiveTask(c, task)
			deleteTask = task.consumed
		}))
	}, func(err error, d time.Duration) {
		log.Fields{
			log.ErrorKey: err,
			"delay":      d,
		}.Warningf(c, "Transient error during subscription Receive loop; retrying...")
	})

	if err := errors.Unwrap(err); err != nil && err != context.Canceled {
		log.WithError(err).Errorf(c, "Failed during Pub/Sub Receive.")
		return err
	}

	log.Debugf(c, "Archivist finished.")
	return nil
}

// GetSettingsLoader is an archivist.SettingsLoader implementation that merges
// global and project-specific settings.
//
// The resulting settings object will be verified by the Archivist.
func (a *application) GetSettingsLoader(acfg *svcconfig.Archivist) archivist.SettingsLoader {
	serviceID := a.ServiceID()

	return func(c context.Context, proj types.ProjectName) (*archivist.Settings, error) {
		// Fold in our project-specific configuration, if valid.
		pcfg, err := a.ProjectConfig(c, proj)
		if err != nil {
			log.Fields{
				log.ErrorKey: err,
				"project":    proj,
			}.Errorf(c, "Failed to fetch project configuration.")
			return nil, err
		}

		indexParam := func(get func(ic *svcconfig.ArchiveIndexConfig) int32) int {
			if ic := pcfg.ArchiveIndexConfig; ic != nil {
				if v := get(ic); v > 0 {
					return int(v)
				}
			}

			if ic := acfg.ArchiveIndexConfig; ic != nil {
				if v := get(ic); v > 0 {
					return int(v)
				}
			}

			return 0
		}

		// Load our base settings.
		//
		// Archival bases are:
		// Staging: gs://<services:gs_staging_bucket>/<project-id>/...
		// Archive: gs://<project:archive_gs_bucket>/<project-id>/...
		st := archivist.Settings{
			GSBase:        gs.MakePath(pcfg.ArchiveGsBucket, "").Concat(serviceID),
			GSStagingBase: gs.MakePath(acfg.GsStagingBucket, "").Concat(serviceID),

			IndexStreamRange: indexParam(func(ic *svcconfig.ArchiveIndexConfig) int32 { return ic.StreamRange }),
			IndexPrefixRange: indexParam(func(ic *svcconfig.ArchiveIndexConfig) int32 { return ic.PrefixRange }),
			IndexByteRange:   indexParam(func(ic *svcconfig.ArchiveIndexConfig) int32 { return ic.ByteRange }),
			AlwaysRender:     (acfg.RenderAllStreams || pcfg.RenderAllStreams),
		}

		// Fold project settings into loaded ones.
		return &st, nil
	}
}

// Entry point.
func main() {
	mathrand.SeedRandomly()
	a := application{
		Service: service.Service{
			Name:               "archivist",
			DefaultAuthOptions: chromeinfra.DefaultAuthOptions(),
		},
	}
	a.Run(context.Background(), a.runArchivist)
}
<commit_msg>[logdog] NACK messages if they are transient failures<commit_after>// Copyright 2016 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"time"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/rand/mathrand"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/gcloud/gs"
	gcps "go.chromium.org/luci/common/gcloud/pubsub"
	log "go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/common/tsmon/distribution"
	"go.chromium.org/luci/common/tsmon/field"
	"go.chromium.org/luci/common/tsmon/metric"
	montypes "go.chromium.org/luci/common/tsmon/types"
	"go.chromium.org/luci/grpc/grpcutil"
	"go.chromium.org/luci/logdog/api/config/svcconfig"
	"go.chromium.org/luci/logdog/common/types"
	"go.chromium.org/luci/logdog/server/archivist"
	"go.chromium.org/luci/logdog/server/bundleServicesClient"
	"go.chromium.org/luci/logdog/server/service"

	"cloud.google.com/go/pubsub"
	"golang.org/x/net/context"

	"go.chromium.org/luci/hardcoded/chromeinfra"
)

var (
	errInvalidConfig = errors.New("invalid configuration")

	// tsTaskProcessingTime measures the amount of time spent processing a single
	// task.
	//
	// The "consumed" field is true if the underlying task was consumed and
	// false if it was not.
	tsTaskProcessingTime = metric.NewCumulativeDistribution("logdog/archivist/task_processing_time_ms",
		"The amount of time (in milliseconds) that a single task takes to process.",
		&montypes.MetricMetadata{Units: montypes.Milliseconds},
		distribution.DefaultBucketer,
		field.Bool("consumed"))
)

// application is the Archivist application state.
type application struct {
	service.Service
}

// run is the main execution function.
func (a *application) runArchivist(c context.Context) error {
	cfg := a.ServiceConfig()

	coordCfg, acfg := cfg.GetCoordinator(), cfg.GetArchivist()
	switch {
	case coordCfg == nil:
		fallthrough

	case acfg == nil:
		return errors.New("missing required config: archivist")
	case acfg.GsStagingBucket == "":
		return errors.New("missing required config: archivist.gs_staging_bucket")
	}

	// Initialize Pub/Sub client.
	//
	// We will initialize both an authenticated Client instance and an
	// authenticated Context, since we need the latter for raw ACK deadline
	// updates.
	taskSub := gcps.Subscription(acfg.Subscription)
	if err := taskSub.Validate(); err != nil {
		log.Fields{
			log.ErrorKey: err,
			"value":      taskSub,
		}.Errorf(c, "Task subscription did not validate.")
		return errors.New("invalid task subscription name")
	}
	psProject, psSubscriptionName := taskSub.Split()

	// New PubSub instance with the authenticated client.
	psClient, err := a.Service.PubSubSubscriberClient(c, psProject)
	if err != nil {
		log.WithError(err).Errorf(c, "Failed to create Pub/Sub client.")
		return err
	}
	sub := psClient.Subscription(psSubscriptionName)
	sub.ReceiveSettings = pubsub.ReceiveSettings{
		MaxExtension:           24 * time.Hour,
		MaxOutstandingMessages: int(acfg.Tasks), // If < 1, default.
		MaxOutstandingBytes:    0,               // Default.
	}

	// Initialize our Storage.
	//
	// NOTE: We're requesting read/write access even though we only need read-only
	// access because GKE doesn't understand the read-only scope:
	// https://www.googleapis.com/auth/bigtable.readonly
	st, err := a.IntermediateStorage(c, true)
	if err != nil {
		log.WithError(err).Errorf(c, "Failed to get storage instance.")
		return err
	}
	defer st.Close()

	// Initialize our Google Storage client.
	gsClient, err := a.GSClient(c)
	if err != nil {
		log.WithError(err).Errorf(c, "Failed to get Google Storage client.")
		return err
	}
	defer gsClient.Close()

	// Initialize a Coordinator client that bundles requests together.
	coordClient := &bundleServicesClient.Client{
		ServicesClient:       a.Coordinator(),
		DelayThreshold:       time.Second,
		BundleCountThreshold: 100,
	}
	defer coordClient.Flush()

	ar := archivist.Archivist{
		Service:        coordClient,
		SettingsLoader: a.GetSettingsLoader(acfg),
		Storage:        st,
		GSClient:       gsClient,
	}

	// Application shutdown will now operate by stopping the Iterator.
	c, cancelFunc := context.WithCancel(c)
	defer cancelFunc()

	// Application shutdown will now operate by cancelling the Collector's
	// shutdown Context.
	a.SetShutdownFunc(cancelFunc)

	// Execute our main subscription pull loop. It will run until the supplied
	// Context is cancelled.
	log.Fields{
		"subscription": taskSub,
	}.Infof(c, "Pulling tasks from Pub/Sub subscription.")

	retryForever := func() retry.Iterator {
		return &retry.ExponentialBackoff{
			Limited: retry.Limited{
				Delay:   200 * time.Millisecond,
				Retries: -1, // Unlimited.
			},
			MaxDelay:   10 * time.Second,
			Multiplier: 2,
		}
	}

	err = retry.Retry(c, transient.Only(retryForever), func() error {
		return grpcutil.WrapIfTransient(sub.Receive(c, func(c context.Context, msg *pubsub.Message) {
			c = log.SetFields(c, log.Fields{
				"messageID": msg.ID,
			})

			// ACK or NACK the message based on whether our task was consumed.
			deleteTask := false
			defer func() {
				// ACK the message if it is completed. If not, NACK it.
				if deleteTask {
					msg.Ack()
				} else {
					msg.Nack()
				}
			}()

			// Time how long task processing takes for metrics.
			startTime := clock.Now(c)
			defer func() {
				duration := clock.Now(c).Sub(startTime)

				if deleteTask {
					log.Fields{
						"duration": duration,
					}.Infof(c, "Task successfully processed; deleting.")
				} else {
					log.Fields{
						"duration": duration,
					}.Infof(c, "Task processing incomplete. Not deleting.")
				}

				// Add to our processing time metric.
				tsTaskProcessingTime.Add(c, duration.Seconds()*1000, deleteTask)
			}()

			task, err := makePubSubArchivistTask(c, psSubscriptionName, msg)
			c = log.SetFields(c, log.Fields{
				"consumed":         task.consumed,
				"subscriptionName": task.subscriptionName,
				"taskTimestamp":    task.timestamp.Format(time.RFC3339Nano),
				"archiveTask":      task.at,
			})
			if task.msg != nil {
				// Log all fields except data.
				c = log.SetFields(c, log.Fields{
					"message": map[string]interface{}{
						"id":          task.msg.ID,
						"attributes":  task.msg.Attributes,
						"publishTime": task.msg.PublishTime.Format(time.RFC3339Nano),
					},
				})
			}
			if err != nil {
				log.WithError(err).Errorf(c, "Failed to unmarshal archive task from message.")
				deleteTask = true
				return
			}

			ar.ArchiveTask(c, task)
			deleteTask = task.consumed
		}))
	}, func(err error, d time.Duration) {
		log.Fields{
			log.ErrorKey: err,
			"delay":      d,
		}.Warningf(c, "Transient error during subscription Receive loop; retrying...")
	})

	if err := errors.Unwrap(err); err != nil && err != context.Canceled {
		log.WithError(err).Errorf(c, "Failed during Pub/Sub Receive.")
		return err
	}

	log.Debugf(c, "Archivist finished.")
	return nil
}

// GetSettingsLoader is an archivist.SettingsLoader implementation that merges
// global and project-specific settings.
//
// The resulting settings object will be verified by the Archivist.
func (a *application) GetSettingsLoader(acfg *svcconfig.Archivist) archivist.SettingsLoader {
	serviceID := a.ServiceID()

	return func(c context.Context, proj types.ProjectName) (*archivist.Settings, error) {
		// Fold in our project-specific configuration, if valid.
		pcfg, err := a.ProjectConfig(c, proj)
		if err != nil {
			log.Fields{
				log.ErrorKey: err,
				"project":    proj,
			}.Errorf(c, "Failed to fetch project configuration.")
			return nil, err
		}

		indexParam := func(get func(ic *svcconfig.ArchiveIndexConfig) int32) int {
			if ic := pcfg.ArchiveIndexConfig; ic != nil {
				if v := get(ic); v > 0 {
					return int(v)
				}
			}

			if ic := acfg.ArchiveIndexConfig; ic != nil {
				if v := get(ic); v > 0 {
					return int(v)
				}
			}

			return 0
		}

		// Load our base settings.
		//
		// Archival bases are:
		// Staging: gs://<services:gs_staging_bucket>/<project-id>/...
		// Archive: gs://<project:archive_gs_bucket>/<project-id>/...
		st := archivist.Settings{
			GSBase:        gs.MakePath(pcfg.ArchiveGsBucket, "").Concat(serviceID),
			GSStagingBase: gs.MakePath(acfg.GsStagingBucket, "").Concat(serviceID),

			IndexStreamRange: indexParam(func(ic *svcconfig.ArchiveIndexConfig) int32 { return ic.StreamRange }),
			IndexPrefixRange: indexParam(func(ic *svcconfig.ArchiveIndexConfig) int32 { return ic.PrefixRange }),
			IndexByteRange:   indexParam(func(ic *svcconfig.ArchiveIndexConfig) int32 { return ic.ByteRange }),
			AlwaysRender:     (acfg.RenderAllStreams || pcfg.RenderAllStreams),
		}

		// Fold project settings into loaded ones.
		return &st, nil
	}
}

// Entry point.
func main() {
	mathrand.SeedRandomly()
	a := application{
		Service: service.Service{
			Name:               "archivist",
			DefaultAuthOptions: chromeinfra.DefaultAuthOptions(),
		},
	}
	a.Run(context.Background(), a.runArchivist)
}
<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      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 controllers

import (
	"context"
	"fmt"

	lropb "google.golang.org/genproto/googleapis/longrunning"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

// GetLROOperation returns LRO operation for the specified namespace instance and operation id.
func GetLROOperation(caClientFactory ConfigAgentClientFactory, ctx context.Context, r client.Reader, namespace, id, instName string) (*lropb.Operation, error) {
	caClient, closeConn, err := caClientFactory.New(ctx, r, namespace, instName)
	if err != nil {
		return nil, err
	}
	defer closeConn()

	return caClient.GetOperation(ctx, &lropb.GetOperationRequest{Name: id})
}

// DeleteLROOperation deletes LRO operation for the specified namespace instance and operation id.
func DeleteLROOperation(caClientFactory ConfigAgentClientFactory, ctx context.Context, r client.Reader, namespace, instName, id string) error {
	caClient, closeConn, err := caClientFactory.New(ctx, r, namespace, instName)
	if err != nil {
		return err
	}
	defer closeConn()

	_, err = caClient.DeleteOperation(ctx, &lropb.DeleteOperationRequest{Name: id})
	return err
}

// Check for LRO job status
// Return (true, nil) if LRO is done without errors.
// Return (true, err) if LRO is done with an error.
// Return (false, nil) if LRO still in progress.
// Return (false, err) if other error occurred.
func IsLROOperationDone(caClientFactory ConfigAgentClientFactory, ctx context.Context, r client.Reader, namespace, id, instName string) (bool, error) {
	operation, err := GetLROOperation(caClientFactory, ctx, r, namespace, id, instName)
	if err != nil {
		return false, err
	}
	if !operation.GetDone() {
		return false, nil
	}

	// handle case when remote LRO completed unsuccessfully
	if operation.GetError() != nil {
		return true, fmt.Errorf("Operation failed with err: %s. %v", operation.GetError().GetMessage(), err)
	}

	return true, nil
}
<commit_msg>Bugfix: restore LRO functionality broken by the recent refactoring change<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      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 controllers

import (
	"context"
	"fmt"

	lropb "google.golang.org/genproto/googleapis/longrunning"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

// GetLROOperation returns LRO operation for the specified namespace instance and operation id.
func GetLROOperation(caClientFactory ConfigAgentClientFactory, ctx context.Context, r client.Reader, namespace, id, instName string) (*lropb.Operation, error) {
	caClient, closeConn, err := caClientFactory.New(ctx, r, namespace, instName)
	if err != nil {
		return nil, err
	}
	defer closeConn()

	return caClient.GetOperation(ctx, &lropb.GetOperationRequest{Name: id})
}

// DeleteLROOperation deletes LRO operation for the specified namespace instance and operation id.
func DeleteLROOperation(caClientFactory ConfigAgentClientFactory, ctx context.Context, r client.Reader, namespace, id, instName string) error {
	caClient, closeConn, err := caClientFactory.New(ctx, r, namespace, instName)
	if err != nil {
		return err
	}
	defer closeConn()

	_, err = caClient.DeleteOperation(ctx, &lropb.DeleteOperationRequest{Name: id})
	return err
}

// Check for LRO job status
// Return (true, nil) if LRO is done without errors.
// Return (true, err) if LRO is done with an error.
// Return (false, nil) if LRO still in progress.
// Return (false, err) if other error occurred.
func IsLROOperationDone(caClientFactory ConfigAgentClientFactory, ctx context.Context, r client.Reader, namespace, id, instName string) (bool, error) {
	operation, err := GetLROOperation(caClientFactory, ctx, r, namespace, id, instName)
	if err != nil {
		return false, err
	}
	if !operation.GetDone() {
		return false, nil
	}

	// handle case when remote LRO completed unsuccessfully
	if operation.GetError() != nil {
		return true, fmt.Errorf("Operation failed with err: %s. %v", operation.GetError().GetMessage(), err)
	}

	return true, nil
}
<|endoftext|>