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

// Package yacr is yet another CSV reader (and writer) with small memory usage.
package yacr

import (
	"bufio"
	"bytes"
	"encoding"
	"fmt"
	"io"
	"reflect"
	"strconv"
)

// Reader provides an interface for reading CSV data
// (compatible with rfc4180 and extended with the option of having a separator other than ",").
// Successive calls to the Scan method will step through the 'fields', skipping the separator/newline between the fields.
// The EndOfRecord method tells when a field is terminated by a line break.
type Reader struct {
	*bufio.Scanner
	sep    byte // values separator
	quoted bool // specify if values may be quoted (when they contain separator or newline)
	guess  bool // try to guess separator based on the file header
	eor    bool // true when the most recent field has been terminated by a newline (not a separator).
	lineno int  // current line number (not record number)
	empty  bool // true when the current line is empty (or a line comment)

	Trim    bool // trim spaces (only on unquoted values). Break rfc4180 rule: "Spaces are considered part of a field and should not be ignored."
	Comment byte // character marking the start of a line comment. When specified, line comment appears as empty line.
	Lazy    bool // specify if quoted values may contains unescaped quote not followed by a separator or a newline
}

// DefaultReader creates a "standard" CSV reader (separator is comma and quoted mode active)
func DefaultReader(rd io.Reader) *Reader {
	return NewReader(rd, ',', true, false)
}

// NewReader returns a new CSV scanner to read from r.
// When quoted is false, values must not contain a separator or newline.
func NewReader(r io.Reader, sep byte, quoted, guess bool) *Reader {
	s := &Reader{bufio.NewScanner(r), sep, quoted, guess, true, 1, false, false, 0, false}
	s.Split(s.ScanField)
	return s
}

// ScanRecord decodes one line fields to values.
// It's like fmt.Scan or database.sql.Rows.Scan.
func (s *Reader) ScanRecord(values ...interface{}) (int, error) {
	for i, value := range values {
		if !s.Scan() {
			return i, s.Err()
		}
		if i == 0 { // skip empty line (or line comment)
			for s.EmptyLine() {
				if !s.Scan() {
					return i, s.Err()
				}
			}
		}
		if err := s.value(value, true); err != nil {
			return i, err
		} else if s.EndOfRecord() != (i == len(values)-1) {
			return i, fmt.Errorf("unexpected number of fields: want %d, got %d", len(values), i+1)
		}
	}
	return len(values), nil
}

// ScanValue advances to the next token and decodes field's content to value.
// The value may point to data that will be overwritten by a subsequent call to Scan.
func (s *Reader) ScanValue(value interface{}) error {
	if !s.Scan() {
		return s.Err()
	}
	return s.value(value, false)
}

// Value decodes field's content to value.
// The value may point to data that will be overwritten by a subsequent call to Scan.
func (s *Reader) Value(value interface{}) error {
	return s.value(value, false)
}
func (s *Reader) value(value interface{}, copied bool) error {
	var err error
	switch value := value.(type) {
	case nil:
	case *string:
		*value = s.Text()
	case *int:
		*value, err = strconv.Atoi(s.Text())
	case *int32:
		var i int64
		i, err = strconv.ParseInt(s.Text(), 10, 32)
		*value = int32(i)
	case *int64:
		*value, err = strconv.ParseInt(s.Text(), 10, 64)
	case *bool:
		*value, err = strconv.ParseBool(s.Text())
	case *float64:
		*value, err = strconv.ParseFloat(s.Text(), 64)
	case *[]byte:
		if copied {
			v := s.Bytes()
			c := make([]byte, len(v))
			copy(c, v)
			*value = c
		} else {
			*value = s.Bytes()
		}
	case encoding.TextUnmarshaler:
		err = value.UnmarshalText(s.Bytes())
	default:
		return s.scanReflect(value)
	}
	return err
}

func (s *Reader) scanReflect(v interface{}) (err error) {
	rv := reflect.ValueOf(v)
	if rv.Kind() != reflect.Ptr || rv.IsNil() {
		return fmt.Errorf("unsupported type %T", v)
	}
	dv := reflect.Indirect(rv)
	switch dv.Kind() {
	case reflect.String:
		dv.SetString(s.Text())
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		var i int64
		i, err = strconv.ParseInt(s.Text(), 10, dv.Type().Bits())
		if err == nil {
			dv.SetInt(i)
		}
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
		var i uint64
		i, err = strconv.ParseUint(s.Text(), 10, dv.Type().Bits())
		if err == nil {
			dv.SetUint(i)
		}
	case reflect.Bool:
		var b bool
		b, err = strconv.ParseBool(s.Text())
		if err == nil {
			dv.SetBool(b)
		}
	case reflect.Float32, reflect.Float64:
		var f float64
		f, err = strconv.ParseFloat(s.Text(), dv.Type().Bits())
		if err == nil {
			dv.SetFloat(f)
		}
	default:
		return fmt.Errorf("unsupported type: %T", v)
	}
	return
}

// LineNumber returns current line number (not record number)
func (s *Reader) LineNumber() int {
	return s.lineno
}

// EndOfRecord returns true when the most recent field has been terminated by a newline (not a separator).
func (s *Reader) EndOfRecord() bool {
	return s.eor
}

// EmptyLine returns true when the current line is empty or a line comment.
func (s *Reader) EmptyLine() bool {
	return s.empty && s.eor
}

// Sep returns the values separator used/guessed
func (s *Reader) Sep() byte {
	return s.sep
}

// ScanField implements bufio.SplitFunc for CSV.
// Lexing is adapted from csv_read_one_field function in SQLite3 shell sources.
func (s *Reader) ScanField(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 && s.eor {
		return 0, nil, nil
	}
	if s.guess {
		s.guess = false
		if b := guess(data); b > 0 {
			s.sep = b
		}
	}
	if s.quoted && len(data) > 0 && data[0] == '"' { // quoted field (may contains separator, newline and escaped quote)
		s.empty = false
		startLineno := s.lineno
		escapedQuotes := 0
		strict := true
		var c, pc, ppc byte
		// Scan until the separator or newline following the closing quote (and ignore escaped quote)
		for i := 1; i < len(data); i++ {
			c = data[i]
			if c == '\n' {
				s.lineno++
			} else if c == '"' {
				if pc == c { // escaped quote
					pc = 0
					escapedQuotes++
					continue
				}
			}
			if pc == '"' && c == s.sep {
				s.eor = false
				return i + 1, unescapeQuotes(data[1:i-1], escapedQuotes, strict), nil
			} else if pc == '"' && c == '\n' {
				s.eor = true
				return i + 1, unescapeQuotes(data[1:i-1], escapedQuotes, strict), nil
			} else if c == '\n' && pc == '\r' && ppc == '"' {
				s.eor = true
				return i + 1, unescapeQuotes(data[1:i-2], escapedQuotes, strict), nil
			}
			if pc == '"' && c != '\r' {
				if s.Lazy {
					strict = false
				} else {
					return 0, nil, fmt.Errorf("unescaped %c character at line %d", pc, s.lineno)
				}
			}
			ppc = pc
			pc = c
		}
		if atEOF {
			if c == '"' {
				s.eor = true
				return len(data), unescapeQuotes(data[1:len(data)-1], escapedQuotes, strict), nil
			}
			// If we're at EOF, we have a non-terminated field.
			return 0, nil, fmt.Errorf("non-terminated quoted field at line %d", startLineno)
		}
	} else if s.eor && s.Comment != 0 && len(data) > 0 && data[0] == s.Comment { // line comment
		s.empty = true
		for i, c := range data {
			if c == '\n' {
				return i + 1, nil, nil
			}
		}
		if atEOF {
			return len(data), nil, nil
		}
	} else { // unquoted field
		// Scan until separator or newline, marking end of field.
		for i, c := range data {
			if c == s.sep {
				s.eor = false
				if s.Trim {
					return i + 1, trim(data[0:i]), nil
				}
				return i + 1, data[0:i], nil
			} else if c == '\n' {
				s.lineno++
				if i > 0 && data[i-1] == '\r' {
					s.empty = s.eor && i == 1 // FIXME empty & trim
					s.eor = true
					if s.Trim {
						return i + 1, trim(data[0 : i-1]), nil
					}
					return i + 1, data[0 : i-1], nil
				}
				s.empty = s.eor && i == 0 // FIXME empty & trim
				s.eor = true
				if s.Trim {
					return i + 1, trim(data[0:i]), nil
				}
				return i + 1, data[0:i], nil
			}
		}
		// If we're at EOF, we have a final field. Return it.
		if atEOF {
			s.empty = false
			s.eor = true
			if s.Trim {
				return len(data), trim(data), nil
			}
			return len(data), data, nil
		}
	}
	// Request more data.
	return 0, nil, nil
}

func unescapeQuotes(b []byte, count int, strict bool) []byte {
	if count == 0 {
		return b
	}
	for i, j := 0, 0; i < len(b); i, j = i+1, j+1 {
		b[j] = b[i]
		if b[i] == '"' && (strict || i < len(b)-1 && b[i+1] == '"') {
			i++
		}
	}
	return b[:len(b)-count]
}

func guess(data []byte) byte {
	seps := []byte{',', ';', '\t', '|', ':'}
	count := make(map[byte]uint)
	for _, b := range data {
		if bytes.IndexByte(seps, b) >= 0 {
			count[b]++
			/*} else if b == '\n' {
			break*/
		}
	}
	var max uint
	var sep byte
	for b, c := range count {
		if c > max {
			max = c
			sep = b
		}
	}
	return sep
}

// bytes.TrimSpace may return nil...
func trim(s []byte) []byte {
	t := bytes.TrimSpace(s)
	if t == nil {
		return s[0:0]
	}
	return t
}
<commit_msg>Correct message. i+1 is the current length, msg should say 1+2. After the fix, sample message is "unexpected number of fields: want 15, got 16 (or more)"<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package yacr is yet another CSV reader (and writer) with small memory usage.
package yacr

import (
	"bufio"
	"bytes"
	"encoding"
	"fmt"
	"io"
	"reflect"
	"strconv"
)

// Reader provides an interface for reading CSV data
// (compatible with rfc4180 and extended with the option of having a separator other than ",").
// Successive calls to the Scan method will step through the 'fields', skipping the separator/newline between the fields.
// The EndOfRecord method tells when a field is terminated by a line break.
type Reader struct {
	*bufio.Scanner
	sep    byte // values separator
	quoted bool // specify if values may be quoted (when they contain separator or newline)
	guess  bool // try to guess separator based on the file header
	eor    bool // true when the most recent field has been terminated by a newline (not a separator).
	lineno int  // current line number (not record number)
	empty  bool // true when the current line is empty (or a line comment)

	Trim    bool // trim spaces (only on unquoted values). Break rfc4180 rule: "Spaces are considered part of a field and should not be ignored."
	Comment byte // character marking the start of a line comment. When specified, line comment appears as empty line.
	Lazy    bool // specify if quoted values may contains unescaped quote not followed by a separator or a newline
}

// DefaultReader creates a "standard" CSV reader (separator is comma and quoted mode active)
func DefaultReader(rd io.Reader) *Reader {
	return NewReader(rd, ',', true, false)
}

// NewReader returns a new CSV scanner to read from r.
// When quoted is false, values must not contain a separator or newline.
func NewReader(r io.Reader, sep byte, quoted, guess bool) *Reader {
	s := &Reader{bufio.NewScanner(r), sep, quoted, guess, true, 1, false, false, 0, false}
	s.Split(s.ScanField)
	return s
}

// ScanRecord decodes one line fields to values.
// It's like fmt.Scan or database.sql.Rows.Scan.
func (s *Reader) ScanRecord(values ...interface{}) (int, error) {
	for i, value := range values {
		if !s.Scan() {
			return i, s.Err()
		}
		if i == 0 { // skip empty line (or line comment)
			for s.EmptyLine() {
				if !s.Scan() {
					return i, s.Err()
				}
			}
		}
		if err := s.value(value, true); err != nil {
			return i, err
		} else if s.EndOfRecord() != (i == len(values)-1) {
			return i, fmt.Errorf("unexpected number of fields: want %d, got %d (or more)", len(values), i+2)
		}
	}
	return len(values), nil
}

// ScanValue advances to the next token and decodes field's content to value.
// The value may point to data that will be overwritten by a subsequent call to Scan.
func (s *Reader) ScanValue(value interface{}) error {
	if !s.Scan() {
		return s.Err()
	}
	return s.value(value, false)
}

// Value decodes field's content to value.
// The value may point to data that will be overwritten by a subsequent call to Scan.
func (s *Reader) Value(value interface{}) error {
	return s.value(value, false)
}
func (s *Reader) value(value interface{}, copied bool) error {
	var err error
	switch value := value.(type) {
	case nil:
	case *string:
		*value = s.Text()
	case *int:
		*value, err = strconv.Atoi(s.Text())
	case *int32:
		var i int64
		i, err = strconv.ParseInt(s.Text(), 10, 32)
		*value = int32(i)
	case *int64:
		*value, err = strconv.ParseInt(s.Text(), 10, 64)
	case *bool:
		*value, err = strconv.ParseBool(s.Text())
	case *float64:
		*value, err = strconv.ParseFloat(s.Text(), 64)
	case *[]byte:
		if copied {
			v := s.Bytes()
			c := make([]byte, len(v))
			copy(c, v)
			*value = c
		} else {
			*value = s.Bytes()
		}
	case encoding.TextUnmarshaler:
		err = value.UnmarshalText(s.Bytes())
	default:
		return s.scanReflect(value)
	}
	return err
}

func (s *Reader) scanReflect(v interface{}) (err error) {
	rv := reflect.ValueOf(v)
	if rv.Kind() != reflect.Ptr || rv.IsNil() {
		return fmt.Errorf("unsupported type %T", v)
	}
	dv := reflect.Indirect(rv)
	switch dv.Kind() {
	case reflect.String:
		dv.SetString(s.Text())
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		var i int64
		i, err = strconv.ParseInt(s.Text(), 10, dv.Type().Bits())
		if err == nil {
			dv.SetInt(i)
		}
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
		var i uint64
		i, err = strconv.ParseUint(s.Text(), 10, dv.Type().Bits())
		if err == nil {
			dv.SetUint(i)
		}
	case reflect.Bool:
		var b bool
		b, err = strconv.ParseBool(s.Text())
		if err == nil {
			dv.SetBool(b)
		}
	case reflect.Float32, reflect.Float64:
		var f float64
		f, err = strconv.ParseFloat(s.Text(), dv.Type().Bits())
		if err == nil {
			dv.SetFloat(f)
		}
	default:
		return fmt.Errorf("unsupported type: %T", v)
	}
	return
}

// LineNumber returns current line number (not record number)
func (s *Reader) LineNumber() int {
	return s.lineno
}

// EndOfRecord returns true when the most recent field has been terminated by a newline (not a separator).
func (s *Reader) EndOfRecord() bool {
	return s.eor
}

// EmptyLine returns true when the current line is empty or a line comment.
func (s *Reader) EmptyLine() bool {
	return s.empty && s.eor
}

// Sep returns the values separator used/guessed
func (s *Reader) Sep() byte {
	return s.sep
}

// ScanField implements bufio.SplitFunc for CSV.
// Lexing is adapted from csv_read_one_field function in SQLite3 shell sources.
func (s *Reader) ScanField(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 && s.eor {
		return 0, nil, nil
	}
	if s.guess {
		s.guess = false
		if b := guess(data); b > 0 {
			s.sep = b
		}
	}
	if s.quoted && len(data) > 0 && data[0] == '"' { // quoted field (may contains separator, newline and escaped quote)
		s.empty = false
		startLineno := s.lineno
		escapedQuotes := 0
		strict := true
		var c, pc, ppc byte
		// Scan until the separator or newline following the closing quote (and ignore escaped quote)
		for i := 1; i < len(data); i++ {
			c = data[i]
			if c == '\n' {
				s.lineno++
			} else if c == '"' {
				if pc == c { // escaped quote
					pc = 0
					escapedQuotes++
					continue
				}
			}
			if pc == '"' && c == s.sep {
				s.eor = false
				return i + 1, unescapeQuotes(data[1:i-1], escapedQuotes, strict), nil
			} else if pc == '"' && c == '\n' {
				s.eor = true
				return i + 1, unescapeQuotes(data[1:i-1], escapedQuotes, strict), nil
			} else if c == '\n' && pc == '\r' && ppc == '"' {
				s.eor = true
				return i + 1, unescapeQuotes(data[1:i-2], escapedQuotes, strict), nil
			}
			if pc == '"' && c != '\r' {
				if s.Lazy {
					strict = false
				} else {
					return 0, nil, fmt.Errorf("unescaped %c character at line %d", pc, s.lineno)
				}
			}
			ppc = pc
			pc = c
		}
		if atEOF {
			if c == '"' {
				s.eor = true
				return len(data), unescapeQuotes(data[1:len(data)-1], escapedQuotes, strict), nil
			}
			// If we're at EOF, we have a non-terminated field.
			return 0, nil, fmt.Errorf("non-terminated quoted field at line %d", startLineno)
		}
	} else if s.eor && s.Comment != 0 && len(data) > 0 && data[0] == s.Comment { // line comment
		s.empty = true
		for i, c := range data {
			if c == '\n' {
				return i + 1, nil, nil
			}
		}
		if atEOF {
			return len(data), nil, nil
		}
	} else { // unquoted field
		// Scan until separator or newline, marking end of field.
		for i, c := range data {
			if c == s.sep {
				s.eor = false
				if s.Trim {
					return i + 1, trim(data[0:i]), nil
				}
				return i + 1, data[0:i], nil
			} else if c == '\n' {
				s.lineno++
				if i > 0 && data[i-1] == '\r' {
					s.empty = s.eor && i == 1 // FIXME empty & trim
					s.eor = true
					if s.Trim {
						return i + 1, trim(data[0 : i-1]), nil
					}
					return i + 1, data[0 : i-1], nil
				}
				s.empty = s.eor && i == 0 // FIXME empty & trim
				s.eor = true
				if s.Trim {
					return i + 1, trim(data[0:i]), nil
				}
				return i + 1, data[0:i], nil
			}
		}
		// If we're at EOF, we have a final field. Return it.
		if atEOF {
			s.empty = false
			s.eor = true
			if s.Trim {
				return len(data), trim(data), nil
			}
			return len(data), data, nil
		}
	}
	// Request more data.
	return 0, nil, nil
}

func unescapeQuotes(b []byte, count int, strict bool) []byte {
	if count == 0 {
		return b
	}
	for i, j := 0, 0; i < len(b); i, j = i+1, j+1 {
		b[j] = b[i]
		if b[i] == '"' && (strict || i < len(b)-1 && b[i+1] == '"') {
			i++
		}
	}
	return b[:len(b)-count]
}

func guess(data []byte) byte {
	seps := []byte{',', ';', '\t', '|', ':'}
	count := make(map[byte]uint)
	for _, b := range data {
		if bytes.IndexByte(seps, b) >= 0 {
			count[b]++
			/*} else if b == '\n' {
			break*/
		}
	}
	var max uint
	var sep byte
	for b, c := range count {
		if c > max {
			max = c
			sep = b
		}
	}
	return sep
}

// bytes.TrimSpace may return nil...
func trim(s []byte) []byte {
	t := bytes.TrimSpace(s)
	if t == nil {
		return s[0:0]
	}
	return t
}
<|endoftext|>


<commit_before>//
// Copyright © 2011 Guy M. Allard
//
// 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 stomp

import (
	"os"
	"strconv"
	"strings"
	"time"
)

// Network reader
func (c *Connection) reader() {
	//
	q := false
	gf := func() {
		q = <-c.rsd
	}
	go gf()

	for {
		f, e := c.readFrame()
		if e != nil {
			if e == os.EOF {
				break
			}
			c.input <- MessageData{Message{"", Headers{}, NULLBUFF}, e}
			break
		}

		if f.Command == "" && q {
			break
		}

		d := MessageData{Message{f.Command, f.Headers, f.Body}, e}
		c.input <- d
		if q {
			break
		}

	}

}

// Frame reader
func (c *Connection) readFrame() (f Frame, e os.Error) {
	f = Frame{"", Headers{}, NULLBUFF}
	e = nil
	// Read f.Command or line ends (maybe heartbeats)
	for {
		s, e := c.rdr.ReadString('\n')
		if s == "" {
			return f, e
		}
		if e != nil {
			return f, e
		}
		if c.hbd != nil {
			c.hbd.lr = time.Nanoseconds() // Latest good read
		}
		f.Command = s[0 : len(s)-1]
		if s != "\n" {
			break
		}
	}
	// Read f.Headers
	for {
		s, e := c.rdr.ReadString('\n')
		if e != nil {
			return f, e
		}
		if c.hbd != nil {
			c.hbd.lr = time.Nanoseconds() // Latest good read
		}
		if s == "\n" {
			break
		}
		//
		s = s[0:len(s)-1]
		p := strings.SplitN(s, ":", 2)
		k := p[0]
		v := p[1]
		if c.protocol != SPL_10 {
			k = decode(k)
			v = decode(v)
		}
		f.Headers = append(f.Headers, k, v)
	}
	// Read f.Body
	if v, ok := f.Headers.Contains("content-length"); ok {
		l, e := strconv.Atoi(strings.TrimSpace(v))
		if e != nil {
			return f, e
		}
		if l == 0 {
			f.Body, e = readUntilNul(c.rdr)
		} else {
			f.Body, e = readBody(c.rdr, l)
		}
	} else {
		// content-length not present
		f.Body, e = readUntilNul(c.rdr)
	}
	if e != nil {
		return f, e
	}
	if c.hbd != nil {
		c.hbd.lr = time.Nanoseconds() // Latest good read
	}
	//
	return f, e
}
<commit_msg>Never decode CONNECTED frame headers during read.<commit_after>//
// Copyright © 2011 Guy M. Allard
//
// 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 stomp

import (
	"os"
	"strconv"
	"strings"
	"time"
)

// Network reader
func (c *Connection) reader() {
	//
	q := false
	gf := func() {
		q = <-c.rsd
	}
	go gf()

	for {
		f, e := c.readFrame()
		if e != nil {
			if e == os.EOF {
				break
			}
			c.input <- MessageData{Message{"", Headers{}, NULLBUFF}, e}
			break
		}

		if f.Command == "" && q {
			break
		}

		d := MessageData{Message{f.Command, f.Headers, f.Body}, e}
		c.input <- d
		if q {
			break
		}

	}

}

// Frame reader
func (c *Connection) readFrame() (f Frame, e os.Error) {
	f = Frame{"", Headers{}, NULLBUFF}
	e = nil
	// Read f.Command or line ends (maybe heartbeats)
	for {
		s, e := c.rdr.ReadString('\n')
		if s == "" {
			return f, e
		}
		if e != nil {
			return f, e
		}
		if c.hbd != nil {
			c.hbd.lr = time.Nanoseconds() // Latest good read
		}
		f.Command = s[0 : len(s)-1]
		if s != "\n" {
			break
		}
	}
	// Read f.Headers
	for {
		s, e := c.rdr.ReadString('\n')
		if e != nil {
			return f, e
		}
		if c.hbd != nil {
			c.hbd.lr = time.Nanoseconds() // Latest good read
		}
		if s == "\n" {
			break
		}
		//
		s = s[0:len(s)-1]
		p := strings.SplitN(s, ":", 2)
		k := p[0]
		v := p[1]
		if c.protocol != SPL_10 && f.Command != CONNECTED {
			k = decode(k)
			v = decode(v)
		}
		f.Headers = append(f.Headers, k, v)
	}
	// Read f.Body
	if v, ok := f.Headers.Contains("content-length"); ok {
		l, e := strconv.Atoi(strings.TrimSpace(v))
		if e != nil {
			return f, e
		}
		if l == 0 {
			f.Body, e = readUntilNul(c.rdr)
		} else {
			f.Body, e = readBody(c.rdr, l)
		}
	} else {
		// content-length not present
		f.Body, e = readUntilNul(c.rdr)
	}
	if e != nil {
		return f, e
	}
	if c.hbd != nil {
		c.hbd.lr = time.Nanoseconds() // Latest good read
	}
	//
	return f, e
}
<|endoftext|>


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

package unix

import (
	"runtime"
	"unsafe"
)

// IoctlRetInt performs an ioctl operation specified by req on a device
// associated with opened file descriptor fd, and returns a non-negative
// integer that is returned by the ioctl syscall.
func IoctlRetInt(fd int, req uint) (int, error) {
	ret, _, err := Syscall(SYS_IOCTL, uintptr(fd), uintptr(req), 0)
	if err != 0 {
		return 0, err
	}
	return int(ret), nil
}

func IoctlGetUint32(fd int, req uint) (uint32, error) {
	var value uint32
	err := ioctlPtr(fd, req, unsafe.Pointer(&value))
	return value, err
}

func IoctlGetRTCTime(fd int) (*RTCTime, error) {
	var value RTCTime
	err := ioctlPtr(fd, RTC_RD_TIME, unsafe.Pointer(&value))
	return &value, err
}

func IoctlSetRTCTime(fd int, value *RTCTime) error {
	err := ioctlPtr(fd, RTC_SET_TIME, unsafe.Pointer(value))
	runtime.KeepAlive(value)
	return err
}

func IoctlGetRTCWkAlrm(fd int) (*RTCWkAlrm, error) {
	var value RTCWkAlrm
	err := ioctlPtr(fd, RTC_WKALM_RD, unsafe.Pointer(&value))
	return &value, err
}

func IoctlSetRTCWkAlrm(fd int, value *RTCWkAlrm) error {
	err := ioctlPtr(fd, RTC_WKALM_SET, unsafe.Pointer(value))
	runtime.KeepAlive(value)
	return err
}

// IoctlGetEthtoolDrvinfo fetches ethtool driver information for the network
// device specified by ifname.
func IoctlGetEthtoolDrvinfo(fd int, ifname string) (*EthtoolDrvinfo, error) {
	ifr, err := newIfreq(ifname)
	if err != nil {
		return nil, err
	}

	value := EthtoolDrvinfo{Cmd: ETHTOOL_GDRVINFO}
	ifrd := ifr.SetData(unsafe.Pointer(&value))

	err = ioctlPtr(fd, SIOCETHTOOL, unsafe.Pointer(&ifrd))
	runtime.KeepAlive(ifrd)
	return &value, err
}

// IoctlGetWatchdogInfo fetches information about a watchdog device from the
// Linux watchdog API. For more information, see:
// https://www.kernel.org/doc/html/latest/watchdog/watchdog-api.html.
func IoctlGetWatchdogInfo(fd int) (*WatchdogInfo, error) {
	var value WatchdogInfo
	err := ioctlPtr(fd, WDIOC_GETSUPPORT, unsafe.Pointer(&value))
	return &value, err
}

// IoctlWatchdogKeepalive issues a keepalive ioctl to a watchdog device. For
// more information, see:
// https://www.kernel.org/doc/html/latest/watchdog/watchdog-api.html.
func IoctlWatchdogKeepalive(fd int) error {
	// arg is ignored and not a pointer, so ioctl is fine instead of ioctlPtr.
	return ioctl(fd, WDIOC_KEEPALIVE, 0)
}

// IoctlFileCloneRange performs an FICLONERANGE ioctl operation to clone the
// range of data conveyed in value to the file associated with the file
// descriptor destFd. See the ioctl_ficlonerange(2) man page for details.
func IoctlFileCloneRange(destFd int, value *FileCloneRange) error {
	err := ioctlPtr(destFd, FICLONERANGE, unsafe.Pointer(value))
	runtime.KeepAlive(value)
	return err
}

// IoctlFileClone performs an FICLONE ioctl operation to clone the entire file
// associated with the file description srcFd to the file associated with the
// file descriptor destFd. See the ioctl_ficlone(2) man page for details.
func IoctlFileClone(destFd, srcFd int) error {
	return ioctl(destFd, FICLONE, uintptr(srcFd))
}

type FileDedupeRange struct {
	Src_offset uint64
	Src_length uint64
	Reserved1  uint16
	Reserved2  uint32
	Info       []FileDedupeRangeInfo
}

type FileDedupeRangeInfo struct {
	Dest_fd       int64
	Dest_offset   uint64
	Bytes_deduped uint64
	Status        int32
	Reserved      uint32
}

// IoctlFileDedupeRange performs an FIDEDUPERANGE ioctl operation to share the
// range of data conveyed in value from the file associated with the file
// descriptor srcFd to the value.Info destinations. See the
// ioctl_fideduperange(2) man page for details.
func IoctlFileDedupeRange(srcFd int, value *FileDedupeRange) error {
	buf := make([]byte, SizeofRawFileDedupeRange+
		len(value.Info)*SizeofRawFileDedupeRangeInfo)
	rawrange := (*RawFileDedupeRange)(unsafe.Pointer(&buf[0]))
	rawrange.Src_offset = value.Src_offset
	rawrange.Src_length = value.Src_length
	rawrange.Dest_count = uint16(len(value.Info))
	rawrange.Reserved1 = value.Reserved1
	rawrange.Reserved2 = value.Reserved2

	for i := range value.Info {
		rawinfo := (*RawFileDedupeRangeInfo)(unsafe.Pointer(
			uintptr(unsafe.Pointer(&buf[0])) + uintptr(SizeofRawFileDedupeRange) +
				uintptr(i*SizeofRawFileDedupeRangeInfo)))
		rawinfo.Dest_fd = value.Info[i].Dest_fd
		rawinfo.Dest_offset = value.Info[i].Dest_offset
		rawinfo.Bytes_deduped = value.Info[i].Bytes_deduped
		rawinfo.Status = value.Info[i].Status
		rawinfo.Reserved = value.Info[i].Reserved
	}

	err := ioctlPtr(srcFd, FIDEDUPERANGE, unsafe.Pointer(&buf[0]))

	// Output
	for i := range value.Info {
		rawinfo := (*RawFileDedupeRangeInfo)(unsafe.Pointer(
			uintptr(unsafe.Pointer(&buf[0])) + uintptr(SizeofRawFileDedupeRange) +
				uintptr(i*SizeofRawFileDedupeRangeInfo)))
		value.Info[i].Dest_fd = rawinfo.Dest_fd
		value.Info[i].Dest_offset = rawinfo.Dest_offset
		value.Info[i].Bytes_deduped = rawinfo.Bytes_deduped
		value.Info[i].Status = rawinfo.Status
		value.Info[i].Reserved = rawinfo.Reserved
	}

	return err
}

func IoctlHIDGetDesc(fd int, value *HIDRawReportDescriptor) error {
	err := ioctlPtr(fd, HIDIOCGRDESC, unsafe.Pointer(value))
	runtime.KeepAlive(value)
	return err
}

func IoctlHIDGetRawInfo(fd int) (*HIDRawDevInfo, error) {
	var value HIDRawDevInfo
	err := ioctlPtr(fd, HIDIOCGRAWINFO, unsafe.Pointer(&value))
	return &value, err
}

func IoctlHIDGetRawName(fd int) (string, error) {
	var value [_HIDIOCGRAWNAME_LEN]byte
	err := ioctlPtr(fd, _HIDIOCGRAWNAME, unsafe.Pointer(&value[0]))
	return ByteSliceToString(value[:]), err
}

func IoctlHIDGetRawPhys(fd int) (string, error) {
	var value [_HIDIOCGRAWPHYS_LEN]byte
	err := ioctlPtr(fd, _HIDIOCGRAWPHYS, unsafe.Pointer(&value[0]))
	return ByteSliceToString(value[:]), err
}

func IoctlHIDGetRawUniq(fd int) (string, error) {
	var value [_HIDIOCGRAWUNIQ_LEN]byte
	err := ioctlPtr(fd, _HIDIOCGRAWUNIQ, unsafe.Pointer(&value[0]))
	return ByteSliceToString(value[:]), err
}
<commit_msg>unix: remove runtime.KeepAlive calls for Linux ioctl wrappers<commit_after>// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package unix

import (
	"unsafe"
)

// IoctlRetInt performs an ioctl operation specified by req on a device
// associated with opened file descriptor fd, and returns a non-negative
// integer that is returned by the ioctl syscall.
func IoctlRetInt(fd int, req uint) (int, error) {
	ret, _, err := Syscall(SYS_IOCTL, uintptr(fd), uintptr(req), 0)
	if err != 0 {
		return 0, err
	}
	return int(ret), nil
}

func IoctlGetUint32(fd int, req uint) (uint32, error) {
	var value uint32
	err := ioctlPtr(fd, req, unsafe.Pointer(&value))
	return value, err
}

func IoctlGetRTCTime(fd int) (*RTCTime, error) {
	var value RTCTime
	err := ioctlPtr(fd, RTC_RD_TIME, unsafe.Pointer(&value))
	return &value, err
}

func IoctlSetRTCTime(fd int, value *RTCTime) error {
	return ioctlPtr(fd, RTC_SET_TIME, unsafe.Pointer(value))
}

func IoctlGetRTCWkAlrm(fd int) (*RTCWkAlrm, error) {
	var value RTCWkAlrm
	err := ioctlPtr(fd, RTC_WKALM_RD, unsafe.Pointer(&value))
	return &value, err
}

func IoctlSetRTCWkAlrm(fd int, value *RTCWkAlrm) error {
	return ioctlPtr(fd, RTC_WKALM_SET, unsafe.Pointer(value))
}

// IoctlGetEthtoolDrvinfo fetches ethtool driver information for the network
// device specified by ifname.
func IoctlGetEthtoolDrvinfo(fd int, ifname string) (*EthtoolDrvinfo, error) {
	ifr, err := newIfreq(ifname)
	if err != nil {
		return nil, err
	}

	value := EthtoolDrvinfo{Cmd: ETHTOOL_GDRVINFO}
	ifrd := ifr.SetData(unsafe.Pointer(&value))

	err = ioctlPtr(fd, SIOCETHTOOL, unsafe.Pointer(&ifrd))
	return &value, err
}

// IoctlGetWatchdogInfo fetches information about a watchdog device from the
// Linux watchdog API. For more information, see:
// https://www.kernel.org/doc/html/latest/watchdog/watchdog-api.html.
func IoctlGetWatchdogInfo(fd int) (*WatchdogInfo, error) {
	var value WatchdogInfo
	err := ioctlPtr(fd, WDIOC_GETSUPPORT, unsafe.Pointer(&value))
	return &value, err
}

// IoctlWatchdogKeepalive issues a keepalive ioctl to a watchdog device. For
// more information, see:
// https://www.kernel.org/doc/html/latest/watchdog/watchdog-api.html.
func IoctlWatchdogKeepalive(fd int) error {
	// arg is ignored and not a pointer, so ioctl is fine instead of ioctlPtr.
	return ioctl(fd, WDIOC_KEEPALIVE, 0)
}

// IoctlFileCloneRange performs an FICLONERANGE ioctl operation to clone the
// range of data conveyed in value to the file associated with the file
// descriptor destFd. See the ioctl_ficlonerange(2) man page for details.
func IoctlFileCloneRange(destFd int, value *FileCloneRange) error {
	return ioctlPtr(destFd, FICLONERANGE, unsafe.Pointer(value))
}

// IoctlFileClone performs an FICLONE ioctl operation to clone the entire file
// associated with the file description srcFd to the file associated with the
// file descriptor destFd. See the ioctl_ficlone(2) man page for details.
func IoctlFileClone(destFd, srcFd int) error {
	return ioctl(destFd, FICLONE, uintptr(srcFd))
}

type FileDedupeRange struct {
	Src_offset uint64
	Src_length uint64
	Reserved1  uint16
	Reserved2  uint32
	Info       []FileDedupeRangeInfo
}

type FileDedupeRangeInfo struct {
	Dest_fd       int64
	Dest_offset   uint64
	Bytes_deduped uint64
	Status        int32
	Reserved      uint32
}

// IoctlFileDedupeRange performs an FIDEDUPERANGE ioctl operation to share the
// range of data conveyed in value from the file associated with the file
// descriptor srcFd to the value.Info destinations. See the
// ioctl_fideduperange(2) man page for details.
func IoctlFileDedupeRange(srcFd int, value *FileDedupeRange) error {
	buf := make([]byte, SizeofRawFileDedupeRange+
		len(value.Info)*SizeofRawFileDedupeRangeInfo)
	rawrange := (*RawFileDedupeRange)(unsafe.Pointer(&buf[0]))
	rawrange.Src_offset = value.Src_offset
	rawrange.Src_length = value.Src_length
	rawrange.Dest_count = uint16(len(value.Info))
	rawrange.Reserved1 = value.Reserved1
	rawrange.Reserved2 = value.Reserved2

	for i := range value.Info {
		rawinfo := (*RawFileDedupeRangeInfo)(unsafe.Pointer(
			uintptr(unsafe.Pointer(&buf[0])) + uintptr(SizeofRawFileDedupeRange) +
				uintptr(i*SizeofRawFileDedupeRangeInfo)))
		rawinfo.Dest_fd = value.Info[i].Dest_fd
		rawinfo.Dest_offset = value.Info[i].Dest_offset
		rawinfo.Bytes_deduped = value.Info[i].Bytes_deduped
		rawinfo.Status = value.Info[i].Status
		rawinfo.Reserved = value.Info[i].Reserved
	}

	err := ioctlPtr(srcFd, FIDEDUPERANGE, unsafe.Pointer(&buf[0]))

	// Output
	for i := range value.Info {
		rawinfo := (*RawFileDedupeRangeInfo)(unsafe.Pointer(
			uintptr(unsafe.Pointer(&buf[0])) + uintptr(SizeofRawFileDedupeRange) +
				uintptr(i*SizeofRawFileDedupeRangeInfo)))
		value.Info[i].Dest_fd = rawinfo.Dest_fd
		value.Info[i].Dest_offset = rawinfo.Dest_offset
		value.Info[i].Bytes_deduped = rawinfo.Bytes_deduped
		value.Info[i].Status = rawinfo.Status
		value.Info[i].Reserved = rawinfo.Reserved
	}

	return err
}

func IoctlHIDGetDesc(fd int, value *HIDRawReportDescriptor) error {
	return ioctlPtr(fd, HIDIOCGRDESC, unsafe.Pointer(value))
}

func IoctlHIDGetRawInfo(fd int) (*HIDRawDevInfo, error) {
	var value HIDRawDevInfo
	err := ioctlPtr(fd, HIDIOCGRAWINFO, unsafe.Pointer(&value))
	return &value, err
}

func IoctlHIDGetRawName(fd int) (string, error) {
	var value [_HIDIOCGRAWNAME_LEN]byte
	err := ioctlPtr(fd, _HIDIOCGRAWNAME, unsafe.Pointer(&value[0]))
	return ByteSliceToString(value[:]), err
}

func IoctlHIDGetRawPhys(fd int) (string, error) {
	var value [_HIDIOCGRAWPHYS_LEN]byte
	err := ioctlPtr(fd, _HIDIOCGRAWPHYS, unsafe.Pointer(&value[0]))
	return ByteSliceToString(value[:]), err
}

func IoctlHIDGetRawUniq(fd int) (string, error) {
	var value [_HIDIOCGRAWUNIQ_LEN]byte
	err := ioctlPtr(fd, _HIDIOCGRAWUNIQ, unsafe.Pointer(&value[0]))
	return ByteSliceToString(value[:]), err
}
<|endoftext|>


<commit_before>package hypervisor

import (
	"errors"
	"net"
)

const consoleTypeUnknown = "UNKNOWN ConsoleType"
const stateUnknown = "UNKNOWN State"
const volumeFormatUnknown = "UNKNOWN VolumeFormat"

var (
	consoleTypeToText = map[ConsoleType]string{
		ConsoleNone:  "none",
		ConsoleDummy: "dummy",
		ConsoleVNC:   "vnc",
	}
	textToConsoleType map[string]ConsoleType

	stateToText = map[State]string{
		StateStarting:      "starting",
		StateRunning:       "running",
		StateFailedToStart: "failed to start",
		StateStopping:      "stopping",
		StateStopped:       "stopped",
		StateDestroying:    "destroying",
		StateMigrating:     "migrating",
		StateExporting:     "exporting",
	}
	textToState map[string]State

	volumeFormatToText = map[VolumeFormat]string{
		VolumeFormatRaw:   "raw",
		VolumeFormatQCOW2: "qcow2",
	}
	textToVolumeFormat map[string]VolumeFormat
)

func init() {
	textToConsoleType = make(map[string]ConsoleType, len(consoleTypeToText))
	for consoleType, text := range consoleTypeToText {
		textToConsoleType[text] = consoleType
	}
	textToState = make(map[string]State, len(stateToText))
	for state, text := range stateToText {
		textToState[text] = state
	}
	textToVolumeFormat = make(map[string]VolumeFormat, len(volumeFormatToText))
	for format, text := range volumeFormatToText {
		textToVolumeFormat[text] = format
	}
}

func shrinkIP(netIP net.IP) net.IP {
	switch len(netIP) {
	case 4:
		return netIP
	case 16:
		if ip4 := netIP.To4(); ip4 == nil {
			return netIP
		} else {
			return ip4
		}
	default:
		return netIP
	}
}

func (left *Address) Equal(right *Address) bool {
	if !left.IpAddress.Equal(right.IpAddress) {
		return false
	}
	if left.MacAddress != right.MacAddress {
		return false
	}
	return true
}

func (address *Address) Shrink() {
	address.IpAddress = shrinkIP(address.IpAddress)
}

func stringSlicesEqual(left, right []string) bool {
	if len(left) != len(right) {
		return false
	}
	for index, leftString := range left {
		if leftString != right[index] {
			return false
		}
	}
	return true
}

func (consoleType *ConsoleType) CheckValid() error {
	if _, ok := consoleTypeToText[*consoleType]; !ok {
		return errors.New(consoleTypeUnknown)
	} else {
		return nil
	}
}

func (consoleType ConsoleType) MarshalText() ([]byte, error) {
	if text := consoleType.String(); text == consoleTypeUnknown {
		return nil, errors.New(text)
	} else {
		return []byte(text), nil
	}
}

func (consoleType *ConsoleType) Set(value string) error {
	if val, ok := textToConsoleType[value]; !ok {
		return errors.New(consoleTypeUnknown)
	} else {
		*consoleType = val
		return nil
	}
}

func (consoleType ConsoleType) String() string {
	if str, ok := consoleTypeToText[consoleType]; !ok {
		return consoleTypeUnknown
	} else {
		return str
	}
}

func (consoleType *ConsoleType) UnmarshalText(text []byte) error {
	txt := string(text)
	if val, ok := textToConsoleType[txt]; ok {
		*consoleType = val
		return nil
	} else {
		return errors.New("unknown ConsoleType: " + txt)
	}
}

func (state State) MarshalText() ([]byte, error) {
	if text := state.String(); text == stateUnknown {
		return nil, errors.New(text)
	} else {
		return []byte(text), nil
	}
}

func (state State) String() string {
	if text, ok := stateToText[state]; ok {
		return text
	} else {
		return stateUnknown
	}
}

func (state *State) UnmarshalText(text []byte) error {
	txt := string(text)
	if val, ok := textToState[txt]; ok {
		*state = val
		return nil
	} else {
		return errors.New("unknown State: " + txt)
	}
}

func (left *Subnet) Equal(right *Subnet) bool {
	if left.Id != right.Id {
		return false
	}
	if !left.IpGateway.Equal(right.IpGateway) {
		return false
	}
	if !left.IpMask.Equal(right.IpMask) {
		return false
	}
	if left.DomainName != right.DomainName {
		return false
	}
	if !IpListsEqual(left.DomainNameServers, right.DomainNameServers) {
		return false
	}
	if left.Manage != right.Manage {
		return false
	}
	if left.VlanId != right.VlanId {
		return false
	}
	if !stringSlicesEqual(left.AllowedGroups, right.AllowedGroups) {
		return false
	}
	if !stringSlicesEqual(left.AllowedUsers, right.AllowedUsers) {
		return false
	}
	return true
}

func IpListsEqual(left, right []net.IP) bool {
	if len(left) != len(right) {
		return false
	}
	for index, leftIP := range left {
		if !leftIP.Equal(right[index]) {
			return false
		}
	}
	return true
}

func (subnet *Subnet) Shrink() {
	subnet.IpGateway = shrinkIP(subnet.IpGateway)
	subnet.IpMask = shrinkIP(subnet.IpMask)
	for index, ipAddr := range subnet.DomainNameServers {
		subnet.DomainNameServers[index] = shrinkIP(ipAddr)
	}
}

func (left *VmInfo) Equal(right *VmInfo) bool {
	if !left.Address.Equal(&right.Address) {
		return false
	}
	if left.ConsoleType != right.ConsoleType {
		return false
	}
	if left.DestroyProtection != right.DestroyProtection {
		return false
	}
	if left.DisableVirtIO != right.DisableVirtIO {
		return false
	}
	if left.Hostname != right.Hostname {
		return false
	}
	if left.ImageName != right.ImageName {
		return false
	}
	if left.ImageURL != right.ImageURL {
		return false
	}
	if left.MemoryInMiB != right.MemoryInMiB {
		return false
	}
	if left.MilliCPUs != right.MilliCPUs {
		return false
	}
	if !stringSlicesEqual(left.OwnerGroups, right.OwnerGroups) {
		return false
	}
	if !stringSlicesEqual(left.OwnerUsers, right.OwnerUsers) {
		return false
	}
	if left.SpreadVolumes != right.SpreadVolumes {
		return false
	}
	if left.State != right.State {
		return false
	}
	if !left.Tags.Equal(right.Tags) {
		return false
	}
	if len(left.SecondaryAddresses) != len(right.SecondaryAddresses) {
		return false
	}
	for index, leftAddress := range left.SecondaryAddresses {
		if !leftAddress.Equal(&right.SecondaryAddresses[index]) {
			return false
		}
	}
	if !stringSlicesEqual(left.SecondarySubnetIDs, right.SecondarySubnetIDs) {
		return false
	}
	if left.SubnetId != right.SubnetId {
		return false
	}
	if left.Uncommitted != right.Uncommitted {
		return false
	}
	if len(left.Volumes) != len(right.Volumes) {
		return false
	}
	for index, leftVolume := range left.Volumes {
		if leftVolume != right.Volumes[index] {
			return false
		}
	}
	return true
}

func (volumeFormat VolumeFormat) MarshalText() ([]byte, error) {
	if text := volumeFormat.String(); text == volumeFormatUnknown {
		return nil, errors.New(text)
	} else {
		return []byte(text), nil
	}
}

func (volumeFormat VolumeFormat) String() string {
	if text, ok := volumeFormatToText[volumeFormat]; ok {
		return text
	} else {
		return volumeFormatUnknown
	}
}

func (volumeFormat *VolumeFormat) UnmarshalText(text []byte) error {
	txt := string(text)
	if val, ok := textToVolumeFormat[txt]; ok {
		*volumeFormat = val
		return nil
	} else {
		return errors.New("unknown VolumeFormat: " + txt)
	}
}
<commit_msg>Export proto/hypervisor.ShrinkIP() function.<commit_after>package hypervisor

import (
	"errors"
	"net"
)

const consoleTypeUnknown = "UNKNOWN ConsoleType"
const stateUnknown = "UNKNOWN State"
const volumeFormatUnknown = "UNKNOWN VolumeFormat"

var (
	consoleTypeToText = map[ConsoleType]string{
		ConsoleNone:  "none",
		ConsoleDummy: "dummy",
		ConsoleVNC:   "vnc",
	}
	textToConsoleType map[string]ConsoleType

	stateToText = map[State]string{
		StateStarting:      "starting",
		StateRunning:       "running",
		StateFailedToStart: "failed to start",
		StateStopping:      "stopping",
		StateStopped:       "stopped",
		StateDestroying:    "destroying",
		StateMigrating:     "migrating",
		StateExporting:     "exporting",
	}
	textToState map[string]State

	volumeFormatToText = map[VolumeFormat]string{
		VolumeFormatRaw:   "raw",
		VolumeFormatQCOW2: "qcow2",
	}
	textToVolumeFormat map[string]VolumeFormat
)

func init() {
	textToConsoleType = make(map[string]ConsoleType, len(consoleTypeToText))
	for consoleType, text := range consoleTypeToText {
		textToConsoleType[text] = consoleType
	}
	textToState = make(map[string]State, len(stateToText))
	for state, text := range stateToText {
		textToState[text] = state
	}
	textToVolumeFormat = make(map[string]VolumeFormat, len(volumeFormatToText))
	for format, text := range volumeFormatToText {
		textToVolumeFormat[text] = format
	}
}

func ShrinkIP(netIP net.IP) net.IP {
	switch len(netIP) {
	case 4:
		return netIP
	case 16:
		if ip4 := netIP.To4(); ip4 == nil {
			return netIP
		} else {
			return ip4
		}
	default:
		return netIP
	}
}

func (left *Address) Equal(right *Address) bool {
	if !left.IpAddress.Equal(right.IpAddress) {
		return false
	}
	if left.MacAddress != right.MacAddress {
		return false
	}
	return true
}

func (address *Address) Shrink() {
	address.IpAddress = ShrinkIP(address.IpAddress)
}

func stringSlicesEqual(left, right []string) bool {
	if len(left) != len(right) {
		return false
	}
	for index, leftString := range left {
		if leftString != right[index] {
			return false
		}
	}
	return true
}

func (consoleType *ConsoleType) CheckValid() error {
	if _, ok := consoleTypeToText[*consoleType]; !ok {
		return errors.New(consoleTypeUnknown)
	} else {
		return nil
	}
}

func (consoleType ConsoleType) MarshalText() ([]byte, error) {
	if text := consoleType.String(); text == consoleTypeUnknown {
		return nil, errors.New(text)
	} else {
		return []byte(text), nil
	}
}

func (consoleType *ConsoleType) Set(value string) error {
	if val, ok := textToConsoleType[value]; !ok {
		return errors.New(consoleTypeUnknown)
	} else {
		*consoleType = val
		return nil
	}
}

func (consoleType ConsoleType) String() string {
	if str, ok := consoleTypeToText[consoleType]; !ok {
		return consoleTypeUnknown
	} else {
		return str
	}
}

func (consoleType *ConsoleType) UnmarshalText(text []byte) error {
	txt := string(text)
	if val, ok := textToConsoleType[txt]; ok {
		*consoleType = val
		return nil
	} else {
		return errors.New("unknown ConsoleType: " + txt)
	}
}

func (state State) MarshalText() ([]byte, error) {
	if text := state.String(); text == stateUnknown {
		return nil, errors.New(text)
	} else {
		return []byte(text), nil
	}
}

func (state State) String() string {
	if text, ok := stateToText[state]; ok {
		return text
	} else {
		return stateUnknown
	}
}

func (state *State) UnmarshalText(text []byte) error {
	txt := string(text)
	if val, ok := textToState[txt]; ok {
		*state = val
		return nil
	} else {
		return errors.New("unknown State: " + txt)
	}
}

func (left *Subnet) Equal(right *Subnet) bool {
	if left.Id != right.Id {
		return false
	}
	if !left.IpGateway.Equal(right.IpGateway) {
		return false
	}
	if !left.IpMask.Equal(right.IpMask) {
		return false
	}
	if left.DomainName != right.DomainName {
		return false
	}
	if !IpListsEqual(left.DomainNameServers, right.DomainNameServers) {
		return false
	}
	if left.Manage != right.Manage {
		return false
	}
	if left.VlanId != right.VlanId {
		return false
	}
	if !stringSlicesEqual(left.AllowedGroups, right.AllowedGroups) {
		return false
	}
	if !stringSlicesEqual(left.AllowedUsers, right.AllowedUsers) {
		return false
	}
	return true
}

func IpListsEqual(left, right []net.IP) bool {
	if len(left) != len(right) {
		return false
	}
	for index, leftIP := range left {
		if !leftIP.Equal(right[index]) {
			return false
		}
	}
	return true
}

func (subnet *Subnet) Shrink() {
	subnet.IpGateway = ShrinkIP(subnet.IpGateway)
	subnet.IpMask = ShrinkIP(subnet.IpMask)
	for index, ipAddr := range subnet.DomainNameServers {
		subnet.DomainNameServers[index] = ShrinkIP(ipAddr)
	}
}

func (left *VmInfo) Equal(right *VmInfo) bool {
	if !left.Address.Equal(&right.Address) {
		return false
	}
	if left.ConsoleType != right.ConsoleType {
		return false
	}
	if left.DestroyProtection != right.DestroyProtection {
		return false
	}
	if left.DisableVirtIO != right.DisableVirtIO {
		return false
	}
	if left.Hostname != right.Hostname {
		return false
	}
	if left.ImageName != right.ImageName {
		return false
	}
	if left.ImageURL != right.ImageURL {
		return false
	}
	if left.MemoryInMiB != right.MemoryInMiB {
		return false
	}
	if left.MilliCPUs != right.MilliCPUs {
		return false
	}
	if !stringSlicesEqual(left.OwnerGroups, right.OwnerGroups) {
		return false
	}
	if !stringSlicesEqual(left.OwnerUsers, right.OwnerUsers) {
		return false
	}
	if left.SpreadVolumes != right.SpreadVolumes {
		return false
	}
	if left.State != right.State {
		return false
	}
	if !left.Tags.Equal(right.Tags) {
		return false
	}
	if len(left.SecondaryAddresses) != len(right.SecondaryAddresses) {
		return false
	}
	for index, leftAddress := range left.SecondaryAddresses {
		if !leftAddress.Equal(&right.SecondaryAddresses[index]) {
			return false
		}
	}
	if !stringSlicesEqual(left.SecondarySubnetIDs, right.SecondarySubnetIDs) {
		return false
	}
	if left.SubnetId != right.SubnetId {
		return false
	}
	if left.Uncommitted != right.Uncommitted {
		return false
	}
	if len(left.Volumes) != len(right.Volumes) {
		return false
	}
	for index, leftVolume := range left.Volumes {
		if leftVolume != right.Volumes[index] {
			return false
		}
	}
	return true
}

func (volumeFormat VolumeFormat) MarshalText() ([]byte, error) {
	if text := volumeFormat.String(); text == volumeFormatUnknown {
		return nil, errors.New(text)
	} else {
		return []byte(text), nil
	}
}

func (volumeFormat VolumeFormat) String() string {
	if text, ok := volumeFormatToText[volumeFormat]; ok {
		return text
	} else {
		return volumeFormatUnknown
	}
}

func (volumeFormat *VolumeFormat) UnmarshalText(text []byte) error {
	txt := string(text)
	if val, ok := textToVolumeFormat[txt]; ok {
		*volumeFormat = val
		return nil
	} else {
		return errors.New("unknown VolumeFormat: " + txt)
	}
}
<|endoftext|>


<commit_before>package shadowsocks

import (
	"context"
	"runtime"
	"time"

	"v2ray.com/core/app/log"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/retry"
	"v2ray.com/core/common/signal"
	"v2ray.com/core/proxy"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/ray"
)

// Client is a inbound handler for Shadowsocks protocol
type Client struct {
	serverPicker protocol.ServerPicker
}

// NewClient create a new Shadowsocks client.
func NewClient(ctx context.Context, config *ClientConfig) (*Client, error) {
	serverList := protocol.NewServerList()
	for _, rec := range config.Server {
		serverList.AddServer(protocol.NewServerSpecFromPB(*rec))
	}
	client := &Client{
		serverPicker: protocol.NewRoundRobinServerPicker(serverList),
	}

	return client, nil
}

// Process implements OutboundHandler.Process().
func (v *Client) Process(ctx context.Context, outboundRay ray.OutboundRay, dialer proxy.Dialer) error {
	destination, ok := proxy.TargetFromContext(ctx)
	if !ok {
		return newError("target not specified")
	}
	network := destination.Network

	var server *protocol.ServerSpec
	var conn internet.Connection

	err := retry.ExponentialBackoff(5, 100).On(func() error {
		server = v.serverPicker.PickServer()
		dest := server.Destination()
		dest.Network = network
		rawConn, err := dialer.Dial(ctx, dest)
		if err != nil {
			return err
		}
		conn = rawConn

		return nil
	})
	if err != nil {
		return newError("failed to find an available destination").AtWarning().Base(err)
	}
	log.Trace(newError("tunneling request to ", destination, " via ", server.Destination()))

	defer conn.Close()

	request := &protocol.RequestHeader{
		Version: Version,
		Address: destination.Address,
		Port:    destination.Port,
	}
	if destination.Network == net.Network_TCP {
		request.Command = protocol.RequestCommandTCP
	} else {
		request.Command = protocol.RequestCommandUDP
	}

	user := server.PickUser()
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		return newError("failed to get a valid user account").AtWarning().Base(err)
	}
	account := rawAccount.(*ShadowsocksAccount)
	request.User = user

	if account.OneTimeAuth == Account_Auto || account.OneTimeAuth == Account_Enabled {
		request.Option |= RequestOptionOneTimeAuth
	}

	ctx, timer := signal.CancelAfterInactivity(ctx, time.Minute*2)

	if request.Command == protocol.RequestCommandTCP {
		bufferedWriter := buf.NewBufferedWriter(conn)
		bodyWriter, err := WriteTCPRequest(request, bufferedWriter)
		if err != nil {
			return newError("failed to write request").Base(err)
		}

		if err := bufferedWriter.SetBuffered(false); err != nil {
			return err
		}

		requestDone := signal.ExecuteAsync(func() error {
			mergedInput := buf.NewMergingReader(outboundRay.OutboundInput())
			if err := buf.PipeUntilEOF(timer, mergedInput, bodyWriter); err != nil {
				return err
			}
			return nil
		})

		responseDone := signal.ExecuteAsync(func() error {
			defer outboundRay.OutboundOutput().Close()

			responseReader, err := ReadTCPResponse(user, conn)
			if err != nil {
				return err
			}

			if err := buf.PipeUntilEOF(timer, responseReader, outboundRay.OutboundOutput()); err != nil {
				return err
			}

			return nil
		})

		if err := signal.ErrorOrFinish2(ctx, requestDone, responseDone); err != nil {
			return newError("connection ends").Base(err)
		}

		return nil
	}

	if request.Command == protocol.RequestCommandUDP {

		writer := &UDPWriter{
			Writer:  conn,
			Request: request,
		}

		requestDone := signal.ExecuteAsync(func() error {
			if err := buf.PipeUntilEOF(timer, outboundRay.OutboundInput(), writer); err != nil {
				return newError("failed to transport all UDP request").Base(err)
			}
			return nil
		})

		responseDone := signal.ExecuteAsync(func() error {
			defer outboundRay.OutboundOutput().Close()

			reader := &UDPReader{
				Reader: conn,
				User:   user,
			}

			if err := buf.PipeUntilEOF(timer, reader, outboundRay.OutboundOutput()); err != nil {
				return newError("failed to transport all UDP response").Base(err)
			}
			return nil
		})

		if err := signal.ErrorOrFinish2(ctx, requestDone, responseDone); err != nil {
			return newError("connection ends").Base(err)
		}

		return nil
	}

	runtime.KeepAlive(timer)

	return nil
}

func init() {
	common.Must(common.RegisterConfig((*ClientConfig)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
		return NewClient(ctx, config.(*ClientConfig))
	}))
}
<commit_msg>check number of shadowsocks servers<commit_after>package shadowsocks

import (
	"context"
	"runtime"
	"time"

	"v2ray.com/core/app/log"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/retry"
	"v2ray.com/core/common/signal"
	"v2ray.com/core/proxy"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/ray"
)

// Client is a inbound handler for Shadowsocks protocol
type Client struct {
	serverPicker protocol.ServerPicker
}

// NewClient create a new Shadowsocks client.
func NewClient(ctx context.Context, config *ClientConfig) (*Client, error) {
	serverList := protocol.NewServerList()
	for _, rec := range config.Server {
		serverList.AddServer(protocol.NewServerSpecFromPB(*rec))
	}
	if serverList.Size() == 0 {
		return nil, newError("0 server")
	}
	client := &Client{
		serverPicker: protocol.NewRoundRobinServerPicker(serverList),
	}

	return client, nil
}

// Process implements OutboundHandler.Process().
func (v *Client) Process(ctx context.Context, outboundRay ray.OutboundRay, dialer proxy.Dialer) error {
	destination, ok := proxy.TargetFromContext(ctx)
	if !ok {
		return newError("target not specified")
	}
	network := destination.Network

	var server *protocol.ServerSpec
	var conn internet.Connection

	err := retry.ExponentialBackoff(5, 100).On(func() error {
		server = v.serverPicker.PickServer()
		dest := server.Destination()
		dest.Network = network
		rawConn, err := dialer.Dial(ctx, dest)
		if err != nil {
			return err
		}
		conn = rawConn

		return nil
	})
	if err != nil {
		return newError("failed to find an available destination").AtWarning().Base(err)
	}
	log.Trace(newError("tunneling request to ", destination, " via ", server.Destination()))

	defer conn.Close()

	request := &protocol.RequestHeader{
		Version: Version,
		Address: destination.Address,
		Port:    destination.Port,
	}
	if destination.Network == net.Network_TCP {
		request.Command = protocol.RequestCommandTCP
	} else {
		request.Command = protocol.RequestCommandUDP
	}

	user := server.PickUser()
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		return newError("failed to get a valid user account").AtWarning().Base(err)
	}
	account := rawAccount.(*ShadowsocksAccount)
	request.User = user

	if account.OneTimeAuth == Account_Auto || account.OneTimeAuth == Account_Enabled {
		request.Option |= RequestOptionOneTimeAuth
	}

	ctx, timer := signal.CancelAfterInactivity(ctx, time.Minute*2)

	if request.Command == protocol.RequestCommandTCP {
		bufferedWriter := buf.NewBufferedWriter(conn)
		bodyWriter, err := WriteTCPRequest(request, bufferedWriter)
		if err != nil {
			return newError("failed to write request").Base(err)
		}

		if err := bufferedWriter.SetBuffered(false); err != nil {
			return err
		}

		requestDone := signal.ExecuteAsync(func() error {
			mergedInput := buf.NewMergingReader(outboundRay.OutboundInput())
			if err := buf.PipeUntilEOF(timer, mergedInput, bodyWriter); err != nil {
				return err
			}
			return nil
		})

		responseDone := signal.ExecuteAsync(func() error {
			defer outboundRay.OutboundOutput().Close()

			responseReader, err := ReadTCPResponse(user, conn)
			if err != nil {
				return err
			}

			if err := buf.PipeUntilEOF(timer, responseReader, outboundRay.OutboundOutput()); err != nil {
				return err
			}

			return nil
		})

		if err := signal.ErrorOrFinish2(ctx, requestDone, responseDone); err != nil {
			return newError("connection ends").Base(err)
		}

		return nil
	}

	if request.Command == protocol.RequestCommandUDP {

		writer := &UDPWriter{
			Writer:  conn,
			Request: request,
		}

		requestDone := signal.ExecuteAsync(func() error {
			if err := buf.PipeUntilEOF(timer, outboundRay.OutboundInput(), writer); err != nil {
				return newError("failed to transport all UDP request").Base(err)
			}
			return nil
		})

		responseDone := signal.ExecuteAsync(func() error {
			defer outboundRay.OutboundOutput().Close()

			reader := &UDPReader{
				Reader: conn,
				User:   user,
			}

			if err := buf.PipeUntilEOF(timer, reader, outboundRay.OutboundOutput()); err != nil {
				return newError("failed to transport all UDP response").Base(err)
			}
			return nil
		})

		if err := signal.ErrorOrFinish2(ctx, requestDone, responseDone); err != nil {
			return newError("connection ends").Base(err)
		}

		return nil
	}

	runtime.KeepAlive(timer)

	return nil
}

func init() {
	common.Must(common.RegisterConfig((*ClientConfig)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
		return NewClient(ctx, config.(*ClientConfig))
	}))
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"regexp"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/organizations"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/validation"
)

func resourceAwsOrganizationsAccount() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsOrganizationsAccountCreate,
		Read:   resourceAwsOrganizationsAccountRead,
		Delete: resourceAwsOrganizationsAccountDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"arn": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"joined_method": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"joined_timestamp": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"status": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"name": {
				Type:         schema.TypeString,
				Required:     true,
				ValidateFunc: validation.StringLenBetween(1, 50),
			},
			"email": {
				Type:         schema.TypeString,
				Required:     true,
				ValidateFunc: validateAwsOrganizationsAccountEmail,
			},
			"iam_user_access_to_billing": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validation.StringInSlice([]string{organizations.IAMUserAccessToBillingAllow, organizations.IAMUserAccessToBillingDeny}, true),
			},
			"role_name": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateAwsOrganizationsAccountRoleName,
			},
		},
	}
}

func resourceAwsOrganizationsAccountCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).organizationsconn

	// Create the account
	createOpts := &organizations.CreateAccountInput{
		AccountName: aws.String(d.Get("name").(string)),
		Email:       aws.String(d.Get("email").(string)),
	}
	if role, ok := d.GetOk("role_name"); ok {
		createOpts.RoleName = aws.String(role.(string))
	}

	if iam_user, ok := d.GetOk("iam_user_access_to_billing"); ok {
		createOpts.IamUserAccessToBilling = aws.String(iam_user.(string))
	}

	log.Printf("[DEBUG] Account create config: %#v", createOpts)

	var err error
	var resp *organizations.CreateAccountOutput
	err = resource.Retry(4*time.Minute, func() *resource.RetryError {
		resp, err = conn.CreateAccount(createOpts)

		if err != nil {
			ec2err, ok := err.(awserr.Error)
			if !ok {
				return resource.NonRetryableError(err)
			}
			if isAWSErr(ec2err, organizations.ErrCodeFinalizingOrganizationException, "") {
				log.Printf("[DEBUG] Trying to create account again: %q", ec2err.Message())
				return resource.RetryableError(err)
			}

			return resource.NonRetryableError(err)
		}

		return nil
	})

	if err != nil {
		return fmt.Errorf("Error creating account: %s", err)
	}
	log.Printf("[DEBUG] Account create response: %#v", resp)

	requestId := *resp.CreateAccountStatus.Id

	// Wait for the account to become available
	log.Printf("[DEBUG] Waiting for account request (%s) to succeed", requestId)

	stateConf := &resource.StateChangeConf{
		Pending:      []string{organizations.CreateAccountStateInProgress},
		Target:       []string{organizations.CreateAccountStateSucceeded},
		Refresh:      resourceAwsOrganizationsAccountStateRefreshFunc(conn, requestId),
		PollInterval: 10 * time.Second,
		Timeout:      5 * time.Minute,
	}
	stateResp, stateErr := stateConf.WaitForState()
	if stateErr != nil {
		return fmt.Errorf(
			"Error waiting for account request (%s) to become available: %s",
			requestId, stateErr)
	}

	// Store the ID
	accountId := stateResp.(*organizations.CreateAccountStatus).AccountId
	d.SetId(*accountId)

	return resourceAwsOrganizationsAccountRead(d, meta)
}

func resourceAwsOrganizationsAccountRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).organizationsconn
	describeOpts := &organizations.DescribeAccountInput{
		AccountId: aws.String(d.Id()),
	}
	resp, err := conn.DescribeAccount(describeOpts)
	if err != nil {
		if isAWSErr(err, organizations.ErrCodeAccountNotFoundException, "") {
			log.Printf("[WARN] Account does not exist, removing from state: %s", d.Id())
			d.SetId("")
			return nil
		}
		return err
	}

	d.Set("arn", resp.Account.Arn)
	d.Set("joined_method", resp.Account.JoinedMethod)
	d.Set("joined_timestamp", resp.Account.JoinedTimestamp)
	d.Set("name", resp.Account.Name)
	d.Set("status", resp.Account.Status)
	return nil
}

func resourceAwsOrganizationsAccountDelete(d *schema.ResourceData, meta interface{}) error {
	log.Printf("[WARN] Organization member accounts must be deleted from the web console.")
	return nil
}

// resourceAwsOrganizationsAccountStateRefreshFunc returns a resource.StateRefreshFunc
// that is used to watch a CreateAccount request
func resourceAwsOrganizationsAccountStateRefreshFunc(conn *organizations.Organizations, id string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		opts := &organizations.DescribeCreateAccountStatusInput{
			CreateAccountRequestId: aws.String(id),
		}
		resp, err := conn.DescribeCreateAccountStatus(opts)
		if err != nil {
			if isAWSErr(err, organizations.ErrCodeCreateAccountStatusNotFoundException, "") {
				resp = nil
			} else {
				log.Printf("Error on OrganizationAccountStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil {
			// Sometimes AWS just has consistency issues and doesn't see
			// our account yet. Return an empty state.
			return nil, "", nil
		}

		accountStatus := resp.CreateAccountStatus
		if *accountStatus.State == organizations.CreateAccountStateFailed {
			return nil, *accountStatus.State, fmt.Errorf(*accountStatus.FailureReason)
		}
		return accountStatus, *accountStatus.State, nil
	}
}

func validateAwsOrganizationsAccountEmail(v interface{}, k string) (ws []string, errors []error) {
	value := v.(string)
	if !regexp.MustCompile(`^[^\s@]+@[^\s@]+\.[^\s@]+$`).MatchString(value) {
		errors = append(errors, fmt.Errorf(
			"%q must be a valid email address", value))
	}

	if len(value) < 6 {
		errors = append(errors, fmt.Errorf(
			"%q cannot be less than 6 characters", value))
	}

	if len(value) > 64 {
		errors = append(errors, fmt.Errorf(
			"%q cannot be greater than 64 characters", value))
	}

	return
}

func validateAwsOrganizationsAccountRoleName(v interface{}, k string) (ws []string, errors []error) {
	value := v.(string)
	if !regexp.MustCompile(`^[\w+=,.@-]{1,64}$`).MatchString(value) {
		errors = append(errors, fmt.Errorf(
			"%q must consist of uppercase letters, lowercase letters, digits with no spaces, and any of the following characters: =,.@-", value))
	}

	return
}
<commit_msg>Must set ForceNew:true since there is no update method defined.<commit_after>package aws

import (
	"fmt"
	"log"
	"regexp"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/organizations"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/validation"
)

func resourceAwsOrganizationsAccount() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsOrganizationsAccountCreate,
		Read:   resourceAwsOrganizationsAccountRead,
		Delete: resourceAwsOrganizationsAccountDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"arn": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"joined_method": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"joined_timestamp": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"status": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"name": {
				ForceNew:     true,
				Type:         schema.TypeString,
				Required:     true,
				ValidateFunc: validation.StringLenBetween(1, 50),
			},
			"email": {
				ForceNew:     true,
				Type:         schema.TypeString,
				Required:     true,
				ValidateFunc: validateAwsOrganizationsAccountEmail,
			},
			"iam_user_access_to_billing": {
				ForceNew:     true,
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validation.StringInSlice([]string{organizations.IAMUserAccessToBillingAllow, organizations.IAMUserAccessToBillingDeny}, true),
			},
			"role_name": {
				ForceNew:     true,
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateAwsOrganizationsAccountRoleName,
			},
		},
	}
}

func resourceAwsOrganizationsAccountCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).organizationsconn

	// Create the account
	createOpts := &organizations.CreateAccountInput{
		AccountName: aws.String(d.Get("name").(string)),
		Email:       aws.String(d.Get("email").(string)),
	}
	if role, ok := d.GetOk("role_name"); ok {
		createOpts.RoleName = aws.String(role.(string))
	}

	if iam_user, ok := d.GetOk("iam_user_access_to_billing"); ok {
		createOpts.IamUserAccessToBilling = aws.String(iam_user.(string))
	}

	log.Printf("[DEBUG] Account create config: %#v", createOpts)

	var err error
	var resp *organizations.CreateAccountOutput
	err = resource.Retry(4*time.Minute, func() *resource.RetryError {
		resp, err = conn.CreateAccount(createOpts)

		if err != nil {
			ec2err, ok := err.(awserr.Error)
			if !ok {
				return resource.NonRetryableError(err)
			}
			if isAWSErr(ec2err, organizations.ErrCodeFinalizingOrganizationException, "") {
				log.Printf("[DEBUG] Trying to create account again: %q", ec2err.Message())
				return resource.RetryableError(err)
			}

			return resource.NonRetryableError(err)
		}

		return nil
	})

	if err != nil {
		return fmt.Errorf("Error creating account: %s", err)
	}
	log.Printf("[DEBUG] Account create response: %#v", resp)

	requestId := *resp.CreateAccountStatus.Id

	// Wait for the account to become available
	log.Printf("[DEBUG] Waiting for account request (%s) to succeed", requestId)

	stateConf := &resource.StateChangeConf{
		Pending:      []string{organizations.CreateAccountStateInProgress},
		Target:       []string{organizations.CreateAccountStateSucceeded},
		Refresh:      resourceAwsOrganizationsAccountStateRefreshFunc(conn, requestId),
		PollInterval: 10 * time.Second,
		Timeout:      5 * time.Minute,
	}
	stateResp, stateErr := stateConf.WaitForState()
	if stateErr != nil {
		return fmt.Errorf(
			"Error waiting for account request (%s) to become available: %s",
			requestId, stateErr)
	}

	// Store the ID
	accountId := stateResp.(*organizations.CreateAccountStatus).AccountId
	d.SetId(*accountId)

	return resourceAwsOrganizationsAccountRead(d, meta)
}

func resourceAwsOrganizationsAccountRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).organizationsconn
	describeOpts := &organizations.DescribeAccountInput{
		AccountId: aws.String(d.Id()),
	}
	resp, err := conn.DescribeAccount(describeOpts)
	if err != nil {
		if isAWSErr(err, organizations.ErrCodeAccountNotFoundException, "") {
			log.Printf("[WARN] Account does not exist, removing from state: %s", d.Id())
			d.SetId("")
			return nil
		}
		return err
	}

	d.Set("arn", resp.Account.Arn)
	d.Set("joined_method", resp.Account.JoinedMethod)
	d.Set("joined_timestamp", resp.Account.JoinedTimestamp)
	d.Set("name", resp.Account.Name)
	d.Set("status", resp.Account.Status)
	return nil
}

func resourceAwsOrganizationsAccountDelete(d *schema.ResourceData, meta interface{}) error {
	log.Printf("[WARN] Organization member accounts must be deleted from the web console.")
	return nil
}

// resourceAwsOrganizationsAccountStateRefreshFunc returns a resource.StateRefreshFunc
// that is used to watch a CreateAccount request
func resourceAwsOrganizationsAccountStateRefreshFunc(conn *organizations.Organizations, id string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		opts := &organizations.DescribeCreateAccountStatusInput{
			CreateAccountRequestId: aws.String(id),
		}
		resp, err := conn.DescribeCreateAccountStatus(opts)
		if err != nil {
			if isAWSErr(err, organizations.ErrCodeCreateAccountStatusNotFoundException, "") {
				resp = nil
			} else {
				log.Printf("Error on OrganizationAccountStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil {
			// Sometimes AWS just has consistency issues and doesn't see
			// our account yet. Return an empty state.
			return nil, "", nil
		}

		accountStatus := resp.CreateAccountStatus
		if *accountStatus.State == organizations.CreateAccountStateFailed {
			return nil, *accountStatus.State, fmt.Errorf(*accountStatus.FailureReason)
		}
		return accountStatus, *accountStatus.State, nil
	}
}

func validateAwsOrganizationsAccountEmail(v interface{}, k string) (ws []string, errors []error) {
	value := v.(string)
	if !regexp.MustCompile(`^[^\s@]+@[^\s@]+\.[^\s@]+$`).MatchString(value) {
		errors = append(errors, fmt.Errorf(
			"%q must be a valid email address", value))
	}

	if len(value) < 6 {
		errors = append(errors, fmt.Errorf(
			"%q cannot be less than 6 characters", value))
	}

	if len(value) > 64 {
		errors = append(errors, fmt.Errorf(
			"%q cannot be greater than 64 characters", value))
	}

	return
}

func validateAwsOrganizationsAccountRoleName(v interface{}, k string) (ws []string, errors []error) {
	value := v.(string)
	if !regexp.MustCompile(`^[\w+=,.@-]{1,64}$`).MatchString(value) {
		errors = append(errors, fmt.Errorf(
			"%q must consist of uppercase letters, lowercase letters, digits with no spaces, and any of the following characters: =,.@-", value))
	}

	return
}
<|endoftext|>


<commit_before>package postgis_test

import (
	"log"
	"testing"

	"github.com/terranodo/tegola"
	"github.com/terranodo/tegola/provider/postgis"
)

func TestNewProvider(t *testing.T) {
	// The database connection string have the following JSON format:
	// { "host" : "host", port
	config := map[string]interface{}{
		postgis.ConfigKeyHost:     "localhost",
		postgis.ConfigKeyPort:     5432,
		postgis.ConfigKeyDB:       "gdey",
		postgis.ConfigKeyUser:     "gdey",
		postgis.ConfigKeyPassword: "",
		postgis.ConfigKeyLayers: map[string]map[string]interface{}{
			"buildings": map[string]interface{}{
				postgis.ConfigKeyTablename: "gis.zoning_base_3857",
			},
		},
	}
	p, err := postgis.NewProvider(config)
	if err != nil {
		t.Errorf("Failed to create a new provider. %v", err)
		return
	}

	tile := tegola.Tile{
		Z: 15,
		X: 12451,
		Y: 18527,
	}
	l, err := p.MVTLayer("buildings", tile, map[string]interface{}{"class": "park"})
	if err != nil {
		t.Errorf("Failed to create mvt layer. %v", err)
		return
	}
	log.Printf("Go to following layer %v\n", l)
}
<commit_msg>Added guard env for now.<commit_after>package postgis_test

import (
	"log"
	"os"
	"testing"

	"github.com/terranodo/tegola"
	"github.com/terranodo/tegola/provider/postgis"
)

func TestNewProvider(t *testing.T) {
	// The database connection string have the following JSON format:
	// { "host" : "host", port
	if os.Getenv("RUN_POSTGRESS_TEST") == "" {
		return
	}

	config := map[string]interface{}{
		postgis.ConfigKeyHost:     "localhost",
		postgis.ConfigKeyPort:     int64(5432),
		postgis.ConfigKeyDB:       "gdey",
		postgis.ConfigKeyUser:     "gdey",
		postgis.ConfigKeyPassword: "",
		postgis.ConfigKeyLayers: map[string]map[string]interface{}{
			"buildings": map[string]interface{}{
				postgis.ConfigKeyTablename: "gis.zoning_base_3857",
			},
		},
	}
	p, err := postgis.NewProvider(config)
	if err != nil {
		t.Errorf("Failed to create a new provider. %v", err)
		return
	}

	tile := tegola.Tile{
		Z: 15,
		X: 12451,
		Y: 18527,
	}
	l, err := p.MVTLayer("buildings", tile, map[string]interface{}{"class": "park"})
	if err != nil {
		t.Errorf("Failed to create mvt layer. %v", err)
		return
	}
	log.Printf("Go to following layer %v\n", l)
}
<|endoftext|>


<commit_before>package mesosreporter

import (
	"log"
	"os"
	"path/filepath"
	"strconv"

	"github.com/dropbox/changes-client/client"
	"github.com/dropbox/changes-client/client/adapter"
	"github.com/dropbox/changes-client/client/reporter"
)

// A reporter that connects and reports to a specific jobstep id.
// Each jobstep id has a number of endpoints associated with it that
// allows the reporter to update the status of logs, snapshots, etc.
type Reporter struct {
	reporter.DefaultReporter
}

func (r *Reporter) PushJobstepStatus(status string, result string) {
	log.Printf("[reporter] Pushing status %s", status)
	form := make(map[string]string)
	form["status"] = status
	if len(result) > 0 {
		form["result"] = result
	}

	hostname, err := os.Hostname()
	if err == nil {
		form["node"] = hostname
	}
	r.PublishChannel <- reporter.ReportPayload{Path: r.JobstepAPIPath(), Data: form, Filename: ""}
}

func (r *Reporter) PushCommandStatus(cID string, status string, retCode int) {
	form := make(map[string]string)
	form["status"] = status
	if retCode >= 0 {
		form["return_code"] = strconv.Itoa(retCode)
	}
	r.PublishChannel <- reporter.ReportPayload{Path: "/commands/" + cID + "/", Data: form, Filename: ""}
}

func (r *Reporter) PushLogChunk(source string, payload []byte) bool {
	// logappend endpoint only works for console logs
	if source != "console" {
		return true
	}
	form := make(map[string]string)
	form["source"] = source
	form["text"] = string(payload)
	r.PublishChannel <- reporter.ReportPayload{Path: r.JobstepAPIPath() + "logappend/", Data: form, Filename: ""}
	return true
}

func (r *Reporter) PushCommandOutput(cID string, status string, retCode int, output []byte) {
	form := make(map[string]string)
	form["status"] = status
	form["output"] = string(output)
	if retCode >= 0 {
		form["return_code"] = strconv.Itoa(retCode)
	}
	r.PublishChannel <- reporter.ReportPayload{Path: "/commands/" + cID + "/", Data: form, Filename: ""}
}

func (r *Reporter) PublishArtifacts(cmd client.ConfigCmd, a adapter.Adapter, clientLog *client.Log) error {
	if len(cmd.Artifacts) == 0 {
		clientLog.Printf("==> Skipping artifact collection")
		return nil
	}

	clientLog.Printf("==> Collecting artifacts matching %s", cmd.Artifacts)

	matches, err := a.CollectArtifacts(cmd.Artifacts, clientLog)
	if err != nil {
		clientLog.Printf("==> ERROR: %s", err)
		return err
	}

	for _, artifact := range matches {
		clientLog.Printf("==> Found: %s", artifact)
	}

	return r.pushArtifacts(matches, a.GetArtifactRoot())
}

func (r *Reporter) pushArtifacts(artifacts []string, root string) error {
	// TODO: PushArtifacts is synchronous due to races with Adapter.Shutdown(), but
	// really what we'd want to do is just say "wait until channel empty, ok continue"
	var firstError error
	for _, artifact := range artifacts {
		name := artifact
		if relativePath, err := filepath.Rel(root, artifact); err == nil {
			name = relativePath
		}
		e := r.SendPayload(reporter.ReportPayload{Path: r.JobstepAPIPath() + "artifacts/",
			Data: map[string]string{"name": name}, Filename: artifact})
		if e != nil && firstError == nil {
			firstError = e
		}
	}
	return firstError
}

func New() reporter.Reporter {
	return &Reporter{}
}

func init() {
	reporter.Register("mesos", New)
}
<commit_msg>Add debug option to not push log chunks in Mesos<commit_after>package mesosreporter

import (
	"log"
	"os"
	"path/filepath"
	"strconv"

	"github.com/dropbox/changes-client/client"
	"github.com/dropbox/changes-client/client/adapter"
	"github.com/dropbox/changes-client/client/reporter"
)

// A reporter that connects and reports to a specific jobstep id.
// Each jobstep id has a number of endpoints associated with it that
// allows the reporter to update the status of logs, snapshots, etc.
type Reporter struct {
	reporter.DefaultReporter
	dontPushLogChunks bool
}

func (r *Reporter) Init(c *client.Config) {
	r.dontPushLogChunks = c.GetDebugConfigBool("mesosDontPushLogChunks", false)
	r.DefaultReporter.Init(c)
}

func (r *Reporter) PushJobstepStatus(status string, result string) {
	log.Printf("[reporter] Pushing status %s", status)
	form := make(map[string]string)
	form["status"] = status
	if len(result) > 0 {
		form["result"] = result
	}

	hostname, err := os.Hostname()
	if err == nil {
		form["node"] = hostname
	}
	r.PublishChannel <- reporter.ReportPayload{Path: r.JobstepAPIPath(), Data: form, Filename: ""}
}

func (r *Reporter) PushCommandStatus(cID string, status string, retCode int) {
	form := make(map[string]string)
	form["status"] = status
	if retCode >= 0 {
		form["return_code"] = strconv.Itoa(retCode)
	}
	r.PublishChannel <- reporter.ReportPayload{Path: "/commands/" + cID + "/", Data: form, Filename: ""}
}

func (r *Reporter) PushLogChunk(source string, payload []byte) bool {
	if r.dontPushLogChunks {
		return true
	}
	// logappend endpoint only works for console logs
	if source != "console" {
		return true
	}
	form := make(map[string]string)
	form["source"] = source
	form["text"] = string(payload)
	r.PublishChannel <- reporter.ReportPayload{Path: r.JobstepAPIPath() + "logappend/", Data: form, Filename: ""}
	return true
}

func (r *Reporter) PushCommandOutput(cID string, status string, retCode int, output []byte) {
	form := make(map[string]string)
	form["status"] = status
	form["output"] = string(output)
	if retCode >= 0 {
		form["return_code"] = strconv.Itoa(retCode)
	}
	r.PublishChannel <- reporter.ReportPayload{Path: "/commands/" + cID + "/", Data: form, Filename: ""}
}

func (r *Reporter) PublishArtifacts(cmd client.ConfigCmd, a adapter.Adapter, clientLog *client.Log) error {
	if len(cmd.Artifacts) == 0 {
		clientLog.Printf("==> Skipping artifact collection")
		return nil
	}

	clientLog.Printf("==> Collecting artifacts matching %s", cmd.Artifacts)

	matches, err := a.CollectArtifacts(cmd.Artifacts, clientLog)
	if err != nil {
		clientLog.Printf("==> ERROR: %s", err)
		return err
	}

	for _, artifact := range matches {
		clientLog.Printf("==> Found: %s", artifact)
	}

	return r.pushArtifacts(matches, a.GetArtifactRoot())
}

func (r *Reporter) pushArtifacts(artifacts []string, root string) error {
	// TODO: PushArtifacts is synchronous due to races with Adapter.Shutdown(), but
	// really what we'd want to do is just say "wait until channel empty, ok continue"
	var firstError error
	for _, artifact := range artifacts {
		name := artifact
		if relativePath, err := filepath.Rel(root, artifact); err == nil {
			name = relativePath
		}
		e := r.SendPayload(reporter.ReportPayload{Path: r.JobstepAPIPath() + "artifacts/",
			Data: map[string]string{"name": name}, Filename: artifact})
		if e != nil && firstError == nil {
			firstError = e
		}
	}
	return firstError
}

func New() reporter.Reporter {
	return &Reporter{}
}

func init() {
	reporter.Register("mesos", New)
}
<|endoftext|>


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

package rpc

import (
	"context"
	"fmt"
	"math"

	"github.com/golang/protobuf/ptypes"

	"go.chromium.org/luci/common/data/strpair"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/common/tsmon/distribution"
	"go.chromium.org/luci/common/tsmon/field"
	"go.chromium.org/luci/common/tsmon/metric"
	"go.chromium.org/luci/common/tsmon/types"
	"go.chromium.org/luci/server/auth"

	"go.chromium.org/luci/buildbucket/appengine/model"
	"go.chromium.org/luci/buildbucket/appengine/tasks"
	taskdefs "go.chromium.org/luci/buildbucket/appengine/tasks/defs"
	pb "go.chromium.org/luci/buildbucket/proto"
)

var (
	// A common set of field definitions for build metrics.
	fieldDefs = map[string]field.Field{
		"bucket":             field.String("bucket"),
		"builder":            field.String("builder"),
		"canary":             field.Bool("canary"),
		"cancelation_reason": field.String("cancelation_reason"),
		"failure_reason":     field.String("failure_reason"),
		"result":             field.String("result"),
		"user_agent":         field.String("user_agent"),
	}

	buildCountCreated = metric.NewCounter(
		"buildbucket/builds/created",
		"Build creation", nil,
		bFields("user_agent")...,
	)
	buildCountStarted = metric.NewCounter(
		"buildbucket/builds/started",
		"Build start", nil,
		bFields("canary")...,
	)
	buildCountCompleted = metric.NewCounter(
		"buildbucket/builds/completed",
		"Build completion, including success, failure and cancellation", nil,
		bFields("result", "failure_reason", "cancelation_reason", "canary")...,
	)

	buildDurationCycle = newbuildDurationMetric(
		"buildbucket/builds/cycle_durations",
		"Duration between build creation and completion",
	)
	buildDurationRun = newbuildDurationMetric(
		"buildbucket/builds/run_durations",
		"Duration between build start and completion",
	)
	buildDurationScheduling = newbuildDurationMetric(
		"buildbucket/builds/scheduling_durations",
		"Duration between build creation and start",
	)
)

func bFields(extraFields ...string) []field.Field {
	fs := make([]field.Field, 2+len(extraFields))
	fs[0], fs[1] = fieldDefs["bucket"], fieldDefs["builder"]
	for i, n := range extraFields {
		f, ok := fieldDefs[n]
		if !ok {
			panic(fmt.Sprintf("unknown build field %q", n))
		}
		fs[i+2] = f
	}
	return fs
}

func newbuildDurationMetric(name, description string, extraFields ...string) metric.CumulativeDistribution {
	fs := []string{"result", "failure_reason", "cancelation_reason", "canary"}
	return metric.NewCumulativeDistribution(
		name, description, &types.MetricMetadata{Units: types.Seconds},
		// Bucketer for 1s..48h range
		distribution.GeometricBucketer(math.Pow(10, 0.053), 100),
		bFields(append(fs, extraFields...)...)...,
	)
}

func getLegacyMetricFields(b *model.Build) (result, failureR, cancelationR string) {
	// The default values are "" instead of UNSET for backwards compatibility.
	switch b.Status {
	case pb.Status_SCHEDULED:
	case pb.Status_STARTED:
	case pb.Status_SUCCESS:
		result = model.Success.String()
	case pb.Status_FAILURE:
		result = model.Failure.String()
		failureR = model.BuildFailure.String()
	case pb.Status_INFRA_FAILURE:
		if b.Proto.StatusDetails.GetTimeout() != nil {
			result = model.Canceled.String()
			cancelationR = model.TimeoutCanceled.String()
		} else {
			result = model.Failure.String()
			failureR = model.InfraFailure.String()
		}
	case pb.Status_CANCELED:
		result = model.Canceled.String()
		cancelationR = model.ExplicitlyCanceled.String()
	default:
		panic(fmt.Sprintf("getLegacyMetricFields: invalid status %q", b.Status))
	}
	return
}

func buildCreated(ctx context.Context, b *model.Build) {
	var ua string
	for _, tag := range b.Tags {
		if k, v := strpair.Parse(tag); k == "user_agent" {
			ua = v
			break
		}
	}
	buildCountCreated.Add(ctx, 1, b.Proto.Builder.Bucket, b.Proto.Builder.Builder, ua)
}

func buildStarted(ctx context.Context, b *model.Build) {
	logging.Infof(ctx, "Build %d: started", b.ID)
	buildCountStarted.Add(ctx, 1, b.Proto.Builder.Bucket, b.Proto.Builder.Builder, b.Proto.Canary)
	if b.Proto.GetStartTime() != nil {
		startT, _ := ptypes.Timestamp(b.Proto.StartTime)
		buildDurationScheduling.Add(
			ctx, startT.Sub(b.CreateTime).Seconds(),
			b.Proto.Builder.Bucket, b.Proto.Builder.Builder, "", "", "", b.Proto.Canary,
		)
	}
}

func buildStarting(ctx context.Context, b *model.Build) error {
	return notifyPubSub(ctx, b)
}

func buildCompleted(ctx context.Context, b *model.Build) {
	r, fr, cr := getLegacyMetricFields(b)
	logging.Infof(ctx, "Build %d: completed by %q with status %q", b.ID, auth.CurrentIdentity(ctx), r)
	buildCountCompleted.Add(ctx, 1, b.Proto.Builder.Bucket, b.Proto.Builder.Builder, r, fr, cr, b.Proto.Canary)

	endT, _ := ptypes.Timestamp(b.Proto.EndTime)
	buildDurationCycle.Add(
		ctx, endT.Sub(b.CreateTime).Seconds(),
		b.Proto.Builder.Bucket, b.Proto.Builder.Builder, r, fr, cr, b.Proto.Canary,
	)
	if b.Proto.StartTime != nil {
		startT, _ := ptypes.Timestamp(b.Proto.StartTime)
		buildDurationRun.Add(
			ctx, endT.Sub(startT).Seconds(),
			b.Proto.Builder.Bucket, b.Proto.Builder.Builder, r, fr, cr, b.Proto.Canary,
		)
	}
}

func buildCompleting(ctx context.Context, b *model.Build) error {
	bqTask := &taskdefs.ExportBigQuery{BuildId: b.ID}
	invTask := &taskdefs.FinalizeResultDB{BuildId: b.ID}
	return parallel.FanOutIn(func(tks chan<- func() error) {
		tks <- func() error { return notifyPubSub(ctx, b) }
		tks <- func() error { return tasks.ExportBigQuery(ctx, bqTask) }
		tks <- func() error { return tasks.FinalizeResultDB(ctx, invTask) }
	})
}
<commit_msg>[buildbucket] speculative mitigation against panic in DS guts.<commit_after>// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package rpc

import (
	"context"
	"fmt"
	"math"

	"github.com/golang/protobuf/ptypes"

	"go.chromium.org/luci/common/data/strpair"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/common/tsmon/distribution"
	"go.chromium.org/luci/common/tsmon/field"
	"go.chromium.org/luci/common/tsmon/metric"
	"go.chromium.org/luci/common/tsmon/types"
	"go.chromium.org/luci/server/auth"

	"go.chromium.org/luci/buildbucket/appengine/model"
	"go.chromium.org/luci/buildbucket/appengine/tasks"
	taskdefs "go.chromium.org/luci/buildbucket/appengine/tasks/defs"
	pb "go.chromium.org/luci/buildbucket/proto"
)

var (
	// A common set of field definitions for build metrics.
	fieldDefs = map[string]field.Field{
		"bucket":             field.String("bucket"),
		"builder":            field.String("builder"),
		"canary":             field.Bool("canary"),
		"cancelation_reason": field.String("cancelation_reason"),
		"failure_reason":     field.String("failure_reason"),
		"result":             field.String("result"),
		"user_agent":         field.String("user_agent"),
	}

	buildCountCreated = metric.NewCounter(
		"buildbucket/builds/created",
		"Build creation", nil,
		bFields("user_agent")...,
	)
	buildCountStarted = metric.NewCounter(
		"buildbucket/builds/started",
		"Build start", nil,
		bFields("canary")...,
	)
	buildCountCompleted = metric.NewCounter(
		"buildbucket/builds/completed",
		"Build completion, including success, failure and cancellation", nil,
		bFields("result", "failure_reason", "cancelation_reason", "canary")...,
	)

	buildDurationCycle = newbuildDurationMetric(
		"buildbucket/builds/cycle_durations",
		"Duration between build creation and completion",
	)
	buildDurationRun = newbuildDurationMetric(
		"buildbucket/builds/run_durations",
		"Duration between build start and completion",
	)
	buildDurationScheduling = newbuildDurationMetric(
		"buildbucket/builds/scheduling_durations",
		"Duration between build creation and start",
	)
)

func bFields(extraFields ...string) []field.Field {
	fs := make([]field.Field, 2+len(extraFields))
	fs[0], fs[1] = fieldDefs["bucket"], fieldDefs["builder"]
	for i, n := range extraFields {
		f, ok := fieldDefs[n]
		if !ok {
			panic(fmt.Sprintf("unknown build field %q", n))
		}
		fs[i+2] = f
	}
	return fs
}

func newbuildDurationMetric(name, description string, extraFields ...string) metric.CumulativeDistribution {
	fs := []string{"result", "failure_reason", "cancelation_reason", "canary"}
	return metric.NewCumulativeDistribution(
		name, description, &types.MetricMetadata{Units: types.Seconds},
		// Bucketer for 1s..48h range
		distribution.GeometricBucketer(math.Pow(10, 0.053), 100),
		bFields(append(fs, extraFields...)...)...,
	)
}

func getLegacyMetricFields(b *model.Build) (result, failureR, cancelationR string) {
	// The default values are "" instead of UNSET for backwards compatibility.
	switch b.Status {
	case pb.Status_SCHEDULED:
	case pb.Status_STARTED:
	case pb.Status_SUCCESS:
		result = model.Success.String()
	case pb.Status_FAILURE:
		result = model.Failure.String()
		failureR = model.BuildFailure.String()
	case pb.Status_INFRA_FAILURE:
		if b.Proto.StatusDetails.GetTimeout() != nil {
			result = model.Canceled.String()
			cancelationR = model.TimeoutCanceled.String()
		} else {
			result = model.Failure.String()
			failureR = model.InfraFailure.String()
		}
	case pb.Status_CANCELED:
		result = model.Canceled.String()
		cancelationR = model.ExplicitlyCanceled.String()
	default:
		panic(fmt.Sprintf("getLegacyMetricFields: invalid status %q", b.Status))
	}
	return
}

func buildCreated(ctx context.Context, b *model.Build) {
	var ua string
	for _, tag := range b.Tags {
		if k, v := strpair.Parse(tag); k == "user_agent" {
			ua = v
			break
		}
	}
	buildCountCreated.Add(ctx, 1, b.Proto.Builder.Bucket, b.Proto.Builder.Builder, ua)
}

func buildStarted(ctx context.Context, b *model.Build) {
	logging.Infof(ctx, "Build %d: started", b.ID)
	buildCountStarted.Add(ctx, 1, b.Proto.Builder.Bucket, b.Proto.Builder.Builder, b.Proto.Canary)
	if b.Proto.GetStartTime() != nil {
		startT, _ := ptypes.Timestamp(b.Proto.StartTime)
		buildDurationScheduling.Add(
			ctx, startT.Sub(b.CreateTime).Seconds(),
			b.Proto.Builder.Bucket, b.Proto.Builder.Builder, "", "", "", b.Proto.Canary,
		)
	}
}

func buildStarting(ctx context.Context, b *model.Build) error {
	return notifyPubSub(ctx, b)
}

func buildCompleted(ctx context.Context, b *model.Build) {
	r, fr, cr := getLegacyMetricFields(b)
	logging.Infof(ctx, "Build %d: completed by %q with status %q", b.ID, auth.CurrentIdentity(ctx), r)
	buildCountCompleted.Add(ctx, 1, b.Proto.Builder.Bucket, b.Proto.Builder.Builder, r, fr, cr, b.Proto.Canary)

	endT, _ := ptypes.Timestamp(b.Proto.EndTime)
	buildDurationCycle.Add(
		ctx, endT.Sub(b.CreateTime).Seconds(),
		b.Proto.Builder.Bucket, b.Proto.Builder.Builder, r, fr, cr, b.Proto.Canary,
	)
	if b.Proto.StartTime != nil {
		startT, _ := ptypes.Timestamp(b.Proto.StartTime)
		buildDurationRun.Add(
			ctx, endT.Sub(startT).Seconds(),
			b.Proto.Builder.Bucket, b.Proto.Builder.Builder, r, fr, cr, b.Proto.Canary,
		)
	}
}

func buildCompleting(ctx context.Context, b *model.Build) error {
	bqTask := &taskdefs.ExportBigQuery{BuildId: b.ID}
	invTask := &taskdefs.FinalizeResultDB{BuildId: b.ID}
	// TODO(tandrii): undo this and find a permanent solution.
	return parallel.WorkPool(1, func(tks chan<- func() error) {
		tks <- func() error { return notifyPubSub(ctx, b) }
		tks <- func() error { return tasks.ExportBigQuery(ctx, bqTask) }
		tks <- func() error { return tasks.FinalizeResultDB(ctx, invTask) }
	})
}
<|endoftext|>