<commit_before>// 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 spinner is a simple package to add a spinner / progress indicator to any terminal application.
package spinner

import (
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"os"
	"runtime"
	"strconv"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/fatih/color"
)

// errInvalidColor is returned when attempting to set an invalid color
var errInvalidColor = errors.New("invalid color")

// validColors holds an array of the only colors allowed
var validColors = map[string]bool{
	// default colors for backwards compatibility
	"black":   true,
	"red":     true,
	"green":   true,
	"yellow":  true,
	"blue":    true,
	"magenta": true,
	"cyan":    true,
	"white":   true,

	// attributes
	"reset":        true,
	"bold":         true,
	"faint":        true,
	"italic":       true,
	"underline":    true,
	"blinkslow":    true,
	"blinkrapid":   true,
	"reversevideo": true,
	"concealed":    true,
	"crossedout":   true,

	// foreground text
	"fgBlack":   true,
	"fgRed":     true,
	"fgGreen":   true,
	"fgYellow":  true,
	"fgBlue":    true,
	"fgMagenta": true,
	"fgCyan":    true,
	"fgWhite":   true,

	// foreground Hi-Intensity text
	"fgHiBlack":   true,
	"fgHiRed":     true,
	"fgHiGreen":   true,
	"fgHiYellow":  true,
	"fgHiBlue":    true,
	"fgHiMagenta": true,
	"fgHiCyan":    true,
	"fgHiWhite":   true,

	// background text
	"bgBlack":   true,
	"bgRed":     true,
	"bgGreen":   true,
	"bgYellow":  true,
	"bgBlue":    true,
	"bgMagenta": true,
	"bgCyan":    true,
	"bgWhite":   true,

	// background Hi-Intensity text
	"bgHiBlack":   true,
	"bgHiRed":     true,
	"bgHiGreen":   true,
	"bgHiYellow":  true,
	"bgHiBlue":    true,
	"bgHiMagenta": true,
	"bgHiCyan":    true,
	"bgHiWhite":   true,
}

// returns a valid color's foreground text color attribute
var colorAttributeMap = map[string]color.Attribute{
	// default colors for backwards compatibility
	"black":   color.FgBlack,
	"red":     color.FgRed,
	"green":   color.FgGreen,
	"yellow":  color.FgYellow,
	"blue":    color.FgBlue,
	"magenta": color.FgMagenta,
	"cyan":    color.FgCyan,
	"white":   color.FgWhite,

	// attributes
	"reset":        color.Reset,
	"bold":         color.Bold,
	"faint":        color.Faint,
	"italic":       color.Italic,
	"underline":    color.Underline,
	"blinkslow":    color.BlinkSlow,
	"blinkrapid":   color.BlinkRapid,
	"reversevideo": color.ReverseVideo,
	"concealed":    color.Concealed,
	"crossedout":   color.CrossedOut,

	// foreground text colors
	"fgBlack":   color.FgBlack,
	"fgRed":     color.FgRed,
	"fgGreen":   color.FgGreen,
	"fgYellow":  color.FgYellow,
	"fgBlue":    color.FgBlue,
	"fgMagenta": color.FgMagenta,
	"fgCyan":    color.FgCyan,
	"fgWhite":   color.FgWhite,

	// foreground Hi-Intensity text colors
	"fgHiBlack":   color.FgHiBlack,
	"fgHiRed":     color.FgHiRed,
	"fgHiGreen":   color.FgHiGreen,
	"fgHiYellow":  color.FgHiYellow,
	"fgHiBlue":    color.FgHiBlue,
	"fgHiMagenta": color.FgHiMagenta,
	"fgHiCyan":    color.FgHiCyan,
	"fgHiWhite":   color.FgHiWhite,

	// background text colors
	"bgBlack":   color.BgBlack,
	"bgRed":     color.BgRed,
	"bgGreen":   color.BgGreen,
	"bgYellow":  color.BgYellow,
	"bgBlue":    color.BgBlue,
	"bgMagenta": color.BgMagenta,
	"bgCyan":    color.BgCyan,
	"bgWhite":   color.BgWhite,

	// background Hi-Intensity text colors
	"bgHiBlack":   color.BgHiBlack,
	"bgHiRed":     color.BgHiRed,
	"bgHiGreen":   color.BgHiGreen,
	"bgHiYellow":  color.BgHiYellow,
	"bgHiBlue":    color.BgHiBlue,
	"bgHiMagenta": color.BgHiMagenta,
	"bgHiCyan":    color.BgHiCyan,
	"bgHiWhite":   color.BgHiWhite,
}

// validColor will make sure the given color is actually allowed.
func validColor(c string) bool {
	if validColors[c] {
		return true
	}
	return false
}

// Spinner struct to hold the provided options.
type Spinner struct {
	mu         *sync.RWMutex                 //
	Delay      time.Duration                 // Delay is the speed of the indicator
	chars      []string                      // chars holds the chosen character set
	Prefix     string                        // Prefix is the text preppended to the indicator
	Suffix     string                        // Suffix is the text appended to the indicator
	FinalMSG   string                        // string displayed after Stop() is called
	lastOutput string                        // last character(set) written
	color      func(a ...interface{}) string // default color is white
	Writer     io.Writer                     // to make testing better, exported so users have access. Use `WithWriter` to update after initialization.
	active     bool                          // active holds the state of the spinner
	stopChan   chan struct{}                 // stopChan is a channel used to stop the indicator
	HideCursor bool                          // hideCursor determines if the cursor is visible
	PreUpdate  func(s *Spinner)              // will be triggered before every spinner update
	PostUpdate func(s *Spinner)              // will be triggered after every spinner update
}

// New provides a pointer to an instance of Spinner with the supplied options.
func New(cs []string, d time.Duration, options ...Option) *Spinner {
	s := &Spinner{
		Delay:    d,
		chars:    cs,
		color:    color.New(color.FgWhite).SprintFunc(),
		mu:       &sync.RWMutex{},
		Writer:   color.Output,
		active:   false,
		stopChan: make(chan struct{}, 1),
	}

	for _, option := range options {
		option(s)
	}
	return s
}

// Option is a function that takes a spinner and applies
// a given configuration.
type Option func(*Spinner)

// Options contains fields to configure the spinner.
type Options struct {
	Color      string
	Suffix     string
	FinalMSG   string
	HideCursor bool
}

// WithColor adds the given color to the spinner.
func WithColor(color string) Option {
	return func(s *Spinner) {
		s.Color(color)
	}
}

// WithSuffix adds the given string to the spinner
// as the suffix.
func WithSuffix(suffix string) Option {
	return func(s *Spinner) {
		s.Suffix = suffix
	}
}

// WithFinalMSG adds the given string ot the spinner
// as the final message to be written.
func WithFinalMSG(finalMsg string) Option {
	return func(s *Spinner) {
		s.FinalMSG = finalMsg
	}
}

// WithHiddenCursor hides the cursor
// if hideCursor = true given.
func WithHiddenCursor(hideCursor bool) Option {
	return func(s *Spinner) {
		s.HideCursor = hideCursor
	}
}

// WithWriter adds the given writer to the spinner. This
// function should be favored over directly assigning to
// the struct value.
func WithWriter(w io.Writer) Option {
	return func(s *Spinner) {
		s.mu.Lock()
		s.Writer = w
		s.mu.Unlock()
	}
}

// Active will return whether or not the spinner is currently active.
func (s *Spinner) Active() bool {
	return s.active
}

// Start will start the indicator.
func (s *Spinner) Start() {
	s.mu.Lock()
	if s.active {
		s.mu.Unlock()
		return
	}
	if s.HideCursor && runtime.GOOS != "windows" {
		// hides the cursor
		fmt.Print("\033[?25l")
	}
	s.active = true
	s.mu.Unlock()

	go func() {
		for {
			for i := 0; i < len(s.chars); i++ {
				select {
				case <-s.stopChan:
					return
				default:
					if !s.active {
						return
					}
					s.mu.Lock()
					s.erase()

					if s.PreUpdate != nil {
						s.PreUpdate(s)
					}

					var outColor string
					if runtime.GOOS == "windows" {
						if s.Writer == os.Stderr {
							outColor = fmt.Sprintf("\r%s%s%s ", s.Prefix, s.chars[i], s.Suffix)
						} else {
							outColor = fmt.Sprintf("\r%s%s%s ", s.Prefix, s.color(s.chars[i]), s.Suffix)
						}
					} else {
						outColor = fmt.Sprintf("%s%s%s ", s.Prefix, s.color(s.chars[i]), s.Suffix)
					}
					outPlain := fmt.Sprintf("%s%s%s ", s.Prefix, s.chars[i], s.Suffix)
					fmt.Fprint(s.Writer, outColor)
					s.lastOutput = outPlain
					delay := s.Delay

					if s.PostUpdate != nil {
						s.PostUpdate(s)
					}

					s.mu.Unlock()
					time.Sleep(delay)
				}
			}
		}
	}()
}

// Stop stops the indicator.
func (s *Spinner) Stop() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.active {
		s.active = false
		if s.HideCursor && runtime.GOOS != "windows" {
			// makes the cursor visible
			fmt.Print("\033[?25h")
		}
		s.erase()
		if s.FinalMSG != "" {
			fmt.Fprintf(s.Writer, s.FinalMSG)
		}
		s.stopChan <- struct{}{}
	}
}

// Restart will stop and start the indicator.
func (s *Spinner) Restart() {
	s.Stop()
	s.Start()
}

// Reverse will reverse the order of the slice assigned to the indicator.
func (s *Spinner) Reverse() {
	s.mu.Lock()
	defer s.mu.Unlock()
	for i, j := 0, len(s.chars)-1; i < j; i, j = i+1, j-1 {
		s.chars[i], s.chars[j] = s.chars[j], s.chars[i]
	}
}

// Color will set the struct field for the given color to be used.
func (s *Spinner) Color(colors ...string) error {
	colorAttributes := make([]color.Attribute, len(colors))

	// Verify colours are valid and place the appropriate attribute in the array
	for index, c := range colors {
		if !validColor(c) {
			return errInvalidColor
		}
		colorAttributes[index] = colorAttributeMap[c]
	}

	s.mu.Lock()
	s.color = color.New(colorAttributes...).SprintFunc()
	s.mu.Unlock()
	s.Restart()
	return nil
}

// UpdateSpeed will set the indicator delay to the given value.
func (s *Spinner) UpdateSpeed(d time.Duration) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.Delay = d
}

// UpdateCharSet will change the current character set to the given one.
func (s *Spinner) UpdateCharSet(cs []string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.chars = cs
}

// erase deletes written characters.
// Caller must already hold s.lock.
func (s *Spinner) erase() {
	n := utf8.RuneCountInString(s.lastOutput)
	if runtime.GOOS == "windows" {
		clearString := "\r"
		for i := 0; i < n; i++ {
			clearString += " "
		}
		clearString += "\r"
		fmt.Fprintf(s.Writer, clearString)
		s.lastOutput = ""
		return
	}
	del, _ := hex.DecodeString("7f")
	for _, c := range []string{"\b", string(del)} {
		for i := 0; i < n; i++ {
			fmt.Fprintf(s.Writer, c)
		}
	}
	fmt.Fprintf(s.Writer, "\033[K") // erases to end of line
	s.lastOutput = ""
}

// Lock allows for manual control to lock the spinner.
func (s *Spinner) Lock() {
	s.mu.Lock()
}

// Unlock allows for manual control to unlock the spinner.
func (s *Spinner) Unlock() {
	s.mu.Unlock()
}

// GenerateNumberSequence will generate a slice of integers at the
// provided length and convert them each to a string.
func GenerateNumberSequence(length int) []string {
	numSeq := make([]string, length)
	for i := 0; i < length; i++ {
		numSeq[i] = strconv.Itoa(i)
	}
	return numSeq
}
<commit_msg>Fix output on iTerm in macOS<commit_after>// 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 spinner is a simple package to add a spinner / progress indicator to any terminal application.
package spinner

import (
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"os"
	"runtime"
	"strconv"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/fatih/color"
)

// errInvalidColor is returned when attempting to set an invalid color
var errInvalidColor = errors.New("invalid color")

// validColors holds an array of the only colors allowed
var validColors = map[string]bool{
	// default colors for backwards compatibility
	"black":   true,
	"red":     true,
	"green":   true,
	"yellow":  true,
	"blue":    true,
	"magenta": true,
	"cyan":    true,
	"white":   true,

	// attributes
	"reset":        true,
	"bold":         true,
	"faint":        true,
	"italic":       true,
	"underline":    true,
	"blinkslow":    true,
	"blinkrapid":   true,
	"reversevideo": true,
	"concealed":    true,
	"crossedout":   true,

	// foreground text
	"fgBlack":   true,
	"fgRed":     true,
	"fgGreen":   true,
	"fgYellow":  true,
	"fgBlue":    true,
	"fgMagenta": true,
	"fgCyan":    true,
	"fgWhite":   true,

	// foreground Hi-Intensity text
	"fgHiBlack":   true,
	"fgHiRed":     true,
	"fgHiGreen":   true,
	"fgHiYellow":  true,
	"fgHiBlue":    true,
	"fgHiMagenta": true,
	"fgHiCyan":    true,
	"fgHiWhite":   true,

	// background text
	"bgBlack":   true,
	"bgRed":     true,
	"bgGreen":   true,
	"bgYellow":  true,
	"bgBlue":    true,
	"bgMagenta": true,
	"bgCyan":    true,
	"bgWhite":   true,

	// background Hi-Intensity text
	"bgHiBlack":   true,
	"bgHiRed":     true,
	"bgHiGreen":   true,
	"bgHiYellow":  true,
	"bgHiBlue":    true,
	"bgHiMagenta": true,
	"bgHiCyan":    true,
	"bgHiWhite":   true,
}

// returns a valid color's foreground text color attribute
var colorAttributeMap = map[string]color.Attribute{
	// default colors for backwards compatibility
	"black":   color.FgBlack,
	"red":     color.FgRed,
	"green":   color.FgGreen,
	"yellow":  color.FgYellow,
	"blue":    color.FgBlue,
	"magenta": color.FgMagenta,
	"cyan":    color.FgCyan,
	"white":   color.FgWhite,

	// attributes
	"reset":        color.Reset,
	"bold":         color.Bold,
	"faint":        color.Faint,
	"italic":       color.Italic,
	"underline":    color.Underline,
	"blinkslow":    color.BlinkSlow,
	"blinkrapid":   color.BlinkRapid,
	"reversevideo": color.ReverseVideo,
	"concealed":    color.Concealed,
	"crossedout":   color.CrossedOut,

	// foreground text colors
	"fgBlack":   color.FgBlack,
	"fgRed":     color.FgRed,
	"fgGreen":   color.FgGreen,
	"fgYellow":  color.FgYellow,
	"fgBlue":    color.FgBlue,
	"fgMagenta": color.FgMagenta,
	"fgCyan":    color.FgCyan,
	"fgWhite":   color.FgWhite,

	// foreground Hi-Intensity text colors
	"fgHiBlack":   color.FgHiBlack,
	"fgHiRed":     color.FgHiRed,
	"fgHiGreen":   color.FgHiGreen,
	"fgHiYellow":  color.FgHiYellow,
	"fgHiBlue":    color.FgHiBlue,
	"fgHiMagenta": color.FgHiMagenta,
	"fgHiCyan":    color.FgHiCyan,
	"fgHiWhite":   color.FgHiWhite,

	// background text colors
	"bgBlack":   color.BgBlack,
	"bgRed":     color.BgRed,
	"bgGreen":   color.BgGreen,
	"bgYellow":  color.BgYellow,
	"bgBlue":    color.BgBlue,
	"bgMagenta": color.BgMagenta,
	"bgCyan":    color.BgCyan,
	"bgWhite":   color.BgWhite,

	// background Hi-Intensity text colors
	"bgHiBlack":   color.BgHiBlack,
	"bgHiRed":     color.BgHiRed,
	"bgHiGreen":   color.BgHiGreen,
	"bgHiYellow":  color.BgHiYellow,
	"bgHiBlue":    color.BgHiBlue,
	"bgHiMagenta": color.BgHiMagenta,
	"bgHiCyan":    color.BgHiCyan,
	"bgHiWhite":   color.BgHiWhite,
}

// validColor will make sure the given color is actually allowed.
func validColor(c string) bool {
	if validColors[c] {
		return true
	}
	return false
}

// Spinner struct to hold the provided options.
type Spinner struct {
	mu         *sync.RWMutex                 //
	Delay      time.Duration                 // Delay is the speed of the indicator
	chars      []string                      // chars holds the chosen character set
	Prefix     string                        // Prefix is the text preppended to the indicator
	Suffix     string                        // Suffix is the text appended to the indicator
	FinalMSG   string                        // string displayed after Stop() is called
	lastOutput string                        // last character(set) written
	color      func(a ...interface{}) string // default color is white
	Writer     io.Writer                     // to make testing better, exported so users have access. Use `WithWriter` to update after initialization.
	active     bool                          // active holds the state of the spinner
	stopChan   chan struct{}                 // stopChan is a channel used to stop the indicator
	HideCursor bool                          // hideCursor determines if the cursor is visible
	PreUpdate  func(s *Spinner)              // will be triggered before every spinner update
	PostUpdate func(s *Spinner)              // will be triggered after every spinner update
}

// New provides a pointer to an instance of Spinner with the supplied options.
func New(cs []string, d time.Duration, options ...Option) *Spinner {
	s := &Spinner{
		Delay:    d,
		chars:    cs,
		color:    color.New(color.FgWhite).SprintFunc(),
		mu:       &sync.RWMutex{},
		Writer:   color.Output,
		active:   false,
		stopChan: make(chan struct{}, 1),
	}

	for _, option := range options {
		option(s)
	}
	return s
}

// Option is a function that takes a spinner and applies
// a given configuration.
type Option func(*Spinner)

// Options contains fields to configure the spinner.
type Options struct {
	Color      string
	Suffix     string
	FinalMSG   string
	HideCursor bool
}

// WithColor adds the given color to the spinner.
func WithColor(color string) Option {
	return func(s *Spinner) {
		s.Color(color)
	}
}

// WithSuffix adds the given string to the spinner
// as the suffix.
func WithSuffix(suffix string) Option {
	return func(s *Spinner) {
		s.Suffix = suffix
	}
}

// WithFinalMSG adds the given string ot the spinner
// as the final message to be written.
func WithFinalMSG(finalMsg string) Option {
	return func(s *Spinner) {
		s.FinalMSG = finalMsg
	}
}

// WithHiddenCursor hides the cursor
// if hideCursor = true given.
func WithHiddenCursor(hideCursor bool) Option {
	return func(s *Spinner) {
		s.HideCursor = hideCursor
	}
}

// WithWriter adds the given writer to the spinner. This
// function should be favored over directly assigning to
// the struct value.
func WithWriter(w io.Writer) Option {
	return func(s *Spinner) {
		s.mu.Lock()
		s.Writer = w
		s.mu.Unlock()
	}
}

// Active will return whether or not the spinner is currently active.
func (s *Spinner) Active() bool {
	return s.active
}

// Start will start the indicator.
func (s *Spinner) Start() {
	s.mu.Lock()
	if s.active {
		s.mu.Unlock()
		return
	}
	if s.HideCursor && runtime.GOOS != "windows" {
		// hides the cursor
		fmt.Print("\033[?25l")
	}
	s.active = true
	s.mu.Unlock()

	go func() {
		for {
			for i := 0; i < len(s.chars); i++ {
				select {
				case <-s.stopChan:
					return
				default:
					if !s.active {
						return
					}
					s.mu.Lock()
					s.erase()

					if s.PreUpdate != nil {
						s.PreUpdate(s)
					}

					var outColor string
					if runtime.GOOS == "windows" {
						if s.Writer == os.Stderr {
							outColor = fmt.Sprintf("\r%s%s%s ", s.Prefix, s.chars[i], s.Suffix)
						} else {
							outColor = fmt.Sprintf("\r%s%s%s ", s.Prefix, s.color(s.chars[i]), s.Suffix)
						}
					} else {
						outColor = fmt.Sprintf("%s%s%s ", s.Prefix, s.color(s.chars[i]), s.Suffix)
					}
					outPlain := fmt.Sprintf("%s%s%s ", s.Prefix, s.chars[i], s.Suffix)
					fmt.Fprint(s.Writer, outColor)
					s.lastOutput = outPlain
					delay := s.Delay

					if s.PostUpdate != nil {
						s.PostUpdate(s)
					}

					s.mu.Unlock()
					time.Sleep(delay)
				}
			}
		}
	}()
}

// Stop stops the indicator.
func (s *Spinner) Stop() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.active {
		s.active = false
		if s.HideCursor && runtime.GOOS != "windows" {
			// makes the cursor visible
			fmt.Print("\033[?25h")
		}
		s.erase()
		if s.FinalMSG != "" {
			fmt.Fprintf(s.Writer, s.FinalMSG)
		}
		s.stopChan <- struct{}{}
	}
}

// Restart will stop and start the indicator.
func (s *Spinner) Restart() {
	s.Stop()
	s.Start()
}

// Reverse will reverse the order of the slice assigned to the indicator.
func (s *Spinner) Reverse() {
	s.mu.Lock()
	defer s.mu.Unlock()
	for i, j := 0, len(s.chars)-1; i < j; i, j = i+1, j-1 {
		s.chars[i], s.chars[j] = s.chars[j], s.chars[i]
	}
}

// Color will set the struct field for the given color to be used.
func (s *Spinner) Color(colors ...string) error {
	colorAttributes := make([]color.Attribute, len(colors))

	// Verify colours are valid and place the appropriate attribute in the array
	for index, c := range colors {
		if !validColor(c) {
			return errInvalidColor
		}
		colorAttributes[index] = colorAttributeMap[c]
	}

	s.mu.Lock()
	s.color = color.New(colorAttributes...).SprintFunc()
	s.mu.Unlock()
	s.Restart()
	return nil
}

// UpdateSpeed will set the indicator delay to the given value.
func (s *Spinner) UpdateSpeed(d time.Duration) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.Delay = d
}

// UpdateCharSet will change the current character set to the given one.
func (s *Spinner) UpdateCharSet(cs []string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.chars = cs
}

// erase deletes written characters.
// Caller must already hold s.lock.
func (s *Spinner) erase() {
	n := utf8.RuneCountInString(s.lastOutput)
	if runtime.GOOS == "windows" {
		clearString := "\r"
		for i := 0; i < n; i++ {
			clearString += " "
		}
		clearString += "\r"
		fmt.Fprintf(s.Writer, clearString)
		s.lastOutput = ""
		return
	}
	del, _ := hex.DecodeString("7f")
	for _, c := range []string{"\b", string(del)} {
		for i := 0; i < n; i++ {
			fmt.Fprintf(s.Writer, c)
		}
	}
	fmt.Fprintf(s.Writer, "\r\033[K") // erases to end of line
	s.lastOutput = ""
}

// Lock allows for manual control to lock the spinner.
func (s *Spinner) Lock() {
	s.mu.Lock()
}

// Unlock allows for manual control to unlock the spinner.
func (s *Spinner) Unlock() {
	s.mu.Unlock()
}

// GenerateNumberSequence will generate a slice of integers at the
// provided length and convert them each to a string.
func GenerateNumberSequence(length int) []string {
	numSeq := make([]string, length)
	for i := 0; i < length; i++ {
		numSeq[i] = strconv.Itoa(i)
	}
	return numSeq
}
<|endoftext|>


<commit_before>package templates

import (
	"html/template"
	"os"
	"path/filepath"
	"strings"

	"github.com/anonx/sunplate/log"
)

var templates = map[string]*template.Template{}

// load gets three input arguments:
// 1. Path to the root of the user application (e.g. "./").
// 2. Path to the views directory relative to the project root (e.g. "views").
// 3. A list of template paths relative to the views directory (in a form of map[string]string).
// It checks whether all the templates exist, parses and registers them.
// It panics if some of the requested templates do not exist or cannot be parsed.
func load(root string, views string, templatePaths map[string]string) {
	log.Trace.Println("Loading templates...")
	root = filepath.Join(root, views)

	// Iterating over all available template paths.
	for _, path := range templatePaths {
		// Find base for the current template
		// (either in the current dir or in one of the previous levels).
		var base, cd string
		for {
			b := filepath.Base(path)
			dir := filepath.Join(path[:len(path)-len(b)], cd)
			cd += "../"

			// Check whether this template is a base. If so, do not load
			// any other bases.
			if b == baseTemplate {
				break
			}

			// Check whether base template exists in the directory.
			base = filepath.Join(dir, baseTemplate)
			if _, ok := templates[base]; ok || contains(templatePaths, base) {
				break
			}
			base = ""

			// Check whether we have unsuccessfully achieved the top level
			// of the path.
			if strings.HasPrefix(dir, "../") {
				break
			}
		}

		log.Trace.Printf("\t%s (%s)", path, base)

		// If the base was found, use it. Otherwise, go without it.
		var err error
		t := template.New(path).Funcs(Funcs).Delims(delimLeft, delimRight)
		if base != "" {
			templates[path], err = t.ParseFiles(
				filepath.Join(root, base),
				filepath.Join(root, path),
			)
			showError(root, base, path, err)
			continue
		}
		templates[path], err = t.ParseFiles(filepath.Join(root, path))
		showError(root, base, path, err)
	}
}

// contains returns true if a requested value found
// in the requested slice of strings.
func contains(lst map[string]string, value string) bool {
	for k := range lst {
		if lst[k] == value {
			return true
		}
	}
	return false
}

// showErrors writes an error to log.
func showError(root, base, path string, err error) {
	if err == nil {
		return
	}
	pwd, _ := os.Getwd()
	log.Error.Panicf(
		`Cannot parse "%s" with "%s" as a base template (pwd "%s"). Error: %v.`,
		filepath.Join(root, path), filepath.Join(root, base), pwd, err,
	)
}
<commit_msg>Make sure infinite loop never happens<commit_after>package templates

import (
	"html/template"
	"os"
	"path/filepath"
	"strings"

	"github.com/anonx/sunplate/log"
)

var templates = map[string]*template.Template{}

// load gets three input arguments:
// 1. Path to the root of the user application (e.g. "./").
// 2. Path to the views directory relative to the project root (e.g. "views").
// 3. A list of template paths relative to the views directory (in a form of map[string]string).
// It checks whether all the templates exist, parses and registers them.
// It panics if some of the requested templates do not exist or cannot be parsed.
func load(root string, views string, templatePaths map[string]string) {
	log.Trace.Println("Loading templates...")
	root = filepath.Join(root, views)

	// Iterating over all available template paths.
	for _, path := range templatePaths {
		// Find base for the current template
		// (either in the current dir or in one of the previous levels).
		var base, cd string
		limit, i := 100, 0
		for {
			if i++; i == limit {
				break
			}

			b := filepath.Base(path)
			dir := filepath.Join(path[:len(path)-len(b)], cd)
			cd += "../"

			// Check whether this template is a base. If so, do not load
			// any other bases.
			if b == baseTemplate {
				break
			}

			// Check whether base template exists in the directory.
			base = filepath.Join(dir, baseTemplate)
			if _, ok := templates[base]; ok || contains(templatePaths, base) {
				break
			}
			base = ""

			// Check whether we have unsuccessfully achieved the top level
			// of the path.
			if strings.HasPrefix(dir, "../") {
				break
			}
		}

		log.Trace.Printf("\t%s (%s)", path, base)

		// If the base was found, use it. Otherwise, go without it.
		var err error
		t := template.New(path).Funcs(Funcs).Delims(delimLeft, delimRight)
		if base != "" {
			templates[path], err = t.ParseFiles(
				filepath.Join(root, base),
				filepath.Join(root, path),
			)
			showError(root, base, path, err)
			continue
		}
		templates[path], err = t.ParseFiles(filepath.Join(root, path))
		showError(root, base, path, err)
	}
}

// contains returns true if a requested value found
// in the requested slice of strings.
func contains(lst map[string]string, value string) bool {
	for k := range lst {
		if lst[k] == value {
			return true
		}
	}
	return false
}

// showErrors writes an error to log.
func showError(root, base, path string, err error) {
	if err == nil {
		return
	}
	pwd, _ := os.Getwd()
	log.Error.Panicf(
		`Cannot parse "%s" with "%s" as a base template (pwd "%s"). Error: %v.`,
		filepath.Join(root, path), filepath.Join(root, base), pwd, err,
	)
}
<|endoftext|>


<commit_before>package clusters

import (
	"fmt"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"strconv"
	"strings"
	"k8s.io/client-go/rest"
)

type ClusterConfig struct {
	MasterCount       int
	WorkerCount       int
	Name              string
	SparkMasterConfig string
	SparkWorkerConfig string
	SparkImage        string
	ExposeWebUI       string
	Metrics           string
}

const Defaultname = "default-oshinko-cluster-config"

var defaultConfig ClusterConfig = ClusterConfig{
	MasterCount:       1,
	WorkerCount:       1,
	Name:              "",
	SparkMasterConfig: "",
	SparkWorkerConfig: "",
	SparkImage:        "",
	ExposeWebUI:       "true",
	Metrics: 	   "false",
}

const failOnMissing = true
const allowMissing = false

const MasterCountMustBeZeroOrOne = "cluster configuration must have a master count of 0 or 1"
const WorkerCountMustBeAtLeastZero = "cluster configuration may not have a worker count less than 0"
const ErrorWhileProcessing = "'%s', %s"
const NamedConfigDoesNotExist = "named config '%s' does not exist"

const SentinelCountValue = -1

// This function is meant to support testability
func GetDefaultConfig() ClusterConfig {
	return defaultConfig
}

func assignConfig(res *ClusterConfig, src ClusterConfig) {
	if src.Name != "" {
		res.Name = src.Name
	}

	if src.MasterCount > SentinelCountValue {
		res.MasterCount = src.MasterCount
	}

	if src.WorkerCount > SentinelCountValue {
		res.WorkerCount = src.WorkerCount
	}

	if src.SparkMasterConfig != "" {
		res.SparkMasterConfig = src.SparkMasterConfig
	}
	if src.SparkWorkerConfig != "" {
		res.SparkWorkerConfig = src.SparkWorkerConfig
	}
	if src.SparkImage != "" {
		res.SparkImage = src.SparkImage
	}
	if src.ExposeWebUI != "" {
		res.ExposeWebUI = src.ExposeWebUI
	}
	if src.Metrics != "" {
		res.Metrics = src.Metrics
	}
}

func checkConfiguration(config ClusterConfig) error {
	var err error
	if config.MasterCount < 0 || config.MasterCount > 1 {
		err = NewClusterError(MasterCountMustBeZeroOrOne, ClusterConfigCode)
	} else if config.WorkerCount < 0 {
		err = NewClusterError(WorkerCountMustBeAtLeastZero, ClusterConfigCode)
	}
	return err
}

func getInt(value, configmapname string) (int, error) {
	i, err := strconv.Atoi(value)
	if err != nil {
		err = NewClusterError(fmt.Sprintf(ErrorWhileProcessing, configmapname, fmt.Sprintf("expected integer, got '%s'", value)), ClusterConfigCode)
	}
	return i, err
}

func process(config *ClusterConfig, name, value, configmapname string) error {

	var err error

	// At present we only have a single level of configs, but if/when we have
	// nested configs then we would descend through the levels beginning here with
	// the first element in the name
	switch name {
	case "mastercount":
		val, err := getInt(value, configmapname + ".mastercount")
		if err != nil {
			return err
		}
		if val > SentinelCountValue {
			config.MasterCount = val
		}
	case "workercount":
		val, err := getInt(value, configmapname + ".workercount")
		if err != nil {
			return err
		}
		if val > SentinelCountValue {
			config.WorkerCount = val
		}
	case "sparkmasterconfig":
		config.SparkMasterConfig = value
	case "sparkworkerconfig":
		config.SparkWorkerConfig = value
	case "sparkimage":
		config.SparkImage = value
	case "exposeui":
		config.ExposeWebUI = value
		_, err = strconv.ParseBool(config.ExposeWebUI)
	case "metrics":
		// default will be "false" if the string is empty
		if value != "" {
			val, err := strconv.ParseBool(value)
			if err == nil {
				//  Support 'true' and default to jolokia
				// Normalize truth values
				if val {
					config.Metrics = "true"
				} else {
					config.Metrics = "false"
				}
			} else {
				if value != "jolokia" && value != "prometheus" {
					msg := fmt.Sprintf("expected 'jolokia' or 'prometheus', got '%s'", value)
					return NewClusterError(fmt.Sprintf(ErrorWhileProcessing, configmapname, msg), ClusterConfigCode)
				}
				config.Metrics = value
			}
		}
	}
	return err
}

func readConfig(name string, res *ClusterConfig, failOnMissing bool, restconfig *rest.Config, namespace string) (found bool, err error) {

	found = false
	cmap, err := getKubeClient(restconfig).CoreV1().ConfigMaps(namespace).Get(name, metav1.GetOptions{})
	if err != nil {
		if strings.Index(err.Error(), "not found") != -1 {
			if !failOnMissing {
				err = nil
			} else {
				err = NewClusterError(fmt.Sprintf(NamedConfigDoesNotExist, name), ClusterConfigCode)
			}
		} else {
			err = NewClusterError(err.Error(), ClientOperationCode)
		}
	}
	if err == nil && cmap != nil {
		// Kube will give us an empty configmap if the named one does not exist,
		// so we test for a Name to see if we foud it
		found = cmap.Name != ""
		for n, v := range cmap.Data {
			err = process(res, strings.Trim(n, "\n"), strings.Trim(v, "\n"), name)
			if err != nil {
				break
			}
		}
	}
	return found, err
}

func loadConfig(name string, restconfig *rest.Config, namespace string) (res ClusterConfig, err error) {
	// If the default config has been modified use those mods.
	res = defaultConfig
	found, err := readConfig(Defaultname, &res, allowMissing, restconfig, namespace)
	if err == nil {
		if name != "" && name != Defaultname {
			_, err = readConfig(name, &res, failOnMissing, restconfig, namespace)
		} else if found {
			res.Name = Defaultname
		}
	}
	return res, err
}

func GetClusterConfig(config *ClusterConfig, restconfig *rest.Config, namespace string) (res ClusterConfig, err error) {

	var name string = ""
	if config != nil {
		name = config.Name
	}
	res, err = loadConfig(name, restconfig, namespace)
	if err == nil && config != nil {
		assignConfig(&res, *config)
	}

	// Check that the final configuration is valid
	if err == nil {
		err = checkConfiguration(res)
	}
	return res, err
}
<commit_msg>issue42<commit_after>package clusters

import (
	"fmt"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"strconv"
	"strings"
	"k8s.io/client-go/rest"
)

type ClusterConfig struct {
	MasterCount       int
	WorkerCount       int
	Name              string
	SparkMasterConfig string
	SparkWorkerConfig string
	SparkImage        string
	ExposeWebUI       string
	Metrics           string
}

const Defaultname = "default-oshinko-cluster-config"

var defaultConfig ClusterConfig = ClusterConfig{
	MasterCount:       1,
	WorkerCount:       1,
	Name:              "",
	SparkMasterConfig: "",
	SparkWorkerConfig: "",
	SparkImage:        "",
	ExposeWebUI:       "true",
	Metrics: 	   "false",
}

const failOnMissing = true
const allowMissing = false

const MasterCountMustBeZeroOrOne = "cluster configuration must have a master count of 0 or 1"
const WorkerCountMustBeAtLeastZero = "cluster configuration may not have a worker count less than 0"
const ErrorWhileProcessing = "'%s', %s"
const NamedConfigDoesNotExist = "named config '%s' does not exist"

const SentinelCountValue = -1

// This function is meant to support testability
func GetDefaultConfig() ClusterConfig {
	return defaultConfig
}

func assignConfig(res *ClusterConfig, src ClusterConfig) {
	if src.Name != "" {
		res.Name = src.Name
	}

	if src.MasterCount > SentinelCountValue {
		res.MasterCount = src.MasterCount
	}

	if src.WorkerCount > SentinelCountValue {
		res.WorkerCount = src.WorkerCount
	}

	if src.SparkMasterConfig != "" {
		res.SparkMasterConfig = src.SparkMasterConfig
	}
	if src.SparkWorkerConfig != "" {
		res.SparkWorkerConfig = src.SparkWorkerConfig
	}
	if src.SparkImage != "" {
		res.SparkImage = src.SparkImage
	}
	if src.ExposeWebUI != "" {
		res.ExposeWebUI = src.ExposeWebUI
	}
	if src.Metrics != "" {
		res.Metrics = src.Metrics
	}
}

func checkConfiguration(config ClusterConfig) error {
	var err error
	if config.MasterCount < 0 || config.MasterCount > 1 {
		err = NewClusterError(MasterCountMustBeZeroOrOne, ClusterConfigCode)
	} else if config.WorkerCount < 0 {
		err = NewClusterError(WorkerCountMustBeAtLeastZero, ClusterConfigCode)
	}
	return err
}

func getInt(value, configmapname string) (int, error) {
	i, err := strconv.Atoi(value)
	if err != nil {
		err = NewClusterError(fmt.Sprintf(ErrorWhileProcessing, configmapname, fmt.Sprintf("expected integer, got '%s'", value)), ClusterConfigCode)
	}
	return i, err
}

func process(config *ClusterConfig, name, value, configmapname string) error {

	var err error

	// At present we only have a single level of configs, but if/when we have
	// nested configs then we would descend through the levels beginning here with
	// the first element in the name
	switch name {
	case "mastercount":
		val, err := getInt(value, configmapname+".mastercount")
		if err != nil {
			return err
		}
		if val > SentinelCountValue {
			config.MasterCount = val
		}
	case "workercount":
		val, err := getInt(value, configmapname+".workercount")
		if err != nil {
			return err
		}
		if val > SentinelCountValue {
			config.WorkerCount = val
		}
	case "sparkmasterconfig":
		config.SparkMasterConfig = value
	case "sparkworkerconfig":
		config.SparkWorkerConfig = value
	case "sparkimage":
		config.SparkImage = value
	case "exposeui":
		config.ExposeWebUI = value
		_, err = strconv.ParseBool(config.ExposeWebUI)
		if err != nil {
			return err
		}
	case "metrics":
		// default will be "false" if the string is empty
		if value != "" {
			val, err := strconv.ParseBool(value)
			if err == nil {
				//  Support 'true' and default to jolokia
				// Normalize truth values
				if val {
					config.Metrics = "true"
				} else {
					config.Metrics = "false"
				}
			} else {
				if value != "jolokia" && value != "prometheus" {
					msg := fmt.Sprintf("expected 'jolokia' or 'prometheus', got '%s'", value)
					return NewClusterError(fmt.Sprintf(ErrorWhileProcessing, configmapname, msg), ClusterConfigCode)
				}
				config.Metrics = value
			}
		}
	default:
		err = NewClusterError(fmt.Sprintf(ErrorWhileProcessing, configmapname, "could not parse config fields"), ClusterConfigCode)
	}
	return err
}

func readConfig(name string, res *ClusterConfig, failOnMissing bool, restconfig *rest.Config, namespace string) (found bool, err error) {

	found = false
	cmap, err := getKubeClient(restconfig).CoreV1().ConfigMaps(namespace).Get(name, metav1.GetOptions{})
	if err != nil {
		if strings.Index(err.Error(), "not found") != -1 {
			if !failOnMissing {
				err = nil
			} else {
				err = NewClusterError(fmt.Sprintf(NamedConfigDoesNotExist, name), ClusterConfigCode)
			}
		} else {
			err = NewClusterError(err.Error(), ClientOperationCode)
		}
	}
	if err == nil && cmap != nil {
		// Kube will give us an empty configmap if the named one does not exist,
		// so we test for a Name to see if we found it
		found = cmap.Name != ""
		for n, v := range cmap.Data {
			err = process(res, strings.Trim(n, "\n"), strings.Trim(v, "\n"), name)
			if err != nil {
				return found, err
			}
		}
	}
	return found, err
}

func loadConfig(name string, restconfig *rest.Config, namespace string) (res ClusterConfig, err error) {
	// If the default config has been modified use those mods.
	res = defaultConfig
	found, err := readConfig(Defaultname, &res, allowMissing, restconfig, namespace)
	if err == nil {
		//process config if it is not named default
		//if there is a newly named config and an error then we create an error
		if name != "" && name != Defaultname {
			found, err = readConfig(name, &res, failOnMissing, restconfig, namespace)
			if !found{
				//then make an error something has gone wrong with the named config when read
				err =NewClusterError(fmt.Sprintf(NamedConfigDoesNotExist, name), ClusterConfigCode)
				return res, err
			}
		} else if found {
			res.Name = Defaultname
		}
	}
	return res, err
}

func GetClusterConfig(config *ClusterConfig, restconfig *rest.Config, namespace string) (res ClusterConfig, err error) {

	var name string = ""
	if config != nil {
		name = config.Name
	}
	res, err = loadConfig(name, restconfig, namespace)
	if err == nil && config != nil {
		assignConfig(&res, *config)
	}

	// Check that the final configuration is valid
	if err == nil {
		err = checkConfiguration(res)
	}
	return res, err
}
<|endoftext|>


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

package adb

import (
	"bufio"
	"bytes"
	"context"
	"encoding/base64"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/google/gapid/core/app"
	"github.com/google/gapid/core/fault"
	"github.com/google/gapid/core/log"
	"github.com/google/gapid/core/os/android"
	"github.com/google/gapid/core/os/device"
	"github.com/google/gapid/core/os/shell"
	"github.com/google/gapid/gapis/perfetto"

	common_pb "protos/perfetto/common"
)

const (
	// ErrDeviceNotRooted is returned by Device.Root when the device is running a
	// production build as is not 'rooted'.
	ErrDeviceNotRooted = fault.Const("Device is not a userdebug build")
	ErrRootFailed      = fault.Const("Device failed to switch to root")

	maxRootAttempts                         = 5
	gpuRenderStagesDataSourceDescriptorName = "gpu.renderstages"

	perfettoPort = NamedFileSystemSocket("/dev/socket/traced_consumer")
)

func isRootSuccessful(line string) bool {
	for _, expected := range []string{
		"adbd is already running as root",
		"* daemon started successfully *",
	} {
		if line == expected {
			return true
		}
	}
	return false
}

// Root restarts adb as root. If the device is running a production build then
// Root will return ErrDeviceNotRooted.
func (b *binding) Root(ctx context.Context) error {
	buf := bytes.Buffer{}
	buf.WriteString("adb root gave output:")
retry:
	for attempt := 0; attempt < maxRootAttempts; attempt++ {
		output, err := b.Command("root").Call(ctx)
		if err != nil {
			return err
		}
		if len(output) == 0 {
			return nil // Assume no output is success
		}
		output = strings.Replace(output, "\r\n", "\n", -1) // Not expected, but let's be safe.
		buf.WriteString(fmt.Sprintf("\n#%d: %v", attempt, output))
		lines := strings.Split(output, "\n")
		for i := len(lines) - 1; i >= 0; i-- {
			line := lines[i]
			if isRootSuccessful(line) {
				return nil // Success
			}
			switch line {
			case "adbd cannot run as root in production builds":
				return ErrDeviceNotRooted
			case "restarting adbd as root":
				time.Sleep(time.Millisecond * 100)
				continue retry
			default:
				// Some output we weren't expecting.
			}
		}
	}
	return log.Err(ctx, ErrRootFailed, buf.String())
}

// IsDebuggableBuild returns true if the device runs a debuggable Android build.
func (b *binding) IsDebuggableBuild(ctx context.Context) (bool, error) {
	output, err := b.Command("shell", "getprop", "ro.debuggable").Call(ctx)
	if err != nil {
		return false, err
	}
	return output == "1", nil
}

// InstallAPK installs the specified APK to the device. If reinstall is true
// and the package is already installed on the device then it will be replaced.
func (b *binding) InstallAPK(ctx context.Context, path string, reinstall bool, grantPermissions bool) error {
	args := []string{}
	if reinstall {
		args = append(args, "-r")
	}
	if grantPermissions && b.Instance().GetConfiguration().GetOS().GetAPIVersion() >= 23 {
		// Starting with API 23, permissions are not granted by default
		// during installation. Before API 23, the flag did not exist.
		args = append(args, "-g")
	}
	args = append(args, path)
	return b.Command("install", args...).Run(ctx)
}

// SELinuxEnforcing returns true if the device is currently in a
// SELinux enforcing mode, or false if the device is currently in a SELinux
// permissive mode.
func (b *binding) SELinuxEnforcing(ctx context.Context) (bool, error) {
	res, err := b.Shell("getenforce").Call(ctx)
	return strings.Contains(strings.ToLower(res), "enforcing"), err
}

// SetSELinuxEnforcing changes the SELinux-enforcing mode.
func (b *binding) SetSELinuxEnforcing(ctx context.Context, enforce bool) error {
	if enforce {
		return b.Shell("setenforce", "1").Run(ctx)
	}
	return b.Shell("setenforce", "0").Run(ctx)
}

// StartActivity launches the specified activity action.
func (b *binding) StartActivity(ctx context.Context, a android.ActivityAction, extras ...android.ActionExtra) error {
	args := append([]string{
		"start",
		"-S", // Force-stop the target app before starting the activity
		"-W", // Wait until the launch finishes
		"-a", a.Name,
		"-n", a.Component(),
	}, extrasFlags(extras)...)
	return b.Shell("am", args...).Run(ctx)
}

// StartActivityForDebug launches the specified activity in debug mode.
func (b *binding) StartActivityForDebug(ctx context.Context, a android.ActivityAction, extras ...android.ActionExtra) error {
	args := append([]string{
		"start",
		"-S", // Force-stop the target app before starting the activity
		"-W", // Wait until the launch finishes
		"-D", // Debug mode
		"-a", a.Name,
		"-n", a.Component(),
	}, extrasFlags(extras)...)
	return b.Shell("am", args...).Run(ctx)
}

// StartService launches the specified service action.
func (b *binding) StartService(ctx context.Context, a android.ServiceAction, extras ...android.ActionExtra) error {
	cmd := "start-foreground-service"
	if b.Instance().GetConfiguration().GetOS().GetAPIVersion() < 26 {
		// "am start-foreground-service" was added in API 26.
		cmd = "startservice"
	}
	args := append([]string{
		cmd,
		"-a", a.Name,
		"-n", a.Component(),
	}, extrasFlags(extras)...)
	return b.Shell("am", args...).Run(ctx)
}

// ForceStop stops everything associated with the given package.
func (b *binding) ForceStop(ctx context.Context, pkg string) error {
	return b.Shell("am", "force-stop", pkg).Run(ctx)
}

// SystemProperty returns the system property in string
func (b *binding) SystemProperty(ctx context.Context, name string) (string, error) {
	res, err := b.Shell("getprop", name).Call(ctx)
	if err != nil {
		return "", log.Errf(ctx, err, "getprop returned error: \n%s", err.Error())
	}
	return res, nil
}

// SetSystemProperty sets the system property with the given string value
func (b *binding) SetSystemProperty(ctx context.Context, name, value string) error {
	if len(value) == 0 {
		value = `""`
	}
	res, err := b.Shell("setprop", name, value).Call(ctx)
	if res != "" {
		return log.Errf(ctx, nil, "setprop returned error: \n%s", res)
	}
	if err != nil {
		return err
	}
	return nil
}

// SystemSetting returns the system setting with the given namespaced key.
func (b *binding) SystemSetting(ctx context.Context, namespace, key string) (string, error) {
	res, err := b.Shell("settings", "get", namespace, key).Call(ctx)
	if err != nil {
		return "", log.Errf(ctx, err, "settings get returned error: \n%s", err.Error())
	}
	return res, nil
}

// SetSystemSetting sets the system setting with with the given namespaced key
// to value.
func (b *binding) SetSystemSetting(ctx context.Context, namespace, key, value string) error {
	res, err := b.Shell("settings", "put", namespace, key, value).Call(ctx)
	if err != nil {
		return log.Errf(ctx, nil, "settings put returned error: \n%s", res)
	}
	return nil
}

// DeleteSystemSetting removes the system setting with with the given namespaced key.
func (b *binding) DeleteSystemSetting(ctx context.Context, namespace, key string) error {
	res, err := b.Shell("settings", "delete", namespace, key).Call(ctx)
	if err != nil {
		return log.Errf(ctx, nil, "settings delete returned error: \n%s", res)
	}
	return nil
}

// TempFile creates a temporary file on the given Device. It returns the
// path to the file, and a function that can be called to clean it up.
func (b *binding) TempFile(ctx context.Context) (string, func(ctx context.Context), error) {
	res, err := b.Shell("mktemp").Call(ctx)
	if err != nil {
		return "", nil, err
	}
	return res, func(ctx context.Context) {
		b.Shell("rm", "-f", res).Call(ctx)
	}, nil
}

// FileContents returns the contents of a given file on the Device.
func (b *binding) FileContents(ctx context.Context, path string) (string, error) {
	return b.Shell("cat", path).Call(ctx)
}

// RemoveFile removes the given file from the device
func (b *binding) RemoveFile(ctx context.Context, path string) error {
	_, err := b.Shell("rm", "-f", path).Call(ctx)
	return err
}

// GetEnv returns the default environment for the Device.
func (b *binding) GetEnv(ctx context.Context) (*shell.Env, error) {
	env, err := b.Shell("env").Call(ctx)
	if err != nil {
		return nil, err
	}
	scanner := bufio.NewScanner(strings.NewReader(env))
	e := &shell.Env{}
	for scanner.Scan() {
		e.Add(scanner.Text())
	}
	return e, nil
}

func (b *binding) SupportsPerfetto(ctx context.Context) bool {
	os := b.Instance().GetConfiguration().GetOS()
	return os.GetAPIVersion() >= 28
}

func (b *binding) ConnectPerfetto(ctx context.Context) (*perfetto.Client, error) {
	if !b.SupportsPerfetto(ctx) {
		return nil, fmt.Errorf("Perfetto is not supported on this device")
	}

	localPort, err := LocalFreeTCPPort()
	if err != nil {
		return nil, err
	}

	if err := b.Forward(ctx, localPort, perfettoPort); err != nil {
		return nil, err
	}
	cleanup := app.Cleanup(func(ctx context.Context) {
		b.RemoveForward(ctx, localPort)
	})

	conn, err := net.Dial("tcp", fmt.Sprintf("localhost:%v", localPort))
	if err != nil {
		cleanup.Invoke(ctx)
		return nil, err
	}
	return perfetto.NewClient(ctx, conn, cleanup)
}

func (b *binding) QueryPerfettoServiceState(ctx context.Context) (*device.PerfettoCapability, error) {
	result := b.To.Configuration.PerfettoCapability
	if b.Instance().GetConfiguration().GetOS().GetAPIVersion() < 29 {
		return result, log.Errf(ctx, nil, "Querying perfetto capability requires Android API >= 29")
	}
	res, err := b.Shell("perfetto", "--query-raw", "|", "base64").Call(ctx)
	if err != nil {
		return result, log.Errf(ctx, err, "adb shell perfetto returned error: %s", res)
	}
	decoded, _ := base64.StdEncoding.DecodeString(res)
	tracingServiceState := &common_pb.TracingServiceState{}
	if err = proto.Unmarshal(decoded, tracingServiceState); err != nil {
		return result, log.Errf(ctx, err, "Unmarshal returned error")
	}

	gpu := result.GpuProfiling
	if gpu == nil {
		gpu = &device.GPUProfiling{}
		result.GpuProfiling = gpu
	}

	dataSources := tracingServiceState.GetDataSources()
	for _, dataSource := range dataSources {
		dataSourceDescriptor := dataSource.GetDsDescriptor()
		if dataSourceDescriptor.GetName() == gpuRenderStagesDataSourceDescriptorName {
			gpu.HasRenderStage = true
			continue
		}
		gpuCounterDescriptor := dataSourceDescriptor.GetGpuCounterDescriptor()
		specs := gpuCounterDescriptor.GetSpecs()
		if len(specs) == 0 {
			continue
		}
		if gpu.GpuCounterDescriptor == nil {
			gpu.GpuCounterDescriptor = &device.GpuCounterDescriptor{}
		}

		// We mirror the Perfetto GpuCounterDescriptor proto into GAPID, hence
		// they are binary format compatible.
		data, err := proto.Marshal(gpuCounterDescriptor)
		if err != nil {
			continue
		}
		proto.UnmarshalMerge(data, gpu.GpuCounterDescriptor)
	}

	return result, nil
}

func extrasFlags(extras []android.ActionExtra) []string {
	flags := []string{}
	for _, e := range extras {
		flags = append(flags, e.Flags()...)
	}
	return flags
}
<commit_msg>Use the new Perfetto client to query the device.<commit_after>// Copyright (C) 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package adb

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/google/gapid/core/app"
	"github.com/google/gapid/core/fault"
	"github.com/google/gapid/core/log"
	"github.com/google/gapid/core/os/android"
	"github.com/google/gapid/core/os/device"
	"github.com/google/gapid/core/os/shell"
	"github.com/google/gapid/gapis/perfetto"

	common_pb "protos/perfetto/common"
)

const (
	// ErrDeviceNotRooted is returned by Device.Root when the device is running a
	// production build as is not 'rooted'.
	ErrDeviceNotRooted = fault.Const("Device is not a userdebug build")
	ErrRootFailed      = fault.Const("Device failed to switch to root")

	maxRootAttempts                         = 5
	gpuRenderStagesDataSourceDescriptorName = "gpu.renderstages"

	perfettoPort = NamedFileSystemSocket("/dev/socket/traced_consumer")
)

func isRootSuccessful(line string) bool {
	for _, expected := range []string{
		"adbd is already running as root",
		"* daemon started successfully *",
	} {
		if line == expected {
			return true
		}
	}
	return false
}

// Root restarts adb as root. If the device is running a production build then
// Root will return ErrDeviceNotRooted.
func (b *binding) Root(ctx context.Context) error {
	buf := bytes.Buffer{}
	buf.WriteString("adb root gave output:")
retry:
	for attempt := 0; attempt < maxRootAttempts; attempt++ {
		output, err := b.Command("root").Call(ctx)
		if err != nil {
			return err
		}
		if len(output) == 0 {
			return nil // Assume no output is success
		}
		output = strings.Replace(output, "\r\n", "\n", -1) // Not expected, but let's be safe.
		buf.WriteString(fmt.Sprintf("\n#%d: %v", attempt, output))
		lines := strings.Split(output, "\n")
		for i := len(lines) - 1; i >= 0; i-- {
			line := lines[i]
			if isRootSuccessful(line) {
				return nil // Success
			}
			switch line {
			case "adbd cannot run as root in production builds":
				return ErrDeviceNotRooted
			case "restarting adbd as root":
				time.Sleep(time.Millisecond * 100)
				continue retry
			default:
				// Some output we weren't expecting.
			}
		}
	}
	return log.Err(ctx, ErrRootFailed, buf.String())
}

// IsDebuggableBuild returns true if the device runs a debuggable Android build.
func (b *binding) IsDebuggableBuild(ctx context.Context) (bool, error) {
	output, err := b.Command("shell", "getprop", "ro.debuggable").Call(ctx)
	if err != nil {
		return false, err
	}
	return output == "1", nil
}

// InstallAPK installs the specified APK to the device. If reinstall is true
// and the package is already installed on the device then it will be replaced.
func (b *binding) InstallAPK(ctx context.Context, path string, reinstall bool, grantPermissions bool) error {
	args := []string{}
	if reinstall {
		args = append(args, "-r")
	}
	if grantPermissions && b.Instance().GetConfiguration().GetOS().GetAPIVersion() >= 23 {
		// Starting with API 23, permissions are not granted by default
		// during installation. Before API 23, the flag did not exist.
		args = append(args, "-g")
	}
	args = append(args, path)
	return b.Command("install", args...).Run(ctx)
}

// SELinuxEnforcing returns true if the device is currently in a
// SELinux enforcing mode, or false if the device is currently in a SELinux
// permissive mode.
func (b *binding) SELinuxEnforcing(ctx context.Context) (bool, error) {
	res, err := b.Shell("getenforce").Call(ctx)
	return strings.Contains(strings.ToLower(res), "enforcing"), err
}

// SetSELinuxEnforcing changes the SELinux-enforcing mode.
func (b *binding) SetSELinuxEnforcing(ctx context.Context, enforce bool) error {
	if enforce {
		return b.Shell("setenforce", "1").Run(ctx)
	}
	return b.Shell("setenforce", "0").Run(ctx)
}

// StartActivity launches the specified activity action.
func (b *binding) StartActivity(ctx context.Context, a android.ActivityAction, extras ...android.ActionExtra) error {
	args := append([]string{
		"start",
		"-S", // Force-stop the target app before starting the activity
		"-W", // Wait until the launch finishes
		"-a", a.Name,
		"-n", a.Component(),
	}, extrasFlags(extras)...)
	return b.Shell("am", args...).Run(ctx)
}

// StartActivityForDebug launches the specified activity in debug mode.
func (b *binding) StartActivityForDebug(ctx context.Context, a android.ActivityAction, extras ...android.ActionExtra) error {
	args := append([]string{
		"start",
		"-S", // Force-stop the target app before starting the activity
		"-W", // Wait until the launch finishes
		"-D", // Debug mode
		"-a", a.Name,
		"-n", a.Component(),
	}, extrasFlags(extras)...)
	return b.Shell("am", args...).Run(ctx)
}

// StartService launches the specified service action.
func (b *binding) StartService(ctx context.Context, a android.ServiceAction, extras ...android.ActionExtra) error {
	cmd := "start-foreground-service"
	if b.Instance().GetConfiguration().GetOS().GetAPIVersion() < 26 {
		// "am start-foreground-service" was added in API 26.
		cmd = "startservice"
	}
	args := append([]string{
		cmd,
		"-a", a.Name,
		"-n", a.Component(),
	}, extrasFlags(extras)...)
	return b.Shell("am", args...).Run(ctx)
}

// ForceStop stops everything associated with the given package.
func (b *binding) ForceStop(ctx context.Context, pkg string) error {
	return b.Shell("am", "force-stop", pkg).Run(ctx)
}

// SystemProperty returns the system property in string
func (b *binding) SystemProperty(ctx context.Context, name string) (string, error) {
	res, err := b.Shell("getprop", name).Call(ctx)
	if err != nil {
		return "", log.Errf(ctx, err, "getprop returned error: \n%s", err.Error())
	}
	return res, nil
}

// SetSystemProperty sets the system property with the given string value
func (b *binding) SetSystemProperty(ctx context.Context, name, value string) error {
	if len(value) == 0 {
		value = `""`
	}
	res, err := b.Shell("setprop", name, value).Call(ctx)
	if res != "" {
		return log.Errf(ctx, nil, "setprop returned error: \n%s", res)
	}
	if err != nil {
		return err
	}
	return nil
}

// SystemSetting returns the system setting with the given namespaced key.
func (b *binding) SystemSetting(ctx context.Context, namespace, key string) (string, error) {
	res, err := b.Shell("settings", "get", namespace, key).Call(ctx)
	if err != nil {
		return "", log.Errf(ctx, err, "settings get returned error: \n%s", err.Error())
	}
	return res, nil
}

// SetSystemSetting sets the system setting with with the given namespaced key
// to value.
func (b *binding) SetSystemSetting(ctx context.Context, namespace, key, value string) error {
	res, err := b.Shell("settings", "put", namespace, key, value).Call(ctx)
	if err != nil {
		return log.Errf(ctx, nil, "settings put returned error: \n%s", res)
	}
	return nil
}

// DeleteSystemSetting removes the system setting with with the given namespaced key.
func (b *binding) DeleteSystemSetting(ctx context.Context, namespace, key string) error {
	res, err := b.Shell("settings", "delete", namespace, key).Call(ctx)
	if err != nil {
		return log.Errf(ctx, nil, "settings delete returned error: \n%s", res)
	}
	return nil
}

// TempFile creates a temporary file on the given Device. It returns the
// path to the file, and a function that can be called to clean it up.
func (b *binding) TempFile(ctx context.Context) (string, func(ctx context.Context), error) {
	res, err := b.Shell("mktemp").Call(ctx)
	if err != nil {
		return "", nil, err
	}
	return res, func(ctx context.Context) {
		b.Shell("rm", "-f", res).Call(ctx)
	}, nil
}

// FileContents returns the contents of a given file on the Device.
func (b *binding) FileContents(ctx context.Context, path string) (string, error) {
	return b.Shell("cat", path).Call(ctx)
}

// RemoveFile removes the given file from the device
func (b *binding) RemoveFile(ctx context.Context, path string) error {
	_, err := b.Shell("rm", "-f", path).Call(ctx)
	return err
}

// GetEnv returns the default environment for the Device.
func (b *binding) GetEnv(ctx context.Context) (*shell.Env, error) {
	env, err := b.Shell("env").Call(ctx)
	if err != nil {
		return nil, err
	}
	scanner := bufio.NewScanner(strings.NewReader(env))
	e := &shell.Env{}
	for scanner.Scan() {
		e.Add(scanner.Text())
	}
	return e, nil
}

func (b *binding) SupportsPerfetto(ctx context.Context) bool {
	os := b.Instance().GetConfiguration().GetOS()
	return os.GetAPIVersion() >= 28
}

func (b *binding) ConnectPerfetto(ctx context.Context) (*perfetto.Client, error) {
	if !b.SupportsPerfetto(ctx) {
		return nil, fmt.Errorf("Perfetto is not supported on this device")
	}

	localPort, err := LocalFreeTCPPort()
	if err != nil {
		return nil, err
	}

	if err := b.Forward(ctx, localPort, perfettoPort); err != nil {
		return nil, err
	}
	cleanup := app.Cleanup(func(ctx context.Context) {
		b.RemoveForward(ctx, localPort)
	})

	conn, err := net.Dial("tcp", fmt.Sprintf("localhost:%v", localPort))
	if err != nil {
		cleanup.Invoke(ctx)
		return nil, err
	}
	return perfetto.NewClient(ctx, conn, cleanup)
}

func (b *binding) QueryPerfettoServiceState(ctx context.Context) (*device.PerfettoCapability, error) {
	result := b.To.Configuration.PerfettoCapability
	c, err := b.ConnectPerfetto(ctx)
	if err != nil {
		return result, err
	}
	defer c.Close(ctx)

	gpu := result.GpuProfiling
	if gpu == nil {
		gpu = &device.GPUProfiling{}
		result.GpuProfiling = gpu
	}

	err = c.Query(ctx, func(s *common_pb.TracingServiceState) error {
		for _, ds := range s.GetDataSources() {
			desc := ds.GetDsDescriptor()
			if desc.GetName() == gpuRenderStagesDataSourceDescriptorName {
				gpu.HasRenderStage = true
				continue
			}
			counters := desc.GetGpuCounterDescriptor().GetSpecs()
			if len(counters) != 0 {
				if gpu.GpuCounterDescriptor == nil {
					gpu.GpuCounterDescriptor = &device.GpuCounterDescriptor{}
				}
				// We mirror the Perfetto GpuCounterDescriptor proto into GAPID, hence
				// they are binary format compatible.
				data, err := proto.Marshal(desc.GetGpuCounterDescriptor())
				if err != nil {
					continue
				}
				proto.UnmarshalMerge(data, gpu.GpuCounterDescriptor)
			}
		}
		return nil
	})
	return result, err
}

func extrasFlags(extras []android.ActionExtra) []string {
	flags := []string{}
	for _, e := range extras {
		flags = append(flags, e.Flags()...)
	}
	return flags
}
<|endoftext|>


<commit_before>package oauth2

import (
  "encoding/json"
  "net/http"
  "net/url"
  "strings"
)

const (
  authorizationCode = "authorization_code"
)

type Storage struct {
  Client       ClientStore
  Code         TokenStore
  Token        TokenStore
  RefreshToken TokenStore
}

type Manager struct {
  CodeLife         int64
  TokenLife        int64
  RefreshTokenLife int64
  AllowGetMethod   bool

  Storage *Storage

  ClientAuthFunc func(r *http.Request, c *Client) bool
}

func (m *Manager) getClient(r *http.Request) (*Client, error) {
  clientId := r.Form.Get("client_id")
  client, err := m.Storage.Client.Read(clientId)
  if err != nil {
    return nil, NewAuthError(nil, E_SERVER_ERROR, err.Error())
  }

  if client == nil {
    return nil, NewAuthError(nil, E_UNAUTHORIZED_CLIENT,
      "Failed to read client")
  }

  return client, nil
}

func validateUri(base, uri string) bool {
  if base == "" || uri == "" {
    return false
  }

  // parse base url
  baseUri, err := url.Parse(base)
  if err != nil {
    return false
  }

  redirectUri, err := url.Parse(uri)
  if err != nil {
    return false
  }

  // must not have fragment
  if baseUri.Fragment != "" || redirectUri.Fragment != "" {
    return false
  }

  // check if urls match
  if baseUri.Scheme == redirectUri.Scheme && baseUri.Host == redirectUri.Host &&
    len(redirectUri.Path) >= len(baseUri.Path) &&
    strings.HasPrefix(redirectUri.Path, baseUri.Path) {
    return true
  }

  return false
}

func (m *Manager) GenerateCode(r *http.Request) (*Token, error) {
  err := r.ParseForm()
  if err != nil {
    return nil, NewAuthError(nil, E_INVALID_REQUEST, err.Error())
  }

  client, err := m.getClient(r)
  if err != nil {
    return nil, err
  }

  if r.Method != "POST" && !m.AllowGetMethod {
    return nil, NewAuthError(client, E_INVALID_REQUEST,
      "Invalid request method")
  }

  responseType := r.Form.Get("response_type")
  if responseType != "code" {
    return nil, NewAuthError(client, E_UNSUPPORTED_RESPONSE_TYPE,
      "Only code response type is supported now")
  }

  redirectUri := r.Form.Get("redirect_uri")
  if redirectUri == "" {
    redirectUri = client.BaseUri
  } else if !validateUri(client.BaseUri, redirectUri) {
    return nil, NewAuthError(client, E_INVALID_REQUEST, "Invalid redirect uri")
  }

  scope := r.Form.Get("scope")
  code := NewToken(client.Id, scope, redirectUri, m.CodeLife)
  return code, nil
}

func (m *Manager) SaveCode(code *Token) error {
  _, err := m.Storage.Code.Save(code)
  if err != nil {
    client, _ := m.Storage.Client.Read(code.ClientId)
    return NewAuthError(client, E_SERVER_ERROR, err.Error())
  }
  return nil
}

func (m *Manager) RedirectUrlWithCode(code *Token) (*url.URL, error) {
  uri, err := url.Parse(code.RedirectUri)
  if err != nil {
    client, _ := m.Storage.Client.Read(code.ClientId)
    return nil, NewAuthError(client, E_INVALID_REQUEST, "Invalid redirect uri")
  }
  q := uri.Query()
  q.Set("state", code.State)
  q.Set("code", code.Value)
  uri.RawQuery = q.Encode()
  return uri, nil
}

func (m *Manager) GenerateToken(r *http.Request) (*Token, *Token, error) {
  err := r.ParseForm()
  if err != nil {
    return nil, nil, NewAuthError(nil, E_INVALID_REQUEST, err.Error())
  }

  client, err := m.getClient(r)
  if err != nil {
    return nil, nil, err
  }

  if !m.ClientAuthFunc(r, client) {
    return nil, nil, NewAuthError(client, E_UNAUTHORIZED_CLIENT,
      "Failed to validate client")
  }

  if r.Method != "POST" && !m.AllowGetMethod {
    return nil, nil, NewAuthError(client, E_INVALID_REQUEST,
      "Invalid request method")
  }

  responseType := r.Form.Get("grant_type")
  if responseType != authorizationCode {
    return nil, nil, NewAuthError(client, E_UNSUPPORTED_GRANT_TYPE,
      "Only authorization code grant type is supported now")
  }

  code, err := m.Storage.Code.Read(r.Form.Get("code"))
  if err != nil {
    return nil, nil, NewAuthError(client, E_SERVER_ERROR, err.Error())
  }
  if code == nil {
    return nil, nil, NewAuthError(client, E_INVALID_GRANT, "Invalid code")
  }

  if code.ClientId != client.Id {
    return nil, nil, NewAuthError(client, E_INVALID_CLIENT, "Client is mismatch")
  }

  redirectUri := r.Form.Get("redirect_uri")
  if redirectUri == "" {
    redirectUri = client.BaseUri
  } else if !validateUri(client.BaseUri, redirectUri) {
    return nil, nil, NewAuthError(client, E_INVALID_REQUEST, "Invalid redirect uri")
  } else if code.RedirectUri != redirectUri {
    return nil, nil, NewAuthError(client, E_INVALID_REQUEST, "Invalid redirect uri")
  }

  token := NewToken(client.Id, code.Scope, redirectUri, m.TokenLife)
  return code, token, nil
}

func (m *Manager) SaveToken(token *Token) error {
  _, err := m.Storage.Token.Save(token)
  if err != nil {
    client, _ := m.Storage.Client.Read(token.ClientId)
    return NewAuthError(client, E_SERVER_ERROR, err.Error())
  }
  return nil
}

func (m *Manager) ResponseWithToken(w http.ResponseWriter,
  token *Token, userData map[string]interface{}) error {
  s := make(map[string]interface{})
  s["scope"] = token.Scope
  s["access_token"] = token.Value
  s["expires_in"] = token.Life
  for k, v := range userData {
    s[k] = v
  }
  o, err := json.Marshal(s)
  if err != nil {
    return err
  }
  w.Header().Set("Content-Type", "application/json;charset=UTF-8")
  w.Header().Set("Cache-Control", "no-store")
  w.Header().Set("Pragma", "no-cache")
  w.Write(o)
  return nil
}

func (m *Manager) ResponseWithError(w http.ResponseWriter, err error) error {
  oerr, ok := err.(*AuthError)
  var errString string
  if ok {
    errString = oerr.ErrorString()
  } else {
    errString = err.Error()
  }
  s := make(map[string]interface{})
  s["error"] = errString
  o, err := json.Marshal(s)
  if err != nil {
    return err
  }
  w.WriteHeader(http.StatusBadRequest)
  w.Header().Set("Content-Type", "application/json;charset=UTF-8")
  w.Header().Set("Cache-Control", "no-store")
  w.Header().Set("Pragma", "no-cache")
  w.Write(o)
  return nil
}
<commit_msg>Always allow post method for code request<commit_after>package oauth2

import (
  "encoding/json"
  "net/http"
  "net/url"
  "strings"
)

const (
  authorizationCode = "authorization_code"
)

type Storage struct {
  Client       ClientStore
  Code         TokenStore
  Token        TokenStore
  RefreshToken TokenStore
}

type Manager struct {
  CodeLife         int64
  TokenLife        int64
  RefreshTokenLife int64
  AllowGetMethod   bool

  Storage *Storage

  ClientAuthFunc func(r *http.Request, c *Client) bool
}

func (m *Manager) getClient(r *http.Request) (*Client, error) {
  clientId := r.Form.Get("client_id")
  client, err := m.Storage.Client.Read(clientId)
  if err != nil {
    return nil, NewAuthError(nil, E_SERVER_ERROR, err.Error())
  }

  if client == nil {
    return nil, NewAuthError(nil, E_UNAUTHORIZED_CLIENT,
      "Failed to read client")
  }

  return client, nil
}

func validateUri(base, uri string) bool {
  if base == "" || uri == "" {
    return false
  }

  // parse base url
  baseUri, err := url.Parse(base)
  if err != nil {
    return false
  }

  redirectUri, err := url.Parse(uri)
  if err != nil {
    return false
  }

  // must not have fragment
  if baseUri.Fragment != "" || redirectUri.Fragment != "" {
    return false
  }

  // check if urls match
  if baseUri.Scheme == redirectUri.Scheme && baseUri.Host == redirectUri.Host &&
    len(redirectUri.Path) >= len(baseUri.Path) &&
    strings.HasPrefix(redirectUri.Path, baseUri.Path) {
    return true
  }

  return false
}

func (m *Manager) GenerateCode(r *http.Request) (*Token, error) {
  err := r.ParseForm()
  if err != nil {
    return nil, NewAuthError(nil, E_INVALID_REQUEST, err.Error())
  }

  client, err := m.getClient(r)
  if err != nil {
    return nil, err
  }

  responseType := r.Form.Get("response_type")
  if responseType != "code" {
    return nil, NewAuthError(client, E_UNSUPPORTED_RESPONSE_TYPE,
      "Only code response type is supported now")
  }

  redirectUri := r.Form.Get("redirect_uri")
  if redirectUri == "" {
    redirectUri = client.BaseUri
  } else if !validateUri(client.BaseUri, redirectUri) {
    return nil, NewAuthError(client, E_INVALID_REQUEST, "Invalid redirect uri")
  }

  scope := r.Form.Get("scope")
  code := NewToken(client.Id, scope, redirectUri, m.CodeLife)
  return code, nil
}

func (m *Manager) SaveCode(code *Token) error {
  _, err := m.Storage.Code.Save(code)
  if err != nil {
    client, _ := m.Storage.Client.Read(code.ClientId)
    return NewAuthError(client, E_SERVER_ERROR, err.Error())
  }
  return nil
}

func (m *Manager) RedirectUrlWithCode(code *Token) (*url.URL, error) {
  uri, err := url.Parse(code.RedirectUri)
  if err != nil {
    client, _ := m.Storage.Client.Read(code.ClientId)
    return nil, NewAuthError(client, E_INVALID_REQUEST, "Invalid redirect uri")
  }
  q := uri.Query()
  q.Set("state", code.State)
  q.Set("code", code.Value)
  uri.RawQuery = q.Encode()
  return uri, nil
}

func (m *Manager) GenerateToken(r *http.Request) (*Token, *Token, error) {
  err := r.ParseForm()
  if err != nil {
    return nil, nil, NewAuthError(nil, E_INVALID_REQUEST, err.Error())
  }

  client, err := m.getClient(r)
  if err != nil {
    return nil, nil, err
  }

  if !m.ClientAuthFunc(r, client) {
    return nil, nil, NewAuthError(client, E_UNAUTHORIZED_CLIENT,
      "Failed to validate client")
  }

  if r.Method != "POST" && !m.AllowGetMethod {
    return nil, nil, NewAuthError(client, E_INVALID_REQUEST,
      "Invalid request method")
  }

  responseType := r.Form.Get("grant_type")
  if responseType != authorizationCode {
    return nil, nil, NewAuthError(client, E_UNSUPPORTED_GRANT_TYPE,
      "Only authorization code grant type is supported now")
  }

  code, err := m.Storage.Code.Read(r.Form.Get("code"))
  if err != nil {
    return nil, nil, NewAuthError(client, E_SERVER_ERROR, err.Error())
  }
  if code == nil {
    return nil, nil, NewAuthError(client, E_INVALID_GRANT, "Invalid code")
  }

  if code.ClientId != client.Id {
    return nil, nil, NewAuthError(client, E_INVALID_CLIENT, "Client is mismatch")
  }

  redirectUri := r.Form.Get("redirect_uri")
  if redirectUri == "" {
    redirectUri = client.BaseUri
  } else if !validateUri(client.BaseUri, redirectUri) {
    return nil, nil, NewAuthError(client, E_INVALID_REQUEST, "Invalid redirect uri")
  } else if code.RedirectUri != redirectUri {
    return nil, nil, NewAuthError(client, E_INVALID_REQUEST, "Invalid redirect uri")
  }

  token := NewToken(client.Id, code.Scope, redirectUri, m.TokenLife)
  return code, token, nil
}

func (m *Manager) SaveToken(token *Token) error {
  _, err := m.Storage.Token.Save(token)
  if err != nil {
    client, _ := m.Storage.Client.Read(token.ClientId)
    return NewAuthError(client, E_SERVER_ERROR, err.Error())
  }
  return nil
}

func (m *Manager) ResponseWithToken(w http.ResponseWriter,
  token *Token, userData map[string]interface{}) error {
  s := make(map[string]interface{})
  s["scope"] = token.Scope
  s["access_token"] = token.Value
  s["expires_in"] = token.Life
  for k, v := range userData {
    s[k] = v
  }
  o, err := json.Marshal(s)
  if err != nil {
    return err
  }
  w.Header().Set("Content-Type", "application/json;charset=UTF-8")
  w.Header().Set("Cache-Control", "no-store")
  w.Header().Set("Pragma", "no-cache")
  w.Write(o)
  return nil
}

func (m *Manager) ResponseWithError(w http.ResponseWriter, err error) error {
  oerr, ok := err.(*AuthError)
  var errString string
  if ok {
    errString = oerr.ErrorString()
  } else {
    errString = err.Error()
  }
  s := make(map[string]interface{})
  s["error"] = errString
  o, err := json.Marshal(s)
  if err != nil {
    return err
  }
  w.WriteHeader(http.StatusBadRequest)
  w.Header().Set("Content-Type", "application/json;charset=UTF-8")
  w.Header().Set("Cache-Control", "no-store")
  w.Header().Set("Pragma", "no-cache")
  w.Write(o)
  return nil
}
<|endoftext|>


<commit_before>package pq

import (
	"bytes"
	"strconv"
)

// Let table = 'schema.device_table',
//     pks = ['id1', 'id2'],
//     data = {'type': 'ios', 'token': 'sometoken', 'userid': 'someuserid'}
//
// upsertQuery generates a query for upsert in the following format:
//
//	WITH updated AS (
//		UPDATE schema.device_table
//			SET ("type", "token", "user_id") =
//			($3, $4, $5)
//			WHERE "id1" = $1 AND "id2" = $2
//			RETURNING *
//		)
//	INSERT schema.device_table
//		("id1", "id2", "type", "token", "user_id")
//		SELECT $1, $2, $3, $4, $5
//		WHERE NOT EXISTS (SELECT * FROM updated);
//
// And args = ['1', '2', 'ios', 'sometoken', 'someuserid']
//
// This approach uses CTE to do an INSERT after UPDATE in one query,
// hoping that the time gap between the UPDATE and INSERT is short
// enough that chance of a concurrent insert is rare.
//
// A complete upsert example is included in postgresql documentation [1],
// but that implementation contains a loop that does not guarantee
// exit. Adding that to poor performance, that implementation is not
// adopted.
//
// More on UPSERT: https://wiki.postgresql.org/wiki/UPSERT#PostgreSQL_.28today.29
//
// [1]: http://www.postgresql.org/docs/9.4/static/plpgsql-control-structures.html#PLPGSQL-UPSERT-EXAMPLE
func upsertQuery(table string, pkData map[string]interface{}, data map[string]interface{}, updateIgnore []string) (sql string, args []interface{}) {
	// extract columns values pair
	pks, pkArgs := extractKeyAndValue(pkData)
	columns, args := extractKeyAndValue(data)
	ignoreIndex := findIgnoreIndex(columns, updateIgnore)

	// generate WITH UPDATE
	b := bytes.Buffer{}
	b.Write([]byte(`WITH updated AS (UPDATE `))
	b.WriteString(table)
	b.Write([]byte(` SET(`))

	for i, column := range columns {
		if ignoreIndex[i] {
			continue
		}
		b.WriteByte('"')
		b.WriteString(column)
		b.Write([]byte(`",`))
	}
	b.Truncate(b.Len() - 1)

	b.Write([]byte(`)=(`))

	for i := len(pks); i < len(pks)+len(columns); i++ {
		if ignoreIndex[i-len(pks)] {
			continue
		}
		b.WriteByte('$')
		b.WriteString(strconv.Itoa(i + 1))
		b.WriteByte(',')
	}
	b.Truncate(b.Len() - 1)

	b.Write([]byte(`) WHERE `))

	for i, pk := range pks {
		b.WriteByte('"')
		b.WriteString(pk)
		b.Write([]byte(`" = $`))
		b.WriteString(strconv.Itoa(i + 1))
		b.Write([]byte(` AND `))
	}
	b.Truncate(b.Len() - 5)

	// generate INSERT
	b.Write([]byte(` RETURNING *) INSERT INTO `))
	b.WriteString(table)
	b.WriteByte('(')

	for _, column := range append(pks, columns...) {
		b.WriteByte('"')
		b.WriteString(column)
		b.Write([]byte(`",`))
	}
	b.Truncate(b.Len() - 1)

	b.Write([]byte(`) SELECT `))

	for i := 0; i < len(pks)+len(columns); i++ {
		b.WriteByte('$')
		b.WriteString(strconv.Itoa(i + 1))
		b.WriteByte(',')
	}
	b.Truncate(b.Len() - 1)

	b.Write([]byte(` WHERE NOT EXISTS (SELECT * FROM updated);`))

	return b.String(), append(pkArgs, args...)
}

func extractKeyAndValue(data map[string]interface{}) (keys []string, values []interface{}) {
	keys = make([]string, len(data), len(data))
	values = make([]interface{}, len(data), len(data))

	i := 0
	for key, value := range data {
		keys[i] = key
		values[i] = value
		i++
	}

	return
}

func findIgnoreIndex(columns []string, ignoreColumns []string) (ignoreIndex []bool) {
	ignoreIndex = make([]bool, len(columns), len(columns))

	for i, column := range columns {
		for _, ignored := range ignoreColumns {
			ignoreIndex[i] = (column == ignored)
			break
		}
	}

	return
}
<commit_msg>Modify upsert query to support empty data<commit_after>package pq

import (
	"bytes"
	"strconv"
)

// Let table = 'schema.device_table',
//     pks = ['id1', 'id2'],
//     data = {'type': 'ios', 'token': 'sometoken', 'userid': 'someuserid'}
//
// upsertQuery generates a query for upsert in the following format:
//
//	WITH updated AS (
//		UPDATE schema.device_table
//			SET ("type", "token", "user_id") =
//			($3, $4, $5)
//			WHERE "id1" = $1 AND "id2" = $2
//			RETURNING *
//		)
//	INSERT INTO schema.device_table
//		("id1", "id2", "type", "token", "user_id")
//		SELECT $1, $2, $3, $4, $5
//		WHERE NOT EXISTS (SELECT * FROM updated);
//
// And args = ['1', '2', 'ios', 'sometoken', 'someuserid']
//
// For empty data, following will be generated
//	WITH updated AS (
//		SELECT "id1", "id2" FROM schema.device_table
//		WHERE "id1" = $1 AND "id2" = $2
//	)
//	INSERT INTO schema.device_table
//		("id1", "id2")
//		SELECT $1, $2
//		WHERE NOT EXISTS (SELECT * FROM updated);
//
// And args = ['1', '2', 'ios', 'sometoken', 'someuserid']
//
// This approach uses CTE to do an INSERT after UPDATE in one query,
// hoping that the time gap between the UPDATE and INSERT is short
// enough that chance of a concurrent insert is rare.
//
// A complete upsert example is included in postgresql documentation [1],
// but that implementation contains a loop that does not guarantee
// exit. Adding that to poor performance, that implementation is not
// adopted.
//
// More on UPSERT: https://wiki.postgresql.org/wiki/UPSERT#PostgreSQL_.28today.29
//
// [1]: http://www.postgresql.org/docs/9.4/static/plpgsql-control-structures.html#PLPGSQL-UPSERT-EXAMPLE
func upsertQuery(table string, pkData map[string]interface{}, data map[string]interface{}, updateIgnore []string) (sql string, args []interface{}) {
	// extract columns values pair
	pks, pkArgs := extractKeyAndValue(pkData)
	columns, args := extractKeyAndValue(data)
	ignoreIndex := findIgnoreIndex(columns, updateIgnore)

	b := bytes.Buffer{}
	if len(columns) > 0 {
		// Generate with UPDATE
		b.Write([]byte(`WITH updated AS (UPDATE `))
		b.WriteString(table)
		b.Write([]byte(` SET(`))

		for i, column := range columns {
			if ignoreIndex[i] {
				continue
			}
			b.WriteByte('"')
			b.WriteString(column)
			b.Write([]byte(`",`))
		}
		b.Truncate(b.Len() - 1)

		b.Write([]byte(`)=(`))

		for i := len(pks); i < len(pks)+len(columns); i++ {
			if ignoreIndex[i-len(pks)] {
				continue
			}
			b.WriteByte('$')
			b.WriteString(strconv.Itoa(i + 1))
			b.WriteByte(',')
		}
		b.Truncate(b.Len() - 1)

		b.Write([]byte(`) WHERE `))

		for i, pk := range pks {
			b.WriteByte('"')
			b.WriteString(pk)
			b.Write([]byte(`" = $`))
			b.WriteString(strconv.Itoa(i + 1))
			b.Write([]byte(` AND `))
		}
		b.Truncate(b.Len() - 5)
		b.Write([]byte(` RETURNING *) `))
	} else {
		// Generate with SELECT
		b.Write([]byte(`WITH updated AS (SELECT `))
		for _, pk := range pks {
			b.WriteByte('"')
			b.WriteString(pk)
			b.Write([]byte(`",`))
		}
		b.Truncate(b.Len() - 1)
		b.Write([]byte(` FROM `))
		b.WriteString(table)
		b.Write([]byte(` WHERE `))
		for i, pk := range pks {
			b.WriteByte('"')
			b.WriteString(pk)
			b.Write([]byte(`" = $`))
			b.WriteString(strconv.Itoa(i + 1))
			b.Write([]byte(` AND `))
		}
		b.Truncate(b.Len() - 5)
		b.Write([]byte(`) `))
	}

	// generate INSERT
	b.Write([]byte(`INSERT INTO `))
	b.WriteString(table)
	b.WriteByte('(')

	for _, column := range append(pks, columns...) {
		b.WriteByte('"')
		b.WriteString(column)
		b.Write([]byte(`",`))
	}
	b.Truncate(b.Len() - 1)

	b.Write([]byte(`) SELECT `))

	for i := 0; i < len(pks)+len(columns); i++ {
		b.WriteByte('$')
		b.WriteString(strconv.Itoa(i + 1))
		b.WriteByte(',')
	}
	b.Truncate(b.Len() - 1)

	b.Write([]byte(` WHERE NOT EXISTS (SELECT * FROM updated);`))

	return b.String(), append(pkArgs, args...)
}

func extractKeyAndValue(data map[string]interface{}) (keys []string, values []interface{}) {
	keys = make([]string, len(data), len(data))
	values = make([]interface{}, len(data), len(data))

	i := 0
	for key, value := range data {
		keys[i] = key
		values[i] = value
		i++
	}

	return
}

func findIgnoreIndex(columns []string, ignoreColumns []string) (ignoreIndex []bool) {
	ignoreIndex = make([]bool, len(columns), len(columns))

	for i, column := range columns {
		for _, ignored := range ignoreColumns {
			ignoreIndex[i] = (column == ignored)
			break
		}
	}

	return
}
<|endoftext|>