<commit_before>// Package tradingdays - Check the day is open or not and at which period time
// 股市開休市判斷（支援非國定假日：颱風假）與當日區間判斷（盤中、盤後、盤後盤）
//
package tradingdays

import (
	"bytes"
	"encoding/csv"
	"io"
	"io/ioutil"
	"net/http"
	"path"
	"runtime"
	"time"

	"github.com/toomore/gogrs/utils"
)

// IsOpen 判斷是否為開休日
func IsOpen(year int, month time.Month, day int) bool {
	d := time.Date(year, month, day, 0, 0, 0, 0, time.UTC)
	if openornot, ok := exceptDays[d.Unix()]; ok {
		return openornot
	}
	if d.In(utils.TaipeiTimeZone).Weekday() == 0 || d.In(utils.TaipeiTimeZone).Weekday() == 6 {
		return false
	}
	return true
}

// FindRecentlyOpened 回傳最近一個開市時間
func FindRecentlyOpened() time.Time {
	d := time.Now().UTC()
	days := d.Day()
	for {
		if IsOpen(d.Year(), d.Month(), days) {
			return time.Date(d.Year(), d.Month(), days, 0, 0, 0, 0, time.UTC)
		}
		days--
	}
}

var exceptDays map[int64]bool
var timeLayout = "2006/1/2"

func readCSV() {
	_, filename, _, _ := runtime.Caller(1)
	data, _ := ioutil.ReadFile(path.Join(path.Dir(filename), "list.csv"))
	processCSV(bytes.NewReader(data))
}

// DownloadCSV 更新開休市表
//
// 從 https://s3-ap-northeast-1.amazonaws.com/toomore/gogrs/list.csv
// 下載表更新，主要發生在非國定假日，如：颱風假。
func DownloadCSV(replace bool) {
	resp, _ := http.Get("https://s3-ap-northeast-1.amazonaws.com/toomore/gogrs/list.csv")
	defer resp.Body.Close()
	if data, err := ioutil.ReadAll(resp.Body); err == nil {
		if replace {
			exceptDays = make(map[int64]bool)
		}
		processCSV(bytes.NewReader(data))
	}
}

func processCSV(data io.Reader) {
	csvdata := csv.NewReader(data)

	for {
		record, err := csvdata.Read()
		if err == io.EOF {
			break
		}
		if t, err := time.ParseInLocation(timeLayout, record[0], time.UTC); err == nil {
			var isopen bool
			if record[1] == "1" {
				isopen = true
			}
			exceptDays[t.Unix()] = isopen
		}
	}
}

// TimePeriod 簡單計算時間是否在當日的開盤前、盤中、盤後盤、收盤的時間
/*
相關時間：
	- 開盤前	00:00 ~ 09:00
	- 盤中	09:00 ~ 13:30
	- 盤後盤	13:30 ~ 14:30
	- 收盤	14:30 ~ 24:00

證交所會在 13:30 後產生一次報表，14:30 盤後盤零股交易後還會再產生一次報表。
*/
type TimePeriod struct {
	target      int64
	zero        int64 // 00:00
	start       int64 // 09:00
	firstclose  int64 // 13:30
	secondclose int64 // 14:40
	night       int64 // 24:00
}

// NewTimePeriod 建立一個時間區間判斷
func NewTimePeriod(date time.Time) *TimePeriod {
	var d = &lazyTime{date: date.In(utils.TaipeiTimeZone)}
	return &TimePeriod{
		target:      d.date.Unix(),
		zero:        d.time(0, 0).Unix(),
		start:       d.time(9, 0).Unix(),
		firstclose:  d.time(13, 30).Unix(),
		secondclose: d.time(14, 30).Unix(),
		night:       d.time(24, 0).Unix(),
	}
}

// AtBefore 檢查是否為 開盤前（0000-0900） 期間
func (t TimePeriod) AtBefore() bool {
	if t.target >= t.zero && t.target < t.start {
		return true
	}

	return false
}

// AtOpen 檢查是否為 開盤（0900-1330） 期間
func (t TimePeriod) AtOpen() bool {
	if t.target >= t.start && t.target < t.firstclose {
		return true
	}
	return false
}

// AtAfterOpen 檢查是否為 盤後盤（1330-1430） 期間
func (t TimePeriod) AtAfterOpen() bool {
	if t.target >= t.firstclose && t.target < t.secondclose {
		return true
	}
	return false
}

// AtClose 檢查是否為 收盤（1430-2400）期間
func (t TimePeriod) AtClose() bool {
	if t.target >= t.secondclose && t.target < t.night {
		return true
	}
	return false
}

type lazyTime struct {
	date time.Time
}

func (t lazyTime) time(hour, min int) time.Time {
	return time.Date(t.date.Year(), t.date.Month(), t.date.Day(), hour, min, 0, 0, t.date.Location())
}

func init() {
	exceptDays = make(map[int64]bool)
	readCSV()
}
<commit_msg>Add TimePeriod into `FindRecentlyOpened`.<commit_after>// Package tradingdays - Check the day is open or not and at which period time
// 股市開休市判斷（支援非國定假日：颱風假）與當日區間判斷（盤中、盤後、盤後盤）
//
package tradingdays

import (
	"bytes"
	"encoding/csv"
	"io"
	"io/ioutil"
	"net/http"
	"path"
	"runtime"
	"time"

	"github.com/toomore/gogrs/utils"
)

// IsOpen 判斷是否為開休日
func IsOpen(year int, month time.Month, day int) bool {
	d := time.Date(year, month, day, 0, 0, 0, 0, time.UTC)
	if openornot, ok := exceptDays[d.Unix()]; ok {
		return openornot
	}
	if d.In(utils.TaipeiTimeZone).Weekday() == 0 || d.In(utils.TaipeiTimeZone).Weekday() == 6 {
		return false
	}
	return true
}

// FindRecentlyOpened 回傳最近一個開市時間（UTC 0）
func FindRecentlyOpened() time.Time {
	d := time.Now().UTC()
	days := d.Day()
	var result time.Time
	var tp *TimePeriod
	for {
		if IsOpen(d.Year(), d.Month(), days) {
			result = time.Date(d.Year(), d.Month(), days, 0, 0, 0, 0, time.UTC)
			tp = NewTimePeriod(time.Date(d.Year(), d.Month(), days, d.Hour(), d.Minute(), d.Second(), d.Nanosecond(), time.UTC))
			if tp.AtAfterOpen() || tp.AtClose() {
				return result
			}
		}
		days--
	}

}

var exceptDays map[int64]bool
var timeLayout = "2006/1/2"

func readCSV() {
	_, filename, _, _ := runtime.Caller(1)
	data, _ := ioutil.ReadFile(path.Join(path.Dir(filename), "list.csv"))
	processCSV(bytes.NewReader(data))
}

// DownloadCSV 更新開休市表
//
// 從 https://s3-ap-northeast-1.amazonaws.com/toomore/gogrs/list.csv
// 下載表更新，主要發生在非國定假日，如：颱風假。
func DownloadCSV(replace bool) {
	resp, _ := http.Get("https://s3-ap-northeast-1.amazonaws.com/toomore/gogrs/list.csv")
	defer resp.Body.Close()
	if data, err := ioutil.ReadAll(resp.Body); err == nil {
		if replace {
			exceptDays = make(map[int64]bool)
		}
		processCSV(bytes.NewReader(data))
	}
}

func processCSV(data io.Reader) {
	csvdata := csv.NewReader(data)

	for {
		record, err := csvdata.Read()
		if err == io.EOF {
			break
		}
		if t, err := time.ParseInLocation(timeLayout, record[0], time.UTC); err == nil {
			var isopen bool
			if record[1] == "1" {
				isopen = true
			}
			exceptDays[t.Unix()] = isopen
		}
	}
}

// TimePeriod 簡單計算時間是否在當日的開盤前、盤中、盤後盤、收盤的時間
/*
相關時間：
	- 開盤前	00:00 ~ 09:00
	- 盤中	09:00 ~ 13:30
	- 盤後盤	13:30 ~ 14:30
	- 收盤	14:30 ~ 24:00

證交所會在 13:30 後產生一次報表，14:30 盤後盤零股交易後還會再產生一次報表。
*/
type TimePeriod struct {
	target      int64
	zero        int64 // 00:00
	start       int64 // 09:00
	firstclose  int64 // 13:30
	secondclose int64 // 14:40
	night       int64 // 24:00
}

// NewTimePeriod 建立一個時間區間判斷
func NewTimePeriod(date time.Time) *TimePeriod {
	var d = &lazyTime{date: date.In(utils.TaipeiTimeZone)}
	return &TimePeriod{
		target:      d.date.Unix(),
		zero:        d.time(0, 0).Unix(),
		start:       d.time(9, 0).Unix(),
		firstclose:  d.time(13, 30).Unix(),
		secondclose: d.time(14, 30).Unix(),
		night:       d.time(24, 0).Unix(),
	}
}

// AtBefore 檢查是否為 開盤前（0000-0900） 期間
func (t TimePeriod) AtBefore() bool {
	if t.target >= t.zero && t.target < t.start {
		return true
	}

	return false
}

// AtOpen 檢查是否為 開盤（0900-1330） 期間
func (t TimePeriod) AtOpen() bool {
	if t.target >= t.start && t.target < t.firstclose {
		return true
	}
	return false
}

// AtAfterOpen 檢查是否為 盤後盤（1330-1430） 期間
func (t TimePeriod) AtAfterOpen() bool {
	if t.target >= t.firstclose && t.target < t.secondclose {
		return true
	}
	return false
}

// AtClose 檢查是否為 收盤（1430-2400）期間
func (t TimePeriod) AtClose() bool {
	if t.target >= t.secondclose && t.target < t.night {
		return true
	}
	return false
}

type lazyTime struct {
	date time.Time
}

func (t lazyTime) time(hour, min int) time.Time {
	return time.Date(t.date.Year(), t.date.Month(), t.date.Day(), hour, min, 0, 0, t.date.Location())
}

func init() {
	exceptDays = make(map[int64]bool)
	readCSV()
}
<|endoftext|>


<commit_before>package main

import (
  "bytes"
  "compress/zlib"
  "crypto/tls"
  "crypto/x509"
  "encoding/binary"
  "encoding/pem"
  "errors"
  "fmt"
  "io"
  "io/ioutil"
  "log"
  "math/rand"
  "net"
  "os"
  "regexp"
  "strconv"
  "time"
)

var hostname string
var hostport_re, _ = regexp.Compile("^(.+):([0-9]+)$")

func init() {
  log.Printf("publisher init\n")
  hostname, _ = os.Hostname()
  rand.Seed(time.Now().UnixNano())
}

func Publishv1(input chan []*FileEvent,
  registrar chan []*FileEvent,
  config *NetworkConfig) {
  var buffer bytes.Buffer
  var compressed_payload []byte
  var socket *tls.Conn
  var last_ack_sequence uint32
  var sequence uint32
  var err error

  socket = connect(config)
  defer socket.Close()

  // TODO(driskell): Make the idle timeout configurable like the network timeout is?
  timer := time.NewTimer(900 * time.Second)

  for {
    select {
    case events := <-input:
      for {
        // Do we need to populate the buffer again? Or do we already have it done?
        if buffer.Len() == 0 {
          sequence = last_ack_sequence
          compressor, _ := zlib.NewWriterLevel(&buffer, 3)

          for _, event := range events {
            sequence += 1
            writeDataFrame(event, sequence, compressor)
          }
          compressor.Flush()
          compressor.Close()

          compressed_payload = buffer.Bytes()
        }

        // Abort if our whole request takes longer than the configured network timeout.
        socket.SetDeadline(time.Now().Add(config.timeout))

        // Set the window size to the length of this payload in events.
        _, err = socket.Write([]byte("1W"))
        if err != nil {
          log.Printf("Socket error, will reconnect: %s\n", err)
          goto RetryPayload
        }
        err = binary.Write(socket, binary.BigEndian, uint32(len(events)))
        if err != nil {
          log.Printf("Socket error, will reconnect: %s\n", err)
          goto RetryPayload
        }

        // Write compressed frame
        _, err = socket.Write([]byte("1C"))
        if err != nil {
          log.Printf("Socket error, will reconnect: %s\n", err)
          goto RetryPayload
        }
        err = binary.Write(socket, binary.BigEndian, uint32(len(compressed_payload)))
        if err != nil {
          log.Printf("Socket error, will reconnect: %s\n", err)
          goto RetryPayload
        }
        _, err = socket.Write(compressed_payload)
        if err != nil {
          log.Printf("Socket error, will reconnect: %s\n", err)
          goto RetryPayload
        }

        // read ack
        for {
          var frame [2]byte

          // Each time we've received a frame, reset the deadline
          socket.SetDeadline(time.Now().Add(config.timeout))

          err = binary.Read(socket, binary.BigEndian, &frame)
          if err != nil {
            log.Printf("Socket error, will reconnect: %s\n", err)
            goto RetryPayload
          }

          if frame == [2]byte{'1', 'A'} {
            var ack_sequence uint32

            // Read the sequence number acked
            err = binary.Read(socket, binary.BigEndian, &ack_sequence)
            if err != nil {
              log.Printf("Socket error, will reconnect: %s\n", err)
              goto RetryPayload
            }

            if sequence == ack_sequence {
              last_ack_sequence = ack_sequence
              // All acknowledged! Stop reading acks
              break
            }

            // NOTE(driskell): If the server is busy and not yet processed anything, we MAY
            // end up receiving an ack for the last sequence in the previous payload, or 0
            if ack_sequence == last_ack_sequence {
              // Just keep waiting
              continue
            } else if ack_sequence - last_ack_sequence > uint32(len(events)) {
              // This is wrong - we've already had an ack for these
              log.Printf("Socket error, will reconnect: Repeated ACK\n")
              goto RetryPayload
            }

            // Send a slice of the acknowledged events downstream and slice what we're still waiting for
            // so that if we encounter an error, we only resend unacknowledged events
            registrar <- events[:ack_sequence - last_ack_sequence]
            events = events[ack_sequence - last_ack_sequence:]
            last_ack_sequence = ack_sequence

            // Reset the events buffer so it gets regenerated if we need to retry the payload
            buffer.Truncate(0)
            continue
          }

          // Unknown frame!
          log.Printf("Socket error, will reconnect: Unknown frame received: %s\n", frame)
          goto RetryPayload
        }

        // Success, stop trying to send the payload.
        break

      RetryPayload:
        // TODO(sissel): Track how frequently we timeout and reconnect. If we're
        // timing out too frequently, there's really no point in timing out since
        // basically everything is slow or down. We'll want to ratchet up the
        // timeout value slowly until things improve, then ratchet it down once
        // things seem healthy.
        time.Sleep(1 * time.Second)
        socket.Close()
        socket = connect(config)
      }

      // Tell the registrar that we've successfully sent the remainder of the events
      registrar <- events

      // Reset the events buffer
      buffer.Truncate(0)

      // Prepare to enter idle by setting a long deadline... if we have more events we'll drop it down again
      socket.SetDeadline(time.Now().Add(1800 * time.Second))

      // Reset the timer
      timer.Reset(900 * time.Second)
    case <-timer.C:
      // We've no events to send - throw a ping (well... window frame) so our connection doesn't idle and die
      err = ping(config, socket)
      if err != nil {
        log.Printf("Socket error during ping, will reconnect: %s\n", err)
        time.Sleep(1 * time.Second)
        socket.Close()
        socket = connect(config)
      }

      // Reset the deadline
      socket.SetDeadline(time.Now().Add(1800 * time.Second))

      // Reset the timer
      timer.Reset(900 * time.Second)
    } /* select */
  } /* for */
} // Publish

func ping(config *NetworkConfig, socket *tls.Conn) error {
  var frame [2]byte

  // Set deadline for this write
  socket.SetDeadline(time.Now().Add(config.timeout))

  // This just keeps connection open through firewalls
  // We don't await for a response so its not a real ping, the protocol does not provide for a real ping
  // And with a complete replacement of protocol happening soon, makes no sense to add new frames and such
  _, err := socket.Write([]byte("1W"))
  if err != nil {
    return err
  }
  err = binary.Write(socket, binary.BigEndian, uint32(*spool_size))
  if err != nil {
    return err
  }

  return nil
}

func connect(config *NetworkConfig) (socket *tls.Conn) {
  var tlsconfig tls.Config

  if len(config.SSLCertificate) > 0 && len(config.SSLKey) > 0 {
    log.Printf("Loading client ssl certificate: %s and %s\n",
      config.SSLCertificate, config.SSLKey)
    cert, err := tls.LoadX509KeyPair(config.SSLCertificate, config.SSLKey)
    if err != nil {
      log.Fatalf("Failed loading client ssl certificate: %s\n", err)
    }
    tlsconfig.Certificates = []tls.Certificate{cert}
  }

  if len(config.SSLCA) > 0 {
    log.Printf("Setting trusted CA from file: %s\n", config.SSLCA)
    tlsconfig.RootCAs = x509.NewCertPool()

    pemdata, err := ioutil.ReadFile(config.SSLCA)
    if err != nil {
      log.Fatalf("Failure reading CA certificate: %s\n", err)
    }

    block, _ := pem.Decode(pemdata)
    if block == nil {
      log.Fatalf("Failed to decode PEM data, is %s a valid cert?\n", config.SSLCA)
    }
    if block.Type != "CERTIFICATE" {
      log.Fatalf("This is not a certificate file: %s\n", config.SSLCA)
    }

    cert, err := x509.ParseCertificate(block.Bytes)
    if err != nil {
      log.Fatalf("Failed to parse a certificate: %s\n", config.SSLCA)
    }
    tlsconfig.RootCAs.AddCert(cert)
  }

  for {
    // Pick a random server from the list.
    hostport := config.Servers[rand.Int()%len(config.Servers)]
    submatch := hostport_re.FindSubmatch([]byte(hostport))
    if submatch == nil {
      log.Fatalf("Invalid host:port given: %s", hostport)
    }
    host := string(submatch[1])
    port := string(submatch[2])
    addresses, err := net.LookupHost(host)

    if err != nil {
      log.Printf("DNS lookup failure \"%s\": %s\n", host, err)
      time.Sleep(1 * time.Second)
      continue
    }

    address := addresses[rand.Int()%len(addresses)]
    addressport := fmt.Sprintf("%s:%s", address, port)

    log.Printf("Connecting to %s (%s) \n", addressport, host)

    tcpsocket, err := net.DialTimeout("tcp", addressport, config.timeout)
    if err != nil {
      log.Printf("Failure connecting to %s: %s\n", address, err)
      time.Sleep(1 * time.Second)
      continue
    }

    socket = tls.Client(tcpsocket, &tlsconfig)
    socket.SetDeadline(time.Now().Add(config.timeout))
    err = socket.Handshake()
    if err != nil {
      log.Printf("Handshake failure with %s: Failed to TLS handshake: %s\n", address, err)
      goto TryNextServer
    }

    log.Printf("Connected with %s\n", address)

    // connected, let's rock and roll.
    return

  TryNextServer:
    time.Sleep(1 * time.Second)
    socket.Close()
    continue
  }
  return
}

func writeDataFrame(event *FileEvent, sequence uint32, output io.Writer) {
  //log.Printf("event: %s\n", *event.Text)
  // header, "2D"
  // Why version 2 data frame? Because server.rb will correctly start returning partial ACKs if we specify version 2
  // This keeps the old logstash forwarders, which broke on partial ACK, working with even the newer server.rb
  // If the newer server.rb receives a 1D it will refuse to send partial ACK, just like before
  output.Write([]byte("2D"))
  // sequence number
  binary.Write(output, binary.BigEndian, uint32(sequence))
  // 'pair' count
  binary.Write(output, binary.BigEndian, uint32(len(*event.Fields)+4))

  writeKV("file", *event.Source, output)
  writeKV("host", hostname, output)
  writeKV("offset", strconv.FormatInt(event.Offset, 10), output)
  writeKV("line", *event.Text, output)
  for k, v := range *event.Fields {
    writeKV(k, v, output)
  }
}

func writeKV(key string, value string, output io.Writer) {
  //log.Printf("kv: %d/%s %d/%s\n", len(key), key, len(value), value)
  binary.Write(output, binary.BigEndian, uint32(len(key)))
  output.Write([]byte(key))
  binary.Write(output, binary.BigEndian, uint32(len(value)))
  output.Write([]byte(value))
}
<commit_msg>Fixed imported but not used<commit_after>package main

import (
  "bytes"
  "compress/zlib"
  "crypto/tls"
  "crypto/x509"
  "encoding/binary"
  "encoding/pem"
  "fmt"
  "io"
  "io/ioutil"
  "log"
  "math/rand"
  "net"
  "os"
  "regexp"
  "strconv"
  "time"
)

var hostname string
var hostport_re, _ = regexp.Compile("^(.+):([0-9]+)$")

func init() {
  log.Printf("publisher init\n")
  hostname, _ = os.Hostname()
  rand.Seed(time.Now().UnixNano())
}

func Publishv1(input chan []*FileEvent,
  registrar chan []*FileEvent,
  config *NetworkConfig) {
  var buffer bytes.Buffer
  var compressed_payload []byte
  var socket *tls.Conn
  var last_ack_sequence uint32
  var sequence uint32
  var err error

  socket = connect(config)
  defer socket.Close()

  // TODO(driskell): Make the idle timeout configurable like the network timeout is?
  timer := time.NewTimer(900 * time.Second)

  for {
    select {
    case events := <-input:
      for {
        // Do we need to populate the buffer again? Or do we already have it done?
        if buffer.Len() == 0 {
          sequence = last_ack_sequence
          compressor, _ := zlib.NewWriterLevel(&buffer, 3)

          for _, event := range events {
            sequence += 1
            writeDataFrame(event, sequence, compressor)
          }
          compressor.Flush()
          compressor.Close()

          compressed_payload = buffer.Bytes()
        }

        // Abort if our whole request takes longer than the configured network timeout.
        socket.SetDeadline(time.Now().Add(config.timeout))

        // Set the window size to the length of this payload in events.
        _, err = socket.Write([]byte("1W"))
        if err != nil {
          log.Printf("Socket error, will reconnect: %s\n", err)
          goto RetryPayload
        }
        err = binary.Write(socket, binary.BigEndian, uint32(len(events)))
        if err != nil {
          log.Printf("Socket error, will reconnect: %s\n", err)
          goto RetryPayload
        }

        // Write compressed frame
        _, err = socket.Write([]byte("1C"))
        if err != nil {
          log.Printf("Socket error, will reconnect: %s\n", err)
          goto RetryPayload
        }
        err = binary.Write(socket, binary.BigEndian, uint32(len(compressed_payload)))
        if err != nil {
          log.Printf("Socket error, will reconnect: %s\n", err)
          goto RetryPayload
        }
        _, err = socket.Write(compressed_payload)
        if err != nil {
          log.Printf("Socket error, will reconnect: %s\n", err)
          goto RetryPayload
        }

        // read ack
        for {
          var frame [2]byte

          // Each time we've received a frame, reset the deadline
          socket.SetDeadline(time.Now().Add(config.timeout))

          err = binary.Read(socket, binary.BigEndian, &frame)
          if err != nil {
            log.Printf("Socket error, will reconnect: %s\n", err)
            goto RetryPayload
          }

          if frame == [2]byte{'1', 'A'} {
            var ack_sequence uint32

            // Read the sequence number acked
            err = binary.Read(socket, binary.BigEndian, &ack_sequence)
            if err != nil {
              log.Printf("Socket error, will reconnect: %s\n", err)
              goto RetryPayload
            }

            if sequence == ack_sequence {
              last_ack_sequence = ack_sequence
              // All acknowledged! Stop reading acks
              break
            }

            // NOTE(driskell): If the server is busy and not yet processed anything, we MAY
            // end up receiving an ack for the last sequence in the previous payload, or 0
            if ack_sequence == last_ack_sequence {
              // Just keep waiting
              continue
            } else if ack_sequence - last_ack_sequence > uint32(len(events)) {
              // This is wrong - we've already had an ack for these
              log.Printf("Socket error, will reconnect: Repeated ACK\n")
              goto RetryPayload
            }

            // Send a slice of the acknowledged events downstream and slice what we're still waiting for
            // so that if we encounter an error, we only resend unacknowledged events
            registrar <- events[:ack_sequence - last_ack_sequence]
            events = events[ack_sequence - last_ack_sequence:]
            last_ack_sequence = ack_sequence

            // Reset the events buffer so it gets regenerated if we need to retry the payload
            buffer.Truncate(0)
            continue
          }

          // Unknown frame!
          log.Printf("Socket error, will reconnect: Unknown frame received: %s\n", frame)
          goto RetryPayload
        }

        // Success, stop trying to send the payload.
        break

      RetryPayload:
        // TODO(sissel): Track how frequently we timeout and reconnect. If we're
        // timing out too frequently, there's really no point in timing out since
        // basically everything is slow or down. We'll want to ratchet up the
        // timeout value slowly until things improve, then ratchet it down once
        // things seem healthy.
        time.Sleep(1 * time.Second)
        socket.Close()
        socket = connect(config)
      }

      // Tell the registrar that we've successfully sent the remainder of the events
      registrar <- events

      // Reset the events buffer
      buffer.Truncate(0)

      // Prepare to enter idle by setting a long deadline... if we have more events we'll drop it down again
      socket.SetDeadline(time.Now().Add(1800 * time.Second))

      // Reset the timer
      timer.Reset(900 * time.Second)
    case <-timer.C:
      // We've no events to send - throw a ping (well... window frame) so our connection doesn't idle and die
      err = ping(config, socket)
      if err != nil {
        log.Printf("Socket error during ping, will reconnect: %s\n", err)
        time.Sleep(1 * time.Second)
        socket.Close()
        socket = connect(config)
      }

      // Reset the deadline
      socket.SetDeadline(time.Now().Add(1800 * time.Second))

      // Reset the timer
      timer.Reset(900 * time.Second)
    } /* select */
  } /* for */
} // Publish

func ping(config *NetworkConfig, socket *tls.Conn) error {
  var frame [2]byte

  // Set deadline for this write
  socket.SetDeadline(time.Now().Add(config.timeout))

  // This just keeps connection open through firewalls
  // We don't await for a response so its not a real ping, the protocol does not provide for a real ping
  // And with a complete replacement of protocol happening soon, makes no sense to add new frames and such
  _, err := socket.Write([]byte("1W"))
  if err != nil {
    return err
  }
  err = binary.Write(socket, binary.BigEndian, uint32(*spool_size))
  if err != nil {
    return err
  }

  return nil
}

func connect(config *NetworkConfig) (socket *tls.Conn) {
  var tlsconfig tls.Config

  if len(config.SSLCertificate) > 0 && len(config.SSLKey) > 0 {
    log.Printf("Loading client ssl certificate: %s and %s\n",
      config.SSLCertificate, config.SSLKey)
    cert, err := tls.LoadX509KeyPair(config.SSLCertificate, config.SSLKey)
    if err != nil {
      log.Fatalf("Failed loading client ssl certificate: %s\n", err)
    }
    tlsconfig.Certificates = []tls.Certificate{cert}
  }

  if len(config.SSLCA) > 0 {
    log.Printf("Setting trusted CA from file: %s\n", config.SSLCA)
    tlsconfig.RootCAs = x509.NewCertPool()

    pemdata, err := ioutil.ReadFile(config.SSLCA)
    if err != nil {
      log.Fatalf("Failure reading CA certificate: %s\n", err)
    }

    block, _ := pem.Decode(pemdata)
    if block == nil {
      log.Fatalf("Failed to decode PEM data, is %s a valid cert?\n", config.SSLCA)
    }
    if block.Type != "CERTIFICATE" {
      log.Fatalf("This is not a certificate file: %s\n", config.SSLCA)
    }

    cert, err := x509.ParseCertificate(block.Bytes)
    if err != nil {
      log.Fatalf("Failed to parse a certificate: %s\n", config.SSLCA)
    }
    tlsconfig.RootCAs.AddCert(cert)
  }

  for {
    // Pick a random server from the list.
    hostport := config.Servers[rand.Int()%len(config.Servers)]
    submatch := hostport_re.FindSubmatch([]byte(hostport))
    if submatch == nil {
      log.Fatalf("Invalid host:port given: %s", hostport)
    }
    host := string(submatch[1])
    port := string(submatch[2])
    addresses, err := net.LookupHost(host)

    if err != nil {
      log.Printf("DNS lookup failure \"%s\": %s\n", host, err)
      time.Sleep(1 * time.Second)
      continue
    }

    address := addresses[rand.Int()%len(addresses)]
    addressport := fmt.Sprintf("%s:%s", address, port)

    log.Printf("Connecting to %s (%s) \n", addressport, host)

    tcpsocket, err := net.DialTimeout("tcp", addressport, config.timeout)
    if err != nil {
      log.Printf("Failure connecting to %s: %s\n", address, err)
      time.Sleep(1 * time.Second)
      continue
    }

    socket = tls.Client(tcpsocket, &tlsconfig)
    socket.SetDeadline(time.Now().Add(config.timeout))
    err = socket.Handshake()
    if err != nil {
      log.Printf("Handshake failure with %s: Failed to TLS handshake: %s\n", address, err)
      goto TryNextServer
    }

    log.Printf("Connected with %s\n", address)

    // connected, let's rock and roll.
    return

  TryNextServer:
    time.Sleep(1 * time.Second)
    socket.Close()
    continue
  }
  return
}

func writeDataFrame(event *FileEvent, sequence uint32, output io.Writer) {
  //log.Printf("event: %s\n", *event.Text)
  // header, "2D"
  // Why version 2 data frame? Because server.rb will correctly start returning partial ACKs if we specify version 2
  // This keeps the old logstash forwarders, which broke on partial ACK, working with even the newer server.rb
  // If the newer server.rb receives a 1D it will refuse to send partial ACK, just like before
  output.Write([]byte("2D"))
  // sequence number
  binary.Write(output, binary.BigEndian, uint32(sequence))
  // 'pair' count
  binary.Write(output, binary.BigEndian, uint32(len(*event.Fields)+4))

  writeKV("file", *event.Source, output)
  writeKV("host", hostname, output)
  writeKV("offset", strconv.FormatInt(event.Offset, 10), output)
  writeKV("line", *event.Text, output)
  for k, v := range *event.Fields {
    writeKV(k, v, output)
  }
}

func writeKV(key string, value string, output io.Writer) {
  //log.Printf("kv: %d/%s %d/%s\n", len(key), key, len(value), value)
  binary.Write(output, binary.BigEndian, uint32(len(key)))
  output.Write([]byte(key))
  binary.Write(output, binary.BigEndian, uint32(len(value)))
  output.Write([]byte(value))
}
<|endoftext|>


<commit_before>package main

import (
  "bytes"
  "compress/zlib"
  "crypto/tls"
  "crypto/x509"
  "encoding/binary"
  "encoding/pem"
  "fmt"
  "io"
  "io/ioutil"
  "log"
  "math/rand"
  "net"
  "os"
  "regexp"
  "strconv"
  "time"
)

var hostname string
var hostport_re, _ = regexp.Compile("^(.+):([0-9]+)$")

func init() {
  log.Printf("publisher init\n")
  hostname, _ = os.Hostname()
  rand.Seed(time.Now().UnixNano())
}

func Publishv1(input chan []*FileEvent,
  registrar chan []*FileEvent,
  config *NetworkConfig) {
  var buffer bytes.Buffer
  var socket *tls.Conn
  var sequence uint32
  var err error

  socket = connect(config)
  defer socket.Close()

  for events := range input {
    buffer.Truncate(0)
    compressor, _ := zlib.NewWriterLevel(&buffer, 3)

    for _, event := range events {
      sequence += 1
      writeDataFrame(event, sequence, compressor)
    }
    compressor.Flush()
    compressor.Close()

    compressed_payload := buffer.Bytes()

    // Send buffer until we're successful...
    oops := func(err error) {
      // TODO(sissel): Track how frequently we timeout and reconnect. If we're
      // timing out too frequently, there's really no point in timing out since
      // basically everything is slow or down. We'll want to ratchet up the
      // timeout value slowly until things improve, then ratchet it down once
      // things seem healthy.
      log.Printf("Socket error, will reconnect: %s\n", err)
      time.Sleep(1 * time.Second)
      socket.Close()
      socket = connect(config)
    }

  SendPayload:
    for {
      // Abort if our whole request takes longer than the configured
      // network timeout.
      socket.SetDeadline(time.Now().Add(config.timeout))

      // Set the window size to the length of this payload in events.
      _, err = socket.Write([]byte("1W"))
      if err != nil {
        oops(err)
        continue
      }
      binary.Write(socket, binary.BigEndian, uint32(len(events)))
      if err != nil {
        oops(err)
        continue
      }

      // Write compressed frame
      socket.Write([]byte("1C"))
      if err != nil {
        oops(err)
        continue
      }
      binary.Write(socket, binary.BigEndian, uint32(len(compressed_payload)))
      if err != nil {
        oops(err)
        continue
      }
      _, err = socket.Write(compressed_payload)
      if err != nil {
        oops(err)
        continue
      }

      // read ack
      response := make([]byte, 0, 6)
      ackbytes := 0
      for ackbytes != 6 {
        n, err := socket.Read(response[len(response):cap(response)])
        if err != nil {
          log.Printf("Read error looking for ack: %s\n", err)
          socket.Close()
          socket = connect(config)
          continue SendPayload // retry sending on new connection
        } else {
          ackbytes += n
        }
      }

      // TODO(sissel): verify ack
      // Success, stop trying to send the payload.
      break
    }

    // Tell the registrar that we've successfully sent these events
    registrar <- events
  } /* for each event payload */
} // Publish

func connect(config *NetworkConfig) (socket *tls.Conn) {
  var tlsconfig tls.Config

  if len(config.SSLCertificate) > 0 && len(config.SSLKey) > 0 {
    log.Printf("Loading client ssl certificate: %s and %s\n",
      config.SSLCertificate, config.SSLKey)
    cert, err := tls.LoadX509KeyPair(config.SSLCertificate, config.SSLKey)
    if err != nil {
      log.Fatalf("Failed loading client ssl certificate: %s\n", err)
    }
    tlsconfig.Certificates = []tls.Certificate{cert}
  }

  if len(config.SSLCA) > 0 {
    log.Printf("Setting trusted CA from file: %s\n", config.SSLCA)
    tlsconfig.RootCAs = x509.NewCertPool()

    pemdata, err := ioutil.ReadFile(config.SSLCA)
    if err != nil {
      log.Fatalf("Failure reading CA certificate: %s\n", err)
    }

    block, _ := pem.Decode(pemdata)
    if block == nil {
      log.Fatalf("Failed to decode PEM data, is %s a valid cert?\n", config.SSLCA)
    }
    if block.Type != "CERTIFICATE" {
      log.Fatalf("This is not a certificate file: %s\n", config.SSLCA)
    }

    cert, err := x509.ParseCertificate(block.Bytes)
    if err != nil {
      log.Fatalf("Failed to parse a certificate: %s\n", config.SSLCA)
    }
    tlsconfig.RootCAs.AddCert(cert)
  }

  for {
    // Pick a random server from the list.
    hostport := config.Servers[rand.Int()%len(config.Servers)]
    submatch := hostport_re.FindSubmatch([]byte(hostport))
    if submatch == nil {
      log.Fatalf("Invalid host:port given: %s", hostport)
    }
    host := string(submatch[1])
    port := string(submatch[2])
    addresses, err := net.LookupHost(host)

    if err != nil {
      log.Printf("DNS lookup failure \"%s\": %s\n", host, err)
      time.Sleep(1 * time.Second)
      continue
    }

    address := addresses[rand.Int()%len(addresses)]
    addressport := fmt.Sprintf("%s:%s", address, port)

    log.Printf("Connecting to %s (%s) \n", addressport, host)

    tcpsocket, err := net.DialTimeout("tcp", addressport, config.timeout)
    if err != nil {
      log.Printf("Failure connecting to %s: %s\n", address, err)
      time.Sleep(1 * time.Second)
      continue
    }

    socket = tls.Client(tcpsocket, &tlsconfig)
    socket.SetDeadline(time.Now().Add(config.timeout))
    err = socket.Handshake()
    if err != nil {
      log.Printf("Failed to tls handshake with %s %s\n", address, err)
      time.Sleep(1 * time.Second)
      socket.Close()
      continue
    }

    log.Printf("Connected to %s\n", address)

    // connected, let's rock and roll.
    return
  }
  return
}

func writeDataFrame(event *FileEvent, sequence uint32, output io.Writer) {
  //log.Printf("event: %s\n", *event.Text)
  // header, "1D"
  output.Write([]byte("1D"))
  // sequence number
  binary.Write(output, binary.BigEndian, uint32(sequence))
  // 'pair' count
  binary.Write(output, binary.BigEndian, uint32(len(*event.Fields)+4))

  writeKV("file", *event.Source, output)
  writeKV("host", hostname, output)
  writeKV("offset", strconv.FormatInt(event.Offset, 10), output)
  writeKV("line", *event.Text, output)
  for k, v := range *event.Fields {
    writeKV(k, v, output)
  }
}

func writeKV(key string, value string, output io.Writer) {
  //log.Printf("kv: %d/%s %d/%s\n", len(key), key, len(value), value)
  binary.Write(output, binary.BigEndian, uint32(len(key)))
  output.Write([]byte(key))
  binary.Write(output, binary.BigEndian, uint32(len(value)))
  output.Write([]byte(value))
}
<commit_msg>support for newer SSL certificate types<commit_after>package main

import (
  "bytes"
  "compress/zlib"
  "crypto/tls"
  "crypto/x509"
  "encoding/binary"
  "encoding/pem"
  "fmt"
  "io"
  "io/ioutil"
  "log"
  "math/rand"
  "net"
  "os"
  "regexp"
  "strconv"
  "time"
)

// Support for newer SSL signature algorithms
import _ "crypto/sha256"
import _ "crypto/sha512"

var hostname string
var hostport_re, _ = regexp.Compile("^(.+):([0-9]+)$")

func init() {
  log.Printf("publisher init\n")
  hostname, _ = os.Hostname()
  rand.Seed(time.Now().UnixNano())
}

func Publishv1(input chan []*FileEvent,
  registrar chan []*FileEvent,
  config *NetworkConfig) {
  var buffer bytes.Buffer
  var socket *tls.Conn
  var sequence uint32
  var err error

  socket = connect(config)
  defer socket.Close()

  for events := range input {
    buffer.Truncate(0)
    compressor, _ := zlib.NewWriterLevel(&buffer, 3)

    for _, event := range events {
      sequence += 1
      writeDataFrame(event, sequence, compressor)
    }
    compressor.Flush()
    compressor.Close()

    compressed_payload := buffer.Bytes()

    // Send buffer until we're successful...
    oops := func(err error) {
      // TODO(sissel): Track how frequently we timeout and reconnect. If we're
      // timing out too frequently, there's really no point in timing out since
      // basically everything is slow or down. We'll want to ratchet up the
      // timeout value slowly until things improve, then ratchet it down once
      // things seem healthy.
      log.Printf("Socket error, will reconnect: %s\n", err)
      time.Sleep(1 * time.Second)
      socket.Close()
      socket = connect(config)
    }

  SendPayload:
    for {
      // Abort if our whole request takes longer than the configured
      // network timeout.
      socket.SetDeadline(time.Now().Add(config.timeout))

      // Set the window size to the length of this payload in events.
      _, err = socket.Write([]byte("1W"))
      if err != nil {
        oops(err)
        continue
      }
      binary.Write(socket, binary.BigEndian, uint32(len(events)))
      if err != nil {
        oops(err)
        continue
      }

      // Write compressed frame
      socket.Write([]byte("1C"))
      if err != nil {
        oops(err)
        continue
      }
      binary.Write(socket, binary.BigEndian, uint32(len(compressed_payload)))
      if err != nil {
        oops(err)
        continue
      }
      _, err = socket.Write(compressed_payload)
      if err != nil {
        oops(err)
        continue
      }

      // read ack
      response := make([]byte, 0, 6)
      ackbytes := 0
      for ackbytes != 6 {
        n, err := socket.Read(response[len(response):cap(response)])
        if err != nil {
          log.Printf("Read error looking for ack: %s\n", err)
          socket.Close()
          socket = connect(config)
          continue SendPayload // retry sending on new connection
        } else {
          ackbytes += n
        }
      }

      // TODO(sissel): verify ack
      // Success, stop trying to send the payload.
      break
    }

    // Tell the registrar that we've successfully sent these events
    registrar <- events
  } /* for each event payload */
} // Publish

func connect(config *NetworkConfig) (socket *tls.Conn) {
  var tlsconfig tls.Config

  if len(config.SSLCertificate) > 0 && len(config.SSLKey) > 0 {
    log.Printf("Loading client ssl certificate: %s and %s\n",
      config.SSLCertificate, config.SSLKey)
    cert, err := tls.LoadX509KeyPair(config.SSLCertificate, config.SSLKey)
    if err != nil {
      log.Fatalf("Failed loading client ssl certificate: %s\n", err)
    }
    tlsconfig.Certificates = []tls.Certificate{cert}
  }

  if len(config.SSLCA) > 0 {
    log.Printf("Setting trusted CA from file: %s\n", config.SSLCA)
    tlsconfig.RootCAs = x509.NewCertPool()

    pemdata, err := ioutil.ReadFile(config.SSLCA)
    if err != nil {
      log.Fatalf("Failure reading CA certificate: %s\n", err)
    }

    block, _ := pem.Decode(pemdata)
    if block == nil {
      log.Fatalf("Failed to decode PEM data, is %s a valid cert?\n", config.SSLCA)
    }
    if block.Type != "CERTIFICATE" {
      log.Fatalf("This is not a certificate file: %s\n", config.SSLCA)
    }

    cert, err := x509.ParseCertificate(block.Bytes)
    if err != nil {
      log.Fatalf("Failed to parse a certificate: %s\n", config.SSLCA)
    }
    tlsconfig.RootCAs.AddCert(cert)
  }

  for {
    // Pick a random server from the list.
    hostport := config.Servers[rand.Int()%len(config.Servers)]
    submatch := hostport_re.FindSubmatch([]byte(hostport))
    if submatch == nil {
      log.Fatalf("Invalid host:port given: %s", hostport)
    }
    host := string(submatch[1])
    port := string(submatch[2])
    addresses, err := net.LookupHost(host)

    if err != nil {
      log.Printf("DNS lookup failure \"%s\": %s\n", host, err)
      time.Sleep(1 * time.Second)
      continue
    }

    address := addresses[rand.Int()%len(addresses)]
    addressport := fmt.Sprintf("%s:%s", address, port)

    log.Printf("Connecting to %s (%s) \n", addressport, host)

    tcpsocket, err := net.DialTimeout("tcp", addressport, config.timeout)
    if err != nil {
      log.Printf("Failure connecting to %s: %s\n", address, err)
      time.Sleep(1 * time.Second)
      continue
    }

    socket = tls.Client(tcpsocket, &tlsconfig)
    socket.SetDeadline(time.Now().Add(config.timeout))
    err = socket.Handshake()
    if err != nil {
      log.Printf("Failed to tls handshake with %s %s\n", address, err)
      time.Sleep(1 * time.Second)
      socket.Close()
      continue
    }

    log.Printf("Connected to %s\n", address)

    // connected, let's rock and roll.
    return
  }
  return
}

func writeDataFrame(event *FileEvent, sequence uint32, output io.Writer) {
  //log.Printf("event: %s\n", *event.Text)
  // header, "1D"
  output.Write([]byte("1D"))
  // sequence number
  binary.Write(output, binary.BigEndian, uint32(sequence))
  // 'pair' count
  binary.Write(output, binary.BigEndian, uint32(len(*event.Fields)+4))

  writeKV("file", *event.Source, output)
  writeKV("host", hostname, output)
  writeKV("offset", strconv.FormatInt(event.Offset, 10), output)
  writeKV("line", *event.Text, output)
  for k, v := range *event.Fields {
    writeKV(k, v, output)
  }
}

func writeKV(key string, value string, output io.Writer) {
  //log.Printf("kv: %d/%s %d/%s\n", len(key), key, len(value), value)
  binary.Write(output, binary.BigEndian, uint32(len(key)))
  output.Write([]byte(key))
  binary.Write(output, binary.BigEndian, uint32(len(value)))
  output.Write([]byte(value))
}
<|endoftext|>


<commit_before>// Copyright 2017 CoreOS, Inc.
//
// 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 positive

import (
	"testing"

	"github.com/coreos/init/tests/coreos-install/register"
	"github.com/coreos/init/tests/coreos-install/util"
)

func init() {
	register.Register(register.Test{
		Name: "Base Test",
		Func: baseTest,
	})
	register.Register(register.Test{
		Name: "Ignition Test",
		Func: baseTest,
		IgnitionConfig: util.StringToPtr(`{
			"ignition": {"version": "2.1.0"}
		}`),
		UseLocalServer: true,
	})
	register.Register(register.Test{
		Name: "CloudConfig Test",
		Func: baseTest,
		CloudConfig: util.StringToPtr(`#cloud-config

		hostname: "coreos1"`),
		UseLocalServer: true,
	})
	register.Register(register.Test{
		Name:    "Alpha 1520.0",
		Func:    baseTest,
		Channel: util.StringToPtr("alpha"),
		Version: util.StringToPtr("1520.0.0"),
	})
	register.Register(register.Test{
		Name:    "Channel Only",
		Func:    baseTest,
		Channel: util.StringToPtr("beta"),
	})
	register.Register(register.Test{
		Name:    "arm64-usr alpha 1367.5.0",
		Func:    baseTest,
		Channel: util.StringToPtr("alpha"),
		Version: util.StringToPtr("1367.5.0"),
		Board:   util.StringToPtr("arm64-usr"),
	})
	register.Register(register.Test{
		Name:    "Version Only",
		Func:    baseTest,
		Version: util.StringToPtr("1409.7.0"),
	})
	register.Register(register.Test{
		Name: "OEM - ami",
		Func: baseTest,
		OEM:  util.StringToPtr("ami"),
	})
	register.Register(register.Test{
		Name: "OEM - cloudstack",
		Func: baseTest,
		OEM:  util.StringToPtr("cloudstack"),
	})
	register.Register(register.Test{
		Name: "OEM - digitalocean",
		Func: baseTest,
		OEM:  util.StringToPtr("digitalocean"),
	})
	register.Register(register.Test{
		Name: "OEM - packet",
		Func: baseTest,
		OEM:  util.StringToPtr("packet"),
	})
	register.Register(register.Test{
		Name: "OEM - rackspace",
		Func: baseTest,
		OEM:  util.StringToPtr("rackspace"),
	})
	register.Register(register.Test{
		Name: "OEM - vmware_raw",
		Func: baseTest,
		OEM:  util.StringToPtr("vmware_raw"),
	})
}

func baseTest(t *testing.T, test register.Test) {
	diskFile, loopDevice := test.CreateDevice(t)
	defer test.CleanupDisk(t, diskFile, loopDevice)

	test.RunCoreOSInstall(t, loopDevice)

	rootDir := test.MountPartitions(t, loopDevice)
	defer test.UnmountPartitions(t, loopDevice)

	test.DefaultChecks(t, rootDir)
}
<commit_msg>tests/coreos-install: update version only test on cl machines<commit_after>// Copyright 2017 CoreOS, Inc.
//
// 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 positive

import (
	"testing"

	"github.com/coreos/init/tests/coreos-install/register"
	"github.com/coreos/init/tests/coreos-install/util"
)

func init() {
	register.Register(register.Test{
		Name: "Base Test",
		Func: baseTest,
	})
	register.Register(register.Test{
		Name: "Ignition Test",
		Func: baseTest,
		IgnitionConfig: util.StringToPtr(`{
			"ignition": {"version": "2.1.0"}
		}`),
		UseLocalServer: true,
	})
	register.Register(register.Test{
		Name: "CloudConfig Test",
		Func: baseTest,
		CloudConfig: util.StringToPtr(`#cloud-config

		hostname: "coreos1"`),
		UseLocalServer: true,
	})
	register.Register(register.Test{
		Name:    "Alpha 1520.0",
		Func:    baseTest,
		Channel: util.StringToPtr("alpha"),
		Version: util.StringToPtr("1520.0.0"),
	})
	register.Register(register.Test{
		Name:    "Channel Only",
		Func:    baseTest,
		Channel: util.StringToPtr("beta"),
	})
	register.Register(register.Test{
		Name:    "arm64-usr alpha 1367.5.0",
		Func:    baseTest,
		Channel: util.StringToPtr("alpha"),
		Version: util.StringToPtr("1367.5.0"),
		Board:   util.StringToPtr("arm64-usr"),
	})
	register.Register(register.Test{
		Name: "Version Only",
		Func: pickVersion,
	})
	register.Register(register.Test{
		Name: "OEM - ami",
		Func: baseTest,
		OEM:  util.StringToPtr("ami"),
	})
	register.Register(register.Test{
		Name: "OEM - cloudstack",
		Func: baseTest,
		OEM:  util.StringToPtr("cloudstack"),
	})
	register.Register(register.Test{
		Name: "OEM - digitalocean",
		Func: baseTest,
		OEM:  util.StringToPtr("digitalocean"),
	})
	register.Register(register.Test{
		Name: "OEM - packet",
		Func: baseTest,
		OEM:  util.StringToPtr("packet"),
	})
	register.Register(register.Test{
		Name: "OEM - rackspace",
		Func: baseTest,
		OEM:  util.StringToPtr("rackspace"),
	})
	register.Register(register.Test{
		Name: "OEM - vmware_raw",
		Func: baseTest,
		OEM:  util.StringToPtr("vmware_raw"),
	})
}

// used by tests which want to test versions without pinning
// a specific channel as on Container Linux machines the
// channel will default to the host machines channel.
//
// Will update the test.Version flag with the selected Version
func pickVersion(t *testing.T, test register.Test) {
	pinnedVersions := map[string]string{
		"alpha":  "1535.0.0",
		"beta":   "1520.3.0",
		"stable": "1465.7.0",
	}

	channel, _, _, err := util.GetDefaultChannelBoardVersion()
	if err != nil {
		t.Fatal(err)
	}

	if version, ok := pinnedVersions[channel]; ok {
		test.Version = util.StringToPtr(version)
	} else {
		t.Fatalf("unknown channel %s", channel)
	}

	baseTest(t, test)
}

func baseTest(t *testing.T, test register.Test) {
	diskFile, loopDevice := test.CreateDevice(t)
	defer test.CleanupDisk(t, diskFile, loopDevice)

	test.RunCoreOSInstall(t, loopDevice)

	rootDir := test.MountPartitions(t, loopDevice)
	defer test.UnmountPartitions(t, loopDevice)

	test.DefaultChecks(t, rootDir)
}
<|endoftext|>


<commit_before>package log5go

import (
	"testing"
	"time"
)

const dateFmt = "2006-01-02T15:04:05"

func TestMinutelyOperations(t *testing.T) {
	traw, _ := time.Parse(dateFmt, "2014-08-11T12:15:22")
	t0, _ := time.Parse(dateFmt, "2014-08-11T12:15:00")
	t1, _ := time.Parse(dateFmt, "2014-08-11T12:16:00")
	t2, _ := time.Parse(dateFmt, "2014-08-11T12:17:00")

	tcalc := calculateNextRollTime(traw, RollMinutely)
	if !tcalc.Equal(t1) {
		t.Errorf("Error getting t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculateNextRollTime(tcalc, RollMinutely)
	if !tcalc.Equal(t2) {
		t.Errorf("Error getting t2. Expected %v but got %v", t2, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollMinutely)
	if !tcalc.Equal(t1) {
		t.Errorf("Error rolling back to t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollMinutely)
	if !tcalc.Equal(t0) {
		t.Errorf("Error getting t0. Expected %v but got %v", t0, tcalc)
	}
}

func TestHourlyOperations(t *testing.T) {
	traw, _ := time.Parse(dateFmt, "2014-08-11T23:15:22")
	t0, _ := time.Parse(dateFmt, "2014-08-11T23:00:00")
	t1, _ := time.Parse(dateFmt, "2014-08-12T00:00:00")
	t2, _ := time.Parse(dateFmt, "2014-08-12T01:00:00")

	tcalc := calculateNextRollTime(traw, RollHourly)
	if !tcalc.Equal(t1) {
		t.Errorf("Error getting t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculateNextRollTime(tcalc, RollHourly)
	if !tcalc.Equal(t2) {
		t.Errorf("Error getting t2. Expected %v but got %v", t2, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollHourly)
	if !tcalc.Equal(t1) {
		t.Errorf("Error rolling back to t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollHourly)
	if !tcalc.Equal(t0) {
		t.Errorf("Error getting t0. Expected %v but got %v", t0, tcalc)
	}
}

func TestDailyOperations(t *testing.T) {
	traw, _ := time.Parse(dateFmt, "2014-08-11T23:15:22")
	t0, _ := time.Parse(dateFmt, "2014-08-11T00:00:00")
	t1, _ := time.Parse(dateFmt, "2014-08-12T00:00:00")
	t2, _ := time.Parse(dateFmt, "2014-08-13T00:00:00")

	tcalc := calculateNextRollTime(traw, RollDaily)
	if !tcalc.Equal(t1) {
		t.Errorf("Error getting t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculateNextRollTime(tcalc, RollDaily)
	if !tcalc.Equal(t2) {
		t.Errorf("Error getting t2. Expected %v but got %v", t2, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollDaily)
	if !tcalc.Equal(t1) {
		t.Errorf("Error rolling back to t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollDaily)
	if !tcalc.Equal(t0) {
		t.Errorf("Error getting t0. Expected %v but got %v", t0, tcalc)
	}
}

func TestWeeklyOperations(t *testing.T) {
	traw, _ := time.Parse(dateFmt, "2014-08-11T23:15:22") // a Monday
	t0, _ := time.Parse(dateFmt, "2014-08-10T00:00:00")
	t1, _ := time.Parse(dateFmt, "2014-08-17T00:00:00")
	t2, _ := time.Parse(dateFmt, "2014-08-24T00:00:00")

	tcalc := calculateNextRollTime(traw, RollWeekly)
	if !tcalc.Equal(t1) {
		t.Errorf("Error getting t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculateNextRollTime(tcalc, RollWeekly)
	if !tcalc.Equal(t2) {
		t.Errorf("Error getting t2. Expected %v but got %v", t2, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollWeekly)
	if !tcalc.Equal(t1) {
		t.Errorf("Error rolling back to t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollWeekly)
	if !tcalc.Equal(t0) {
		t.Errorf("Error getting t0. Expected %v but got %v", t0, tcalc)
	}
}

func TestHourlyHandlesDSTStart(t *testing.T) {
	t0, _ := time.Parse(time.RFC822, "09 Mar 14 00:15 PST")
	t1, _ := time.Parse(time.RFC822, "09 Mar 14 01:00 PST")
	t3, _ := time.Parse(time.RFC822, "09 Mar 14 03:00 PDT")

	tnext := calculateNextRollTime(t0, RollHourly)
	if !tnext.Equal(t1) {
		t.Errorf("Error getting t1. Expected %v but got %v", t1, tnext)
	}

	tnext = calculateNextRollTime(tnext, RollHourly)
	if !tnext.Equal(t3) {
		t.Errorf("Error getting t3. Expected %v but got %v", t3, tnext)
	}

	tprev := calculatePreviousRollTime(tnext, RollHourly)
	if !tprev.Equal(t1) {
		t.Errorf("Error rolling back to t1. Expected %v but got %v", t1, tprev)
	}
}

func TestHourlyHandlesDSTEnd(t *testing.T) {
	t0, _ := time.Parse(time.RFC822, "02 Nov 14 00:15 PDT")
	t1, _ := time.Parse(time.RFC822, "02 Nov 14 01:00 PDT")
	t2, _ := time.Parse(time.RFC822, "02 Nov 14 02:00 PST")

	tnext := calculateNextRollTime(t0, RollHourly)
	if !tnext.Equal(t1) {
		t.Errorf("Error getting t1. Expected %v but got %v", t1, tnext)
	}

	tnext = calculateNextRollTime(tnext, RollHourly)
	if !tnext.Equal(t2) {
		t.Errorf("Error getting t2. Expected %v but got %v", t2, tnext)
	}

	tprev := calculatePreviousRollTime(tnext, RollHourly)
	if !tprev.Equal(t1) {
		t.Errorf("Error rolling back to t1. Expected %v but got %v", t1, tprev)
	}
}

func TestDailyHandlesDST(t *testing.T) {
	// DST start
	traw, _ := time.Parse(time.RFC822, "09 Mar 14 12:15 PDT")
	t0, _ := time.Parse(time.RFC822, "09 Mar 14 00:00 PST")
	t1, _ := time.Parse(time.RFC822, "10 Mar 14 00:00 PDT")

	tcalc := calculateNextRollTime(traw, RollDaily)
	if !tcalc.Equal(t1) {
		t.Errorf("Error getting t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollDaily)
	if !tcalc.Equal(t0) {
		t.Errorf("Error rolling back to t0. Expected %v but got %v", t0, tcalc)
	}

	// DST end
	traw, _ = time.Parse(time.RFC822, "02 Nov 14 12:15 PST")
	t0, _ = time.Parse(time.RFC822, "02 Nov 14 00:00 PDT")
	t1, _ = time.Parse(time.RFC822, "03 Nov 14 00:00 PST")

	tcalc = calculateNextRollTime(traw, RollDaily)
	if !tcalc.Equal(t1) {
		t.Errorf("Error getting t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollDaily)
	if !tcalc.Equal(t0) {
		t.Errorf("Error rolling back to t0. Expected %v but got %v", t0, tcalc)
	}
}
<commit_msg>deliberately breaking a test to test travis-ci<commit_after>package log5go

import (
	"testing"
	"time"
)

const dateFmt = "2006-01-02T15:04:05"

func TestMinutelyOperations(t *testing.T) {
	traw, _ := time.Parse(dateFmt, "2014-08-11T12:15:22")
	t0, _ := time.Parse(dateFmt, "2014-08-11T12:15:00")
	t1, _ := time.Parse(dateFmt, "2014-08-11T12:16:00")
	t2, _ := time.Parse(dateFmt, "2014-08-11T12:17:00")

	tcalc := calculateNextRollTime(traw, RollMinutely)
	if !tcalc.Equal(t1) {
		t.Errorf("Error getting t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculateNextRollTime(tcalc, RollMinutely)
	if !tcalc.Equal(t2) {
		t.Errorf("Error getting t2. Expected %v but got %v", t2, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollMinutely)
	if !tcalc.Equal(t1) {
		t.Errorf("Error rolling back to t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollMinutely)
	if !tcalc.Equal(t0) {
		t.Errorf("Error getting t0. Expected %v but got %v", t0, tcalc)
	}
}

func TestHourlyOperations(t *testing.T) {
	traw, _ := time.Parse(dateFmt, "2014-08-11T23:15:22")
	t0, _ := time.Parse(dateFmt, "2014-08-11T23:00:00")
	t1, _ := time.Parse(dateFmt, "2014-08-12T00:00:00")
	t2, _ := time.Parse(dateFmt, "2014-08-12T01:00:00")

	tcalc := calculateNextRollTime(traw, RollHourly)
	if !tcalc.Equal(t1) {
		t.Errorf("Error getting t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculateNextRollTime(tcalc, RollHourly)
	if !tcalc.Equal(t2) {
		t.Errorf("Error getting t2. Expected %v but got %v", t2, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollHourly)
	if !tcalc.Equal(t1) {
		t.Errorf("Error rolling back to t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollHourly)
	if !tcalc.Equal(t0) {
		t.Errorf("Error getting t0. Expected %v but got %v", t0, tcalc)
	}
}

func TestDailyOperations(t *testing.T) {
	traw, _ := time.Parse(dateFmt, "2014-08-11T23:15:22")
	t0, _ := time.Parse(dateFmt, "2014-08-11T00:00:00")
	t1, _ := time.Parse(dateFmt, "2014-08-12T00:00:00")
	t2, _ := time.Parse(dateFmt, "2014-08-13T00:00:00")

	tcalc := calculateNextRollTime(traw, RollDaily)
	if !tcalc.Equal(t1) {
		t.Errorf("Error getting t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculateNextRollTime(tcalc, RollDaily)
	if !tcalc.Equal(t2) {
		t.Errorf("Error getting t2. Expected %v but got %v", t2, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollDaily)
	if !tcalc.Equal(t1) {
		t.Errorf("Error rolling back to t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollDaily)
	if !tcalc.Equal(t0) {
		t.Errorf("Error getting t0. Expected %v but got %v", t0, tcalc)
	}
}

func TestWeeklyOperations(t *testing.T) {
	traw, _ := time.Parse(dateFmt, "2014-08-11T23:15:22") // a Monday
	t0, _ := time.Parse(dateFmt, "2014-08-10T00:00:00")
	t1, _ := time.Parse(dateFmt, "2014-08-17T00:00:00")
	t2, _ := time.Parse(dateFmt, "2014-08-24T00:00:00")

	tcalc := calculateNextRollTime(traw, RollWeekly)
	if !tcalc.Equal(t1) {
		t.Errorf("Error getting t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculateNextRollTime(tcalc, RollWeekly)
	if !tcalc.Equal(t2) {
		t.Errorf("Error getting t2. Expected %v but got %v", t2, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollWeekly)
	if !tcalc.Equal(t1) {
		t.Errorf("Error rolling back to t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollWeekly)
	if !tcalc.Equal(t0) {
		t.Errorf("Error getting t0. Expected %v but got %v", t0, tcalc)
	}
}

func TestHourlyHandlesDSTStart(t *testing.T) {

	// XXX: deliberate break. Testing:
	tt := time.Now()
	tz, _ := tt.Zone()
	if tz != "XXX" {
		t.Errorf("expected some BS but got %s", tz)
	}

	t0, _ := time.Parse(time.RFC822, "09 Mar 14 00:15 PST")
	t1, _ := time.Parse(time.RFC822, "09 Mar 14 01:00 PST")
	t3, _ := time.Parse(time.RFC822, "09 Mar 14 03:00 PDT")

	tnext := calculateNextRollTime(t0, RollHourly)
	if !tnext.Equal(t1) {
		t.Errorf("Error getting t1. Expected %v but got %v", t1, tnext)
	}

	tnext = calculateNextRollTime(tnext, RollHourly)
	if !tnext.Equal(t3) {
		t.Errorf("Error getting t3. Expected %v but got %v", t3, tnext)
	}

	tprev := calculatePreviousRollTime(tnext, RollHourly)
	if !tprev.Equal(t1) {
		t.Errorf("Error rolling back to t1. Expected %v but got %v", t1, tprev)
	}
}

func TestHourlyHandlesDSTEnd(t *testing.T) {
	t0, _ := time.Parse(time.RFC822, "02 Nov 14 00:15 PDT")
	t1, _ := time.Parse(time.RFC822, "02 Nov 14 01:00 PDT")
	t2, _ := time.Parse(time.RFC822, "02 Nov 14 02:00 PST")

	tnext := calculateNextRollTime(t0, RollHourly)
	if !tnext.Equal(t1) {
		t.Errorf("Error getting t1. Expected %v but got %v", t1, tnext)
	}

	tnext = calculateNextRollTime(tnext, RollHourly)
	if !tnext.Equal(t2) {
		t.Errorf("Error getting t2. Expected %v but got %v", t2, tnext)
	}

	tprev := calculatePreviousRollTime(tnext, RollHourly)
	if !tprev.Equal(t1) {
		t.Errorf("Error rolling back to t1. Expected %v but got %v", t1, tprev)
	}
}

func TestDailyHandlesDST(t *testing.T) {
	// DST start
	traw, _ := time.Parse(time.RFC822, "09 Mar 14 12:15 PDT")
	t0, _ := time.Parse(time.RFC822, "09 Mar 14 00:00 PST")
	t1, _ := time.Parse(time.RFC822, "10 Mar 14 00:00 PDT")

	tcalc := calculateNextRollTime(traw, RollDaily)
	if !tcalc.Equal(t1) {
		t.Errorf("Error getting t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollDaily)
	if !tcalc.Equal(t0) {
		t.Errorf("Error rolling back to t0. Expected %v but got %v", t0, tcalc)
	}

	// DST end
	traw, _ = time.Parse(time.RFC822, "02 Nov 14 12:15 PST")
	t0, _ = time.Parse(time.RFC822, "02 Nov 14 00:00 PDT")
	t1, _ = time.Parse(time.RFC822, "03 Nov 14 00:00 PST")

	tcalc = calculateNextRollTime(traw, RollDaily)
	if !tcalc.Equal(t1) {
		t.Errorf("Error getting t1. Expected %v but got %v", t1, tcalc)
	}

	tcalc = calculatePreviousRollTime(tcalc, RollDaily)
	if !tcalc.Equal(t0) {
		t.Errorf("Error rolling back to t0. Expected %v but got %v", t0, tcalc)
	}
}
<|endoftext|>


<commit_before>package firebase

import (
	"errors"
	"fmt"
	"time"

	"github.com/SermoDigital/jose/crypto"
	"github.com/SermoDigital/jose/jws"
	"github.com/SermoDigital/jose/jwt"
)

// clientCertURL is the URL containing the public keys for the Google certs
// (whose private keys are used to sign Firebase Auth ID Tokens).
const clientCertURL = "https://www.googleapis.com/robot/v1/metadata/x509/securetoken@system.gserviceaccount.com"

// defaultAcceptableExpSkew is the default expiry leeway.
const defaultAcceptableExpSkew = 300 * time.Second

func verify(projectID, tokenString string) (*Token, error) {
	decodedJWT, err := jws.ParseJWT([]byte(tokenString))
	if err != nil {
		return nil, err
	}
	decodedJWS, ok := decodedJWT.(jws.JWS)
	if !ok {
		return nil, errors.New("Firebase Auth ID Token cannot be decoded")
	}

	keys := func(j jws.JWS) ([]interface{}, error) {
		certs := &Certificates{URL: clientCertURL}
		kid, ok := j.Header().Get("kid").(string)
		if !ok {
			return nil, errors.New("Firebase Auth ID Token has no 'kid' claim")
		}
		cert, err := certs.Cert(kid)
		if err != nil {
			return nil, err
		}
		return []interface{}{cert.PublicKey}, nil
	}

	err = decodedJWS.VerifyCallback(keys,
		[]crypto.SigningMethod{crypto.SigningMethodRS256},
		&jws.SigningOpts{Number: 1, Indices: []int{0}})
	if err != nil {
		return nil, err
	}

	ks, _ := keys(decodedJWS)
	key := ks[0]
	if err := decodedJWT.Validate(key, crypto.SigningMethodRS256, validator(projectID)); err != nil {
		return nil, err
	}

	return &Token{delegate: decodedJWT}, nil
}

func validator(projectID string) *jwt.Validator {
	v := &jwt.Validator{}
	v.EXP = defaultAcceptableExpSkew
	v.SetAudience(projectID)
	v.SetIssuer(fmt.Sprintf("https://securetoken.google.com/%s", projectID))
	v.Fn = func(claims jwt.Claims) error {
		subject, ok := claims.Subject()
		if !ok || len(subject) == 0 || len(subject) > 128 {
			return jwt.ErrInvalidSUBClaim
		}
		return nil
	}
	return v
}
<commit_msg>lookup kid from protected headers of a firebase token<commit_after>package firebase

import (
	"errors"
	"fmt"
	"time"

	"github.com/SermoDigital/jose/crypto"
	"github.com/SermoDigital/jose/jws"
	"github.com/SermoDigital/jose/jwt"
)

// clientCertURL is the URL containing the public keys for the Google certs
// (whose private keys are used to sign Firebase Auth ID Tokens).
const clientCertURL = "https://www.googleapis.com/robot/v1/metadata/x509/securetoken@system.gserviceaccount.com"

// defaultAcceptableExpSkew is the default expiry leeway.
const defaultAcceptableExpSkew = 300 * time.Second

func verify(projectID, tokenString string) (*Token, error) {
	decodedJWT, err := jws.ParseJWT([]byte(tokenString))
	if err != nil {
		return nil, err
	}
	decodedJWS, ok := decodedJWT.(jws.JWS)
	if !ok {
		return nil, errors.New("Firebase Auth ID Token cannot be decoded")
	}

	keys := func(j jws.JWS) ([]interface{}, error) {
		certs := &Certificates{URL: clientCertURL}
		kid, ok := j.Protected().Get("kid").(string)
		if !ok {
			return nil, errors.New("Firebase Auth ID Token has no 'kid' claim")
		}
		cert, err := certs.Cert(kid)
		if err != nil {
			return nil, err
		}
		return []interface{}{cert.PublicKey}, nil
	}

	err = decodedJWS.VerifyCallback(keys,
		[]crypto.SigningMethod{crypto.SigningMethodRS256},
		&jws.SigningOpts{Number: 1, Indices: []int{0}})
	if err != nil {
		return nil, err
	}

	ks, _ := keys(decodedJWS)
	key := ks[0]
	if err := decodedJWT.Validate(key, crypto.SigningMethodRS256, validator(projectID)); err != nil {
		return nil, err
	}

	return &Token{delegate: decodedJWT}, nil
}

func validator(projectID string) *jwt.Validator {
	v := &jwt.Validator{}
	v.EXP = defaultAcceptableExpSkew
	v.SetAudience(projectID)
	v.SetIssuer(fmt.Sprintf("https://securetoken.google.com/%s", projectID))
	v.Fn = func(claims jwt.Claims) error {
		subject, ok := claims.Subject()
		if !ok || len(subject) == 0 || len(subject) > 128 {
			return jwt.ErrInvalidSUBClaim
		}
		return nil
	}
	return v
}
<|endoftext|>


<commit_before>package specs

/*
	Helper functions for client repos.
*/

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"reflect"
	"strings"

	"github.com/pmezard/go-difflib/difflib"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/sklog"
)

var (
	// Flags.
	test = flag.Bool("test", false, "Run in test mode: verify that the output hasn't changed.")
)

// GetCheckoutRoot returns the path of the root of the checkout.
func GetCheckoutRoot() (string, error) {
	cwd, err := os.Getwd()
	if err != nil {
		return "", err
	}
	for {
		if _, err := os.Stat(cwd); err != nil {
			return "", err
		}
		// TODO(borenet): Should we verify that this is the
		// correct checkout and not something else?

		// Check for infra/bots dir.
		s, err := os.Stat(path.Join(cwd, "infra", "bots"))
		if err == nil && s.IsDir() {
			return cwd, nil
		}
		// Check for .git dir.
		s, err = os.Stat(path.Join(cwd, ".git"))
		if err == nil && s.IsDir() {
			return cwd, nil
		}

		// Stop if we're at the filesystem root.
		if cwd == string(filepath.Separator) {
			return "", fmt.Errorf("Unable to find repository root.")
		}
		cwd = filepath.Clean(path.Join(cwd, ".."))
	}
}

// TasksCfgBuilder is a helper struct used for building a TasksCfg.
type TasksCfgBuilder struct {
	assetsDir    string
	cfg          *TasksCfg
	cipdPackages map[string]*CipdPackage
	root         string
}

// NewTasksCfgBuilder returns a TasksCfgBuilder instance.
func NewTasksCfgBuilder() (*TasksCfgBuilder, error) {
	common.Init()

	// Create the config.
	cfg := &TasksCfg{
		Jobs:  map[string]*JobSpec{},
		Tasks: map[string]*TaskSpec{},
	}

	root, err := GetCheckoutRoot()
	if err != nil {
		return nil, err
	}

	return &TasksCfgBuilder{
		cfg:          cfg,
		cipdPackages: map[string]*CipdPackage{},
		root:         root,
	}, nil
}

// MustNewTasksCfgBuilder returns a TasksCfgBuilder instance. Panics on error.
func MustNewTasksCfgBuilder() *TasksCfgBuilder {
	b, err := NewTasksCfgBuilder()
	if err != nil {
		sklog.Fatal(err)
	}
	return b
}

// CheckoutRoot returns the path to the root of the client checkout.
func (b *TasksCfgBuilder) CheckoutRoot() string {
	return b.root
}

// SetAssetsDir sets the directory path used for assets.
func (b *TasksCfgBuilder) SetAssetsDir(assetsDir string) {
	b.assetsDir = assetsDir
}

// AddTask adds a TaskSpec to the TasksCfgBuilder. Returns an error if the
// config already contains a Task with the same name and a different
// implementation.
func (b *TasksCfgBuilder) AddTask(name string, t *TaskSpec) error {
	if old, ok := b.cfg.Tasks[name]; ok {
		if !reflect.DeepEqual(old, t) {
			return fmt.Errorf("Config already contains a Task named %q with a different implementation!\nHave:\n%v\n\nGot:\n%v", name, old, t)
		}
		return nil
	}
	b.cfg.Tasks[name] = t
	return nil
}

// MustAddTask adds a TaskSpec to the TasksCfgBuilder and panics on failure.
func (b *TasksCfgBuilder) MustAddTask(name string, t *TaskSpec) {
	if err := b.AddTask(name, t); err != nil {
		sklog.Fatal(err)
	}
}

// AddJob adds a JobSpec to the TasksCfgBuilder.
func (b *TasksCfgBuilder) AddJob(name string, j *JobSpec) error {
	if _, ok := b.cfg.Jobs[name]; ok {
		return fmt.Errorf("Config already contains a Job named %q", name)
	}
	b.cfg.Jobs[name] = j
	return nil
}

// MustAddJob adds a JobSpec to the TasksCfgBuilder and panics on failure.
func (b *TasksCfgBuilder) MustAddJob(name string, j *JobSpec) {
	if err := b.AddJob(name, j); err != nil {
		sklog.Fatal(err)
	}
}

// GetCipdPackageFromAsset reads the version information for the given asset
// and returns a CipdPackage instance.
func (b *TasksCfgBuilder) GetCipdPackageFromAsset(assetName string) (*CipdPackage, error) {
	if pkg, ok := b.cipdPackages[assetName]; ok {
		return pkg, nil
	}
	assetsDir := b.assetsDir
	if assetsDir == "" {
		assetsDir = path.Join(b.root, "infra", "bots", "assets")
	}
	versionFile := path.Join(assetsDir, assetName, "VERSION")
	contents, err := ioutil.ReadFile(versionFile)
	if err != nil {
		return nil, err
	}
	version := strings.TrimSpace(string(contents))
	pkg := &CipdPackage{
		Name:    fmt.Sprintf("skia/bots/%s", assetName),
		Path:    assetName,
		Version: fmt.Sprintf("version:%s", version),
	}
	if assetName == "win_toolchain" {
		pkg.Path = "t" // Workaround for path length limit on Windows.
	}
	b.cipdPackages[assetName] = pkg
	return pkg, nil
}

// MustGetCipdPackageFromAsset reads the version information for the given asset
// and returns a CipdPackage instance. Panics on failure.
func (b *TasksCfgBuilder) MustGetCipdPackageFromAsset(assetName string) *CipdPackage {
	pkg, err := b.GetCipdPackageFromAsset(assetName)
	if err != nil {
		sklog.Fatal(err)
	}
	return pkg
}

// Finish validates and writes out the TasksCfg, or, if the --test flag is
// provided, verifies that the contents have not changed.
func (b *TasksCfgBuilder) Finish() error {
	// Validate the config.
	if err := b.cfg.Validate(); err != nil {
		return err
	}

	enc, err := EncodeTasksCfg(b.cfg)
	if err != nil {
		return err
	}

	// Write the tasks.json file.
	outFile := path.Join(b.root, TASKS_CFG_FILE)
	if *test {
		// Don't write the file; read it and compare.
		expect, err := ioutil.ReadFile(outFile)
		if err != nil {
			return err
		}
		if !bytes.Equal(expect, enc) {
			diff, err := difflib.GetUnifiedDiffString(difflib.UnifiedDiff{
				A:        difflib.SplitLines(string(expect)),
				B:        difflib.SplitLines(string(enc)),
				FromFile: "Expected",
				ToFile:   "Actual",
				Context:  3,
				Eol:      "\n",
			})
			if err != nil {
				diff = fmt.Sprintf("<failed to obtain diff: %s>", err)
			}
			return fmt.Errorf(`Expected no changes, but changes were found:

%s

You may need to run:

	$ go run infra/bots/gen_tasks.go

`, diff)
		}
	} else {
		if err := ioutil.WriteFile(outFile, enc, os.ModePerm); err != nil {
			return err
		}
	}
	return nil
}

// MustFinish validates and writes out the TasksCfg, or, if the --test flag is
// provided, verifies that the contents have not changed. Panics on failure.
func (b *TasksCfgBuilder) MustFinish() {
	if err := b.Finish(); err != nil {
		sklog.Error(err)
		os.Exit(1)
	}
}
<commit_msg>Remove win_toolchain name substitution<commit_after>package specs

/*
	Helper functions for client repos.
*/

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"reflect"
	"strings"

	"github.com/pmezard/go-difflib/difflib"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/sklog"
)

var (
	// Flags.
	test = flag.Bool("test", false, "Run in test mode: verify that the output hasn't changed.")
)

// GetCheckoutRoot returns the path of the root of the checkout.
func GetCheckoutRoot() (string, error) {
	cwd, err := os.Getwd()
	if err != nil {
		return "", err
	}
	for {
		if _, err := os.Stat(cwd); err != nil {
			return "", err
		}
		// TODO(borenet): Should we verify that this is the
		// correct checkout and not something else?

		// Check for infra/bots dir.
		s, err := os.Stat(path.Join(cwd, "infra", "bots"))
		if err == nil && s.IsDir() {
			return cwd, nil
		}
		// Check for .git dir.
		s, err = os.Stat(path.Join(cwd, ".git"))
		if err == nil && s.IsDir() {
			return cwd, nil
		}

		// Stop if we're at the filesystem root.
		if cwd == string(filepath.Separator) {
			return "", fmt.Errorf("Unable to find repository root.")
		}
		cwd = filepath.Clean(path.Join(cwd, ".."))
	}
}

// TasksCfgBuilder is a helper struct used for building a TasksCfg.
type TasksCfgBuilder struct {
	assetsDir    string
	cfg          *TasksCfg
	cipdPackages map[string]*CipdPackage
	root         string
}

// NewTasksCfgBuilder returns a TasksCfgBuilder instance.
func NewTasksCfgBuilder() (*TasksCfgBuilder, error) {
	common.Init()

	// Create the config.
	cfg := &TasksCfg{
		Jobs:  map[string]*JobSpec{},
		Tasks: map[string]*TaskSpec{},
	}

	root, err := GetCheckoutRoot()
	if err != nil {
		return nil, err
	}

	return &TasksCfgBuilder{
		cfg:          cfg,
		cipdPackages: map[string]*CipdPackage{},
		root:         root,
	}, nil
}

// MustNewTasksCfgBuilder returns a TasksCfgBuilder instance. Panics on error.
func MustNewTasksCfgBuilder() *TasksCfgBuilder {
	b, err := NewTasksCfgBuilder()
	if err != nil {
		sklog.Fatal(err)
	}
	return b
}

// CheckoutRoot returns the path to the root of the client checkout.
func (b *TasksCfgBuilder) CheckoutRoot() string {
	return b.root
}

// SetAssetsDir sets the directory path used for assets.
func (b *TasksCfgBuilder) SetAssetsDir(assetsDir string) {
	b.assetsDir = assetsDir
}

// AddTask adds a TaskSpec to the TasksCfgBuilder. Returns an error if the
// config already contains a Task with the same name and a different
// implementation.
func (b *TasksCfgBuilder) AddTask(name string, t *TaskSpec) error {
	if old, ok := b.cfg.Tasks[name]; ok {
		if !reflect.DeepEqual(old, t) {
			return fmt.Errorf("Config already contains a Task named %q with a different implementation!\nHave:\n%v\n\nGot:\n%v", name, old, t)
		}
		return nil
	}
	b.cfg.Tasks[name] = t
	return nil
}

// MustAddTask adds a TaskSpec to the TasksCfgBuilder and panics on failure.
func (b *TasksCfgBuilder) MustAddTask(name string, t *TaskSpec) {
	if err := b.AddTask(name, t); err != nil {
		sklog.Fatal(err)
	}
}

// AddJob adds a JobSpec to the TasksCfgBuilder.
func (b *TasksCfgBuilder) AddJob(name string, j *JobSpec) error {
	if _, ok := b.cfg.Jobs[name]; ok {
		return fmt.Errorf("Config already contains a Job named %q", name)
	}
	b.cfg.Jobs[name] = j
	return nil
}

// MustAddJob adds a JobSpec to the TasksCfgBuilder and panics on failure.
func (b *TasksCfgBuilder) MustAddJob(name string, j *JobSpec) {
	if err := b.AddJob(name, j); err != nil {
		sklog.Fatal(err)
	}
}

// GetCipdPackageFromAsset reads the version information for the given asset
// and returns a CipdPackage instance.
func (b *TasksCfgBuilder) GetCipdPackageFromAsset(assetName string) (*CipdPackage, error) {
	if pkg, ok := b.cipdPackages[assetName]; ok {
		return pkg, nil
	}
	assetsDir := b.assetsDir
	if assetsDir == "" {
		assetsDir = path.Join(b.root, "infra", "bots", "assets")
	}
	versionFile := path.Join(assetsDir, assetName, "VERSION")
	contents, err := ioutil.ReadFile(versionFile)
	if err != nil {
		return nil, err
	}
	version := strings.TrimSpace(string(contents))
	pkg := &CipdPackage{
		Name:    fmt.Sprintf("skia/bots/%s", assetName),
		Path:    assetName,
		Version: fmt.Sprintf("version:%s", version),
	}
	b.cipdPackages[assetName] = pkg
	return pkg, nil
}

// MustGetCipdPackageFromAsset reads the version information for the given asset
// and returns a CipdPackage instance. Panics on failure.
func (b *TasksCfgBuilder) MustGetCipdPackageFromAsset(assetName string) *CipdPackage {
	pkg, err := b.GetCipdPackageFromAsset(assetName)
	if err != nil {
		sklog.Fatal(err)
	}
	return pkg
}

// Finish validates and writes out the TasksCfg, or, if the --test flag is
// provided, verifies that the contents have not changed.
func (b *TasksCfgBuilder) Finish() error {
	// Validate the config.
	if err := b.cfg.Validate(); err != nil {
		return err
	}

	enc, err := EncodeTasksCfg(b.cfg)
	if err != nil {
		return err
	}

	// Write the tasks.json file.
	outFile := path.Join(b.root, TASKS_CFG_FILE)
	if *test {
		// Don't write the file; read it and compare.
		expect, err := ioutil.ReadFile(outFile)
		if err != nil {
			return err
		}
		if !bytes.Equal(expect, enc) {
			diff, err := difflib.GetUnifiedDiffString(difflib.UnifiedDiff{
				A:        difflib.SplitLines(string(expect)),
				B:        difflib.SplitLines(string(enc)),
				FromFile: "Expected",
				ToFile:   "Actual",
				Context:  3,
				Eol:      "\n",
			})
			if err != nil {
				diff = fmt.Sprintf("<failed to obtain diff: %s>", err)
			}
			return fmt.Errorf(`Expected no changes, but changes were found:

%s

You may need to run:

	$ go run infra/bots/gen_tasks.go

`, diff)
		}
	} else {
		if err := ioutil.WriteFile(outFile, enc, os.ModePerm); err != nil {
			return err
		}
	}
	return nil
}

// MustFinish validates and writes out the TasksCfg, or, if the --test flag is
// provided, verifies that the contents have not changed. Panics on failure.
func (b *TasksCfgBuilder) MustFinish() {
	if err := b.Finish(); err != nil {
		sklog.Error(err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package redis

import (
	"fmt"
	"testing"
	"time"
)

var (
	network = "tcp"
	address = "192.168.84.250:6379"
)

func dial() (*Redis, error) {
	return DialTimeout(network, address, 0, "", 5*time.Second, 5)
}

func TestDial(t *testing.T) {
	_, err := dial()
	if err != nil {
		t.Error(err)
	}
}

func TestDialFail(t *testing.T) {
	_, err := DialTimeout(network, address+"0", 0, "", 5*time.Second, 5)
	if err == nil {
		t.Error(err)
	}
}

func TestDiaURL(t *testing.T) {
	rawurl := fmt.Sprintf("redis://%s/1?size=5&timeout=10s", address)
	r, err := DialURL(rawurl)
	if err != nil {
		t.Fatal(err)
	}
	if r.db != 1 || r.size != 5 || r.timeout != 10*time.Second {
		t.Fail()
	}
}

func TestDialURLFail(t *testing.T) {
	rawurl := fmt.Sprintf("redis://tester:password@%s/1", address)
	_, err := DialURL(rawurl)
	if err == nil {
		t.Fail()
	}
}

func TestAuth(t *testing.T) {
	r, _ := dial()
	if err := r.Auth("password"); err == nil {
		t.Fail()
	}
}

func TestClientList(t *testing.T) {
	r, _ := dial()
	_, err := r.ClientList()
	if err != nil {
		t.Error(err)
	}
}

func TestAppend(t *testing.T) {
	r, _ := dial()
	r.Del("key")
	n, err := r.Append("key", "value")
	if err != nil {
		t.Error(err)
	}
	if n != 5 {
		t.Fail()
	}
	n, err = r.Append("key", "value")
	if err != nil {
		t.Error(err)
	}
	if n != 10 {
		t.Fail()
	}
	r.Del("key")
	r.LPush("key", "value")
	if _, err := r.Append("key", "value"); err == nil {
		t.Error(err)
	}
}

func TestBLPop(t *testing.T) {
	r, _ := dial()
	r.Del("key")
	result, err := r.BLPop([]string{"key"}, 1)
	if err != nil {
		t.Error(err)
	}
	if len(result) != 0 {
		t.Fail()
	}
	r.LPush("key", "value")
	result, err = r.BLPop([]string{"key"}, 0)
	if err != nil {
		t.Error(err)
	}
	if len(result) == 0 {
		t.Fail()
	}
	if result[0] != "key" || result[1] != "value" {
		t.Fail()
	}
}

func TestDBSize(t *testing.T) {
	r, _ := dial()
	r.FlushDB()
	n, err := r.DBSize()
	if err != nil {
		t.Error(err)
	}
	if n != 0 {
		t.Fail()
	}
}

func TestEval(t *testing.T) {
	r, _ := dial()
	rp, err := r.Eval("return {KEYS[1], KEYS[2], ARGV[1], ARGV[2]}", []string{"key1", "key2"}, []string{"arg1", "arg2"})
	if err != nil {
		t.Error(err)
	}
	l, err := rp.ListValue()
	if err != nil {
		t.Error(err)
	}
	if l[0] != "key1" || l[3] != "arg2" {
		t.Fail()
	}
	rp, err = r.Eval("return redis.call('set','foo','bar')", nil, nil)
	if err != nil {
		t.Error(err)
	}
	if err := rp.OKValue(); err != nil {
		t.Error(err)
	}
	rp, err = r.Eval("return 10", nil, nil)
	if err != nil {
		t.Error(err)
	}
	n, err := rp.IntegerValue()
	if err != nil {
		t.Error(err)
	}
	if n != 10 {
		t.Fail()
	}
	rp, err = r.Eval("return {1,2,{3,'Hello World!'}}", nil, nil)
	if err != nil {
		t.Error(err)
	}
	if len(rp.Multi) != 3 {
		t.Fail()
	}
	if rp.Multi[2].Multi[0].Integer != 3 {
		t.Fail()
	}
	if s, err := rp.Multi[2].Multi[1].StringValue(); err != nil || s != "Hello World!" {
		t.Fail()
	}
}

func TestExists(t *testing.T) {
	r, _ := dial()
	r.Del("key")
	b, err := r.Exists("key")
	if err != nil {
		t.Error(err)
	}
	if b {
		t.Fail()
	}
}

func TestGet(t *testing.T) {
	r, _ := dial()
	r.Del("key")
	data, err := r.Get("key")
	if err != nil {
		t.Error(err)
	}
	if data != nil {
		t.Fail()
	}
}

func TestHGetAll(t *testing.T) {
	r, _ := dial()
	r.Del("key")
	pairs := map[string]string{"name": "foo", "attr": "bar"}
	if err := r.HMSet("key", pairs); err != nil {
		t.Error(err)
	}
	data, err := r.HGetAll("key")
	if err != nil {
		t.Error(err)
	}
	if data["name"] != "foo" {
		t.Fail()
	}
}

func TestHMGet(t *testing.T) {
	r, _ := dial()
	r.Del("key")
	r.HSet("key", "field", "value")
	data, err := r.HMGet("key", "field", "nofield")
	if err != nil {
		t.Error(err)
	}
	if string(data[0]) != "value" {
		t.Fail()
	}
	if data[1] != nil {
		t.Fail()
	}
}

func TestKeys(t *testing.T) {
	r, _ := dial()
	r.FlushDB()
	keys, err := r.Keys("*")
	if err != nil {
		t.Error(err)
	}
	if len(keys) != 0 {
		t.Fail()
	}
	r.Set("key", "value", 0, 0, false, false)
	keys, err = r.Keys("*")
	if err != nil {
		t.Error(err)
	}
	if len(keys) != 1 || keys[0] != "key" {
		t.Fail()
	}
}

func TestTransaction(t *testing.T) {
	r, _ := dial()
	transaction, err := r.Transaction()
	if err != nil {
		t.Error(err)
	}
	defer transaction.Close()
	transaction.Command("DEL", "key")
	transaction.Command("SET", "key", 1)
	transaction.Command("INCR", "key")
	transaction.Command("GET", "key")
	result, err := transaction.Exec()
	if err != nil {
		t.Error(err)
	}
	if len(result) != 4 {
		t.Fail()
	}
	if s, err := result[3].StringValue(); err != nil || s != "2" {
		t.Fail()
	}
}
<commit_msg>test publish<commit_after>package redis

import (
	"fmt"
	"testing"
	"time"
)

var (
	network = "tcp"
	address = "192.168.84.250:6379"
)

func dial() (*Redis, error) {
	return DialTimeout(network, address, 0, "", 5*time.Second, 5)
}

func TestDial(t *testing.T) {
	_, err := dial()
	if err != nil {
		t.Error(err)
	}
}

func TestDialFail(t *testing.T) {
	_, err := DialTimeout(network, address+"0", 0, "", 5*time.Second, 5)
	if err == nil {
		t.Error(err)
	}
}

func TestDiaURL(t *testing.T) {
	rawurl := fmt.Sprintf("redis://%s/1?size=5&timeout=10s", address)
	r, err := DialURL(rawurl)
	if err != nil {
		t.Fatal(err)
	}
	if r.db != 1 || r.size != 5 || r.timeout != 10*time.Second {
		t.Fail()
	}
}

func TestDialURLFail(t *testing.T) {
	rawurl := fmt.Sprintf("redis://tester:password@%s/1", address)
	_, err := DialURL(rawurl)
	if err == nil {
		t.Fail()
	}
}

func TestAuth(t *testing.T) {
	r, _ := dial()
	if err := r.Auth("password"); err == nil {
		t.Fail()
	}
}

func TestClientList(t *testing.T) {
	r, _ := dial()
	_, err := r.ClientList()
	if err != nil {
		t.Error(err)
	}
}

func TestAppend(t *testing.T) {
	r, _ := dial()
	r.Del("key")
	n, err := r.Append("key", "value")
	if err != nil {
		t.Error(err)
	}
	if n != 5 {
		t.Fail()
	}
	n, err = r.Append("key", "value")
	if err != nil {
		t.Error(err)
	}
	if n != 10 {
		t.Fail()
	}
	r.Del("key")
	r.LPush("key", "value")
	if _, err := r.Append("key", "value"); err == nil {
		t.Error(err)
	}
}

func TestBLPop(t *testing.T) {
	r, _ := dial()
	r.Del("key")
	result, err := r.BLPop([]string{"key"}, 1)
	if err != nil {
		t.Error(err)
	}
	if len(result) != 0 {
		t.Fail()
	}
	r.LPush("key", "value")
	result, err = r.BLPop([]string{"key"}, 0)
	if err != nil {
		t.Error(err)
	}
	if len(result) == 0 {
		t.Fail()
	}
	if result[0] != "key" || result[1] != "value" {
		t.Fail()
	}
}

func TestDBSize(t *testing.T) {
	r, _ := dial()
	r.FlushDB()
	n, err := r.DBSize()
	if err != nil {
		t.Error(err)
	}
	if n != 0 {
		t.Fail()
	}
}

func TestEval(t *testing.T) {
	r, _ := dial()
	rp, err := r.Eval("return {KEYS[1], KEYS[2], ARGV[1], ARGV[2]}", []string{"key1", "key2"}, []string{"arg1", "arg2"})
	if err != nil {
		t.Error(err)
	}
	l, err := rp.ListValue()
	if err != nil {
		t.Error(err)
	}
	if l[0] != "key1" || l[3] != "arg2" {
		t.Fail()
	}
	rp, err = r.Eval("return redis.call('set','foo','bar')", nil, nil)
	if err != nil {
		t.Error(err)
	}
	if err := rp.OKValue(); err != nil {
		t.Error(err)
	}
	rp, err = r.Eval("return 10", nil, nil)
	if err != nil {
		t.Error(err)
	}
	n, err := rp.IntegerValue()
	if err != nil {
		t.Error(err)
	}
	if n != 10 {
		t.Fail()
	}
	rp, err = r.Eval("return {1,2,{3,'Hello World!'}}", nil, nil)
	if err != nil {
		t.Error(err)
	}
	if len(rp.Multi) != 3 {
		t.Fail()
	}
	if rp.Multi[2].Multi[0].Integer != 3 {
		t.Fail()
	}
	if s, err := rp.Multi[2].Multi[1].StringValue(); err != nil || s != "Hello World!" {
		t.Fail()
	}
}

func TestExists(t *testing.T) {
	r, _ := dial()
	r.Del("key")
	b, err := r.Exists("key")
	if err != nil {
		t.Error(err)
	}
	if b {
		t.Fail()
	}
}

func TestGet(t *testing.T) {
	r, _ := dial()
	r.Del("key")
	data, err := r.Get("key")
	if err != nil {
		t.Error(err)
	}
	if data != nil {
		t.Fail()
	}
}

func TestHGetAll(t *testing.T) {
	r, _ := dial()
	r.Del("key")
	pairs := map[string]string{"name": "foo", "attr": "bar"}
	if err := r.HMSet("key", pairs); err != nil {
		t.Error(err)
	}
	data, err := r.HGetAll("key")
	if err != nil {
		t.Error(err)
	}
	if data["name"] != "foo" {
		t.Fail()
	}
}

func TestHMGet(t *testing.T) {
	r, _ := dial()
	r.Del("key")
	r.HSet("key", "field", "value")
	data, err := r.HMGet("key", "field", "nofield")
	if err != nil {
		t.Error(err)
	}
	if string(data[0]) != "value" {
		t.Fail()
	}
	if data[1] != nil {
		t.Fail()
	}
}

func TestKeys(t *testing.T) {
	r, _ := dial()
	r.FlushDB()
	keys, err := r.Keys("*")
	if err != nil {
		t.Error(err)
	}
	if len(keys) != 0 {
		t.Fail()
	}
	r.Set("key", "value", 0, 0, false, false)
	keys, err = r.Keys("*")
	if err != nil {
		t.Error(err)
	}
	if len(keys) != 1 || keys[0] != "key" {
		t.Fail()
	}
}

func TestTransaction(t *testing.T) {
	r, _ := dial()
	transaction, err := r.Transaction()
	if err != nil {
		t.Error(err)
	}
	defer transaction.Close()
	transaction.Command("DEL", "key")
	transaction.Command("SET", "key", 1)
	transaction.Command("INCR", "key")
	transaction.Command("GET", "key")
	result, err := transaction.Exec()
	if err != nil {
		t.Error(err)
	}
	if len(result) != 4 {
		t.Fail()
	}
	if s, err := result[3].StringValue(); err != nil || s != "2" {
		t.Fail()
	}
}

func TestPublish(t *testing.T) {
	r, _ := dial()
	if _, err := r.Publish("key", "value"); err != nil {
		t.Error(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"os/exec"
	"io"
	"os"
	"io/ioutil"
	"strings"
	"syscall"
	"fmt"
	"bytes"
    "crypto/rand"
    "encoding/base64"
    "encoding/gob"
    "log"
    "net"

	"github.com/flynn/go-discover/discover"
    "github.com/flynn/lorne/types"
    "github.com/flynn/sampi/client"
    "github.com/flynn/sampi/types"
    "github.com/titanous/go-dockerclient"
)

func main() {
	root := "/var/lib/demo/apps"
	hostname := shell("curl -s icanhazip.com")

	client, err := discover.NewClient()
	if err != nil {
		panic(err)
	}
	set, _ := client.Services("shelf")
	addrs := set.OnlineAddrs()
	if len(addrs) < 1 {
		panic("Shelf is not discoverable")
	}
	shelfHost := addrs[0]

	app := os.Args[2]
	os.MkdirAll(root + "/" + app, 0755)

	fmt.Printf("-----> Building %s on %s ...\n", app, hostname)

	scheduleAndAttach(docker.Config{
        Image:        "flynn/slugbuilder",
        Cmd:          []string{"http://"+shelfHost+"/"+app+".tgz",},
        Tty:          false,
        AttachStdin:  true,
        AttachStdout: true,
        AttachStderr: true,
        OpenStdin:    true,
        StdinOnce:    true,
	})

	/*fmt.Printf("-----> Deploying %s ...\n", app)
	if _, err := os.Stat(root + "/" + app + "/CONTAINER"); err == nil {
    	oldid := readFile(root + "/" + app + "/CONTAINER")
    	shell("docker kill " + oldid)
	}

	id := shell("docker run -d -p 5000 -e PORT=5000 -e SLUG_URL=http://"+shelfHost+"/"+app+".tgz flynn/slugrunner start web")
	writeFile(root + "/" + app + "/CONTAINER", id)
	port := shell("docker port "+id+" 5000 | sed 's/0.0.0.0://'")
	writeFile(root + "/" + app + "/PORT", port)
	writeFile(root + "/" + app + "/URL", "http://"+hostname+":"+port)

	fmt.Printf("=====> Application deployed:\n")
	fmt.Printf("       %s\n", readFile(root + "/" + app + "/URL"))*/
	fmt.Println("")

}

func readFile(filename string) string {
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		panic(err)
	}
	return string(data)
}

func writeFile(filename, data string) {
	err := ioutil.WriteFile(filename, []byte(data), 0644)
	if err != nil {
		panic(err)
	}
}

func shell(cmdline string) string {
	out, err := exec.Command("bash", "-c", cmdline).Output()
	if err != nil {
		panic(err)
	}
	return strings.Trim(string(out), " \n")
}

func attachCmd(cmd *exec.Cmd, stdout, stderr io.Writer, stdin io.Reader) chan error {
	errCh := make(chan error)

	stdinIn, err := cmd.StdinPipe()
	if err != nil {
		panic(err)
	}
	stdoutOut, err := cmd.StdoutPipe()
	if err != nil {
		panic(err)
	}
	stderrOut, err := cmd.StderrPipe()
	if err != nil {
		panic(err)
	}

	go func() {
		_, e := io.Copy(stdinIn, stdin)
		errCh <- e
	}()
	go func() {
		_, e := io.Copy(stdout, stdoutOut)
		errCh <- e
	}()
	go func() {
		_, e := io.Copy(stderr, stderrOut)
		errCh <- e
	}()

	return errCh
}

func exitStatusCh(cmd *exec.Cmd) chan uint {
	exitCh := make(chan uint)
	go func() {
		err := cmd.Wait()
		if err != nil {
			if exiterr, ok := err.(*exec.ExitError); ok {
				// There is no plattform independent way to retrieve
				// the exit code, but the following will work on Unix
				if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
					exitCh <- uint(status.ExitStatus())
				}
			} else {
				panic(err)
			}
			return
		}
		exitCh <- uint(0)
	}()
	return exitCh
}

func scheduleAndAttach(config docker.Config) {
        disc, err := discover.NewClient()
        if err != nil {
                log.Fatal(err)
        }

        scheduler, err := client.New()
        if err != nil {
                log.Fatal(err)
        }

        state, err := scheduler.State()
        if err != nil {
                log.Fatal(err)
        }

        var firstHost string
        for k := range state {
                firstHost = k
                break
        }
        if firstHost == "" {
                log.Fatal("no hosts")
        }

        id := randomID()

        services, err := disc.Services("flynn-lorne-attach." + firstHost)
        if err != nil {
                log.Fatal(err)
        }
        conn, err := net.Dial("tcp", services.OnlineAddrs()[0])
        if err != nil {
                log.Fatal(err)
        }
        err = gob.NewEncoder(conn).Encode(&lorne.AttachReq{
                JobID:  id,
                Flags:  lorne.AttachFlagStdout | lorne.AttachFlagStderr | lorne.AttachFlagStdin | lorne.AttachFlagStream,
        })
        if err != nil {
                log.Fatal(err)
        }
        attachState := make([]byte, 1)
        if _, err := conn.Read(attachState); err != nil {
                log.Fatal(err)
        }
        switch attachState[0] {
        case lorne.AttachError:
                log.Fatal("attach error")
        }

        schedReq := &sampi.ScheduleReq{
                Incremental: true,
                HostJobs: map[string][]*sampi.Job{firstHost: {{ID: id, Config: &config}}},
        }
        if _, err := scheduler.Schedule(schedReq); err != nil {
                log.Fatal(err)
        }

        if _, err := conn.Read(attachState); err != nil {
                log.Fatal(err)
        }

        go io.Copy(conn, os.Stdin)
        if _, err := io.Copy(os.Stdout, conn); err != nil {
                log.Fatal(err)
        }
}

func randomID() string {
        b := make([]byte, 16)
        enc := make([]byte, 24)
        _, err := io.ReadFull(rand.Reader, b)
        if err != nil {
                panic(err) // This shouldn't ever happen, right?
        }
        base64.URLEncoding.Encode(enc, b)
        return string(bytes.TrimRight(enc, "="))
}
<commit_msg>demo: gofmt<commit_after>package main

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"encoding/gob"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"strings"
	"syscall"

	"github.com/flynn/go-discover/discover"
	"github.com/flynn/lorne/types"
	"github.com/flynn/sampi/client"
	"github.com/flynn/sampi/types"
	"github.com/titanous/go-dockerclient"
)

func main() {
	root := "/var/lib/demo/apps"
	hostname := shell("curl -s icanhazip.com")

	client, err := discover.NewClient()
	if err != nil {
		panic(err)
	}
	set, _ := client.Services("shelf")
	addrs := set.OnlineAddrs()
	if len(addrs) < 1 {
		panic("Shelf is not discoverable")
	}
	shelfHost := addrs[0]

	app := os.Args[2]
	os.MkdirAll(root+"/"+app, 0755)

	fmt.Printf("-----> Building %s on %s ...\n", app, hostname)

	scheduleAndAttach(docker.Config{
		Image:        "flynn/slugbuilder",
		Cmd:          []string{"http://" + shelfHost + "/" + app + ".tgz"},
		Tty:          false,
		AttachStdin:  true,
		AttachStdout: true,
		AttachStderr: true,
		OpenStdin:    true,
		StdinOnce:    true,
	})

	/*fmt.Printf("-----> Deploying %s ...\n", app)
		if _, err := os.Stat(root + "/" + app + "/CONTAINER"); err == nil {
	    	oldid := readFile(root + "/" + app + "/CONTAINER")
	    	shell("docker kill " + oldid)
		}

		id := shell("docker run -d -p 5000 -e PORT=5000 -e SLUG_URL=http://"+shelfHost+"/"+app+".tgz flynn/slugrunner start web")
		writeFile(root + "/" + app + "/CONTAINER", id)
		port := shell("docker port "+id+" 5000 | sed 's/0.0.0.0://'")
		writeFile(root + "/" + app + "/PORT", port)
		writeFile(root + "/" + app + "/URL", "http://"+hostname+":"+port)

		fmt.Printf("=====> Application deployed:\n")
		fmt.Printf("       %s\n", readFile(root + "/" + app + "/URL"))*/
	fmt.Println("")

}

func readFile(filename string) string {
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		panic(err)
	}
	return string(data)
}

func writeFile(filename, data string) {
	err := ioutil.WriteFile(filename, []byte(data), 0644)
	if err != nil {
		panic(err)
	}
}

func shell(cmdline string) string {
	out, err := exec.Command("bash", "-c", cmdline).Output()
	if err != nil {
		panic(err)
	}
	return strings.Trim(string(out), " \n")
}

func attachCmd(cmd *exec.Cmd, stdout, stderr io.Writer, stdin io.Reader) chan error {
	errCh := make(chan error)

	stdinIn, err := cmd.StdinPipe()
	if err != nil {
		panic(err)
	}
	stdoutOut, err := cmd.StdoutPipe()
	if err != nil {
		panic(err)
	}
	stderrOut, err := cmd.StderrPipe()
	if err != nil {
		panic(err)
	}

	go func() {
		_, e := io.Copy(stdinIn, stdin)
		errCh <- e
	}()
	go func() {
		_, e := io.Copy(stdout, stdoutOut)
		errCh <- e
	}()
	go func() {
		_, e := io.Copy(stderr, stderrOut)
		errCh <- e
	}()

	return errCh
}

func exitStatusCh(cmd *exec.Cmd) chan uint {
	exitCh := make(chan uint)
	go func() {
		err := cmd.Wait()
		if err != nil {
			if exiterr, ok := err.(*exec.ExitError); ok {
				// There is no plattform independent way to retrieve
				// the exit code, but the following will work on Unix
				if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
					exitCh <- uint(status.ExitStatus())
				}
			} else {
				panic(err)
			}
			return
		}
		exitCh <- uint(0)
	}()
	return exitCh
}

func scheduleAndAttach(config docker.Config) {
	disc, err := discover.NewClient()
	if err != nil {
		log.Fatal(err)
	}

	scheduler, err := client.New()
	if err != nil {
		log.Fatal(err)
	}

	state, err := scheduler.State()
	if err != nil {
		log.Fatal(err)
	}

	var firstHost string
	for k := range state {
		firstHost = k
		break
	}
	if firstHost == "" {
		log.Fatal("no hosts")
	}

	id := randomID()

	services, err := disc.Services("flynn-lorne-attach." + firstHost)
	if err != nil {
		log.Fatal(err)
	}
	conn, err := net.Dial("tcp", services.OnlineAddrs()[0])
	if err != nil {
		log.Fatal(err)
	}
	err = gob.NewEncoder(conn).Encode(&lorne.AttachReq{
		JobID: id,
		Flags: lorne.AttachFlagStdout | lorne.AttachFlagStderr | lorne.AttachFlagStdin | lorne.AttachFlagStream,
	})
	if err != nil {
		log.Fatal(err)
	}
	attachState := make([]byte, 1)
	if _, err := conn.Read(attachState); err != nil {
		log.Fatal(err)
	}
	switch attachState[0] {
	case lorne.AttachError:
		log.Fatal("attach error")
	}

	schedReq := &sampi.ScheduleReq{
		Incremental: true,
		HostJobs:    map[string][]*sampi.Job{firstHost: {{ID: id, Config: &config}}},
	}
	if _, err := scheduler.Schedule(schedReq); err != nil {
		log.Fatal(err)
	}

	if _, err := conn.Read(attachState); err != nil {
		log.Fatal(err)
	}

	go io.Copy(conn, os.Stdin)
	if _, err := io.Copy(os.Stdout, conn); err != nil {
		log.Fatal(err)
	}
}

func randomID() string {
	b := make([]byte, 16)
	enc := make([]byte, 24)
	_, err := io.ReadFull(rand.Reader, b)
	if err != nil {
		panic(err) // This shouldn't ever happen, right?
	}
	base64.URLEncoding.Encode(enc, b)
	return string(bytes.TrimRight(enc, "="))
}
<|endoftext|>