<commit_before>package api

import (
	"database/sql"
	"fmt"
	"image/color"
	"image/png"
	"log"
	"net/http"
	"net/url"
	"strconv"

	"github.com/mdlayher/wavepipe/data"

	"github.com/gorilla/context"
	"github.com/gorilla/mux"
	"github.com/hashicorp/golang-lru"
	"github.com/mdlayher/waveform"
	"github.com/nfnt/resize"
	"github.com/unrolled/render"
)

// waveformLRU is a LRU cache which stores a fixed number of recently computed waveform image
// values, and evicts the least-recently-used entries when they are not used
var waveformLRU *lru.Cache

func init() {
	// Initialize fixed-capacity LRU cache
	var err error
	waveformLRU, err = lru.New(20)
	if err != nil {
		panic(err)
	}
}

// GetWaveform generates and returns a waveform image from wavepipe.  On success, this API will
// return a binary stream. On failure, it will return a JSON error.
func GetWaveform(w http.ResponseWriter, r *http.Request) {
	// Retrieve render
	ren := context.Get(r, CtxRender).(*render.Render)

	// Check API version
	if version, ok := mux.Vars(r)["version"]; ok {
		// Check if this API call is supported in the advertised version
		if !apiVersionSet.Has(version) {
			ren.JSON(w, 400, errRes(400, "unsupported API version: "+version))
			return
		}
	}

	// Check for an ID parameter
	pID, ok := mux.Vars(r)["id"]
	if !ok {
		ren.JSON(w, 400, errRes(400, "no integer song ID provided"))
		return
	}

	// Verify valid integer ID
	id, err := strconv.Atoi(pID)
	if err != nil {
		ren.JSON(w, 400, errRes(400, "invalid integer song ID"))
		return
	}

	// Check for a pre-existing set of waveform values
	var values []float64
	if tmpValues, ok := waveformLRU.Get(id); ok {
		// Use existing values
		values = tmpValues.([]float64)
	} else {
		// No existing waveform
		// Attempt to load the song with matching ID
		song := &data.Song{ID: id}
		if err := song.Load(); err != nil {
			// Check for invalid ID
			if err == sql.ErrNoRows {
				ren.JSON(w, 404, errRes(404, "song ID not found"))
				return
			}

			// All other errors
			log.Println(err)
			ren.JSON(w, 500, serverErr)
			return
		}

		// Open song's backing stream
		stream, err := song.Stream()
		if err != nil {
			log.Println(err)
			ren.JSON(w, 500, serverErr)
			return
		}

		// Generate a waveform object
		wave, err := waveform.New(stream, waveform.Resolution(4))
		if err != nil {
			log.Println(err)
			ren.JSON(w, 500, serverErr)
			return
		}

		// Compute waveform values from input stream
		tmpValues, err := wave.Compute()
		if err != nil {
			// If unknown format, return JSON error
			if err == waveform.ErrFormat {
				ren.JSON(w, 501, errRes(501, "unsupported audio format"))
				return
			}

			log.Println(err)
			ren.JSON(w, 500, serverErr)
			return
		}

		// Cache values, and use for generating current and future images
		waveformLRU.Add(id, tmpValues)
		values = tmpValues
	}

	// Check for optional color parameters
	var cR, cG, cB uint8

	// Background color
	var bgColor color.Color = color.White
	if bgColorStr := r.URL.Query().Get("bg"); bgColorStr != "" {
		// Convert %23 to #
		bgColorStr, err := url.QueryUnescape(bgColorStr)
		if err == nil {
			cR, cG, cB = hexToRGB(bgColorStr)
			bgColor = color.RGBA{cR, cG, cB, 255}
		}
	}

	// Foreground color
	var fgColor color.Color = color.Black
	if fgColorStr := r.URL.Query().Get("fg"); fgColorStr != "" {
		// Convert %23 to #
		fgColorStr, err := url.QueryUnescape(fgColorStr)
		if err == nil {
			cR, cG, cB = hexToRGB(fgColorStr)
			fgColor = color.RGBA{cR, cG, cB, 255}
		}
	}

	// Alternate color; follow foreground color by default
	var altColor color.Color = fgColor
	if altColorStr := r.URL.Query().Get("alt"); altColorStr != "" {
		// Convert %23 to #
		altColorStr, err := url.QueryUnescape(altColorStr)
		if err == nil {
			cR, cG, cB = hexToRGB(altColorStr)
			altColor = color.RGBA{cR, cG, cB, 255}
		}
	}

	// If requested, resize the image to the specified width
	var sizeX, sizeY int
	if strSize := r.URL.Query().Get("size"); strSize != "" {
		// Check for dimensions in two integers
		if _, err := fmt.Sscanf(strSize, "%dx%d", &sizeX, &sizeY); err != nil {
			ren.JSON(w, 400, errRes(400, "invalid x-separated integer pair for size"))
			return
		}
	}

	// Generate waveform object with sane defaults and user settings
	wave, err := waveform.New(nil,
		waveform.Colors(fgColor, bgColor, altColor),
		waveform.Scale(5, 4),
		waveform.Sharpness(1),
		waveform.ScaleClipping(),
	)
	if err != nil {
		log.Println(err)
		ren.JSON(w, 500, serverErr)
		return
	}

	// Generate waveform image from computed values, with specified options
	img := wave.Draw(values)

	// If a resize option was set, perform it now
	if sizeX > 0 {
		// Perform image resize
		img = resize.Resize(uint(sizeX), uint(sizeY), img, resize.NearestNeighbor)
	}

	// Encode as PNG to HTTP writer
	if err := png.Encode(w, img); err != nil {
		log.Println(err)
	}
}

// hexToRGB converts a hex string to a RGB triple.
// Credit: https://code.google.com/p/gorilla/source/browse/color/hex.go?r=ef489f63418265a7249b1d53bdc358b09a4a2ea0
func hexToRGB(h string) (uint8, uint8, uint8) {
	if len(h) > 0 && h[0] == '#' {
		h = h[1:]
	}
	if len(h) == 3 {
		h = h[:1] + h[:1] + h[1:2] + h[1:2] + h[2:] + h[2:]
	}
	if len(h) == 6 {
		if rgb, err := strconv.ParseUint(string(h), 16, 32); err == nil {
			return uint8(rgb >> 16), uint8((rgb >> 8) & 0xFF), uint8(rgb & 0xFF)
		}
	}
	return 0, 0, 0
}
<commit_msg>api/waveform: update to new mdlayher/waveform API<commit_after>package api

import (
	"database/sql"
	"fmt"
	"image/color"
	"image/png"
	"log"
	"net/http"
	"net/url"
	"strconv"

	"github.com/mdlayher/wavepipe/data"

	"github.com/gorilla/context"
	"github.com/gorilla/mux"
	"github.com/hashicorp/golang-lru"
	"github.com/mdlayher/waveform"
	"github.com/nfnt/resize"
	"github.com/unrolled/render"
)

// waveformLRU is a LRU cache which stores a fixed number of recently computed waveform image
// values, and evicts the least-recently-used entries when they are not used
var waveformLRU *lru.Cache

func init() {
	// Initialize fixed-capacity LRU cache
	var err error
	waveformLRU, err = lru.New(20)
	if err != nil {
		panic(err)
	}
}

// GetWaveform generates and returns a waveform image from wavepipe.  On success, this API will
// return a binary stream. On failure, it will return a JSON error.
func GetWaveform(w http.ResponseWriter, r *http.Request) {
	// Retrieve render
	ren := context.Get(r, CtxRender).(*render.Render)

	// Check API version
	if version, ok := mux.Vars(r)["version"]; ok {
		// Check if this API call is supported in the advertised version
		if !apiVersionSet.Has(version) {
			ren.JSON(w, 400, errRes(400, "unsupported API version: "+version))
			return
		}
	}

	// Check for an ID parameter
	pID, ok := mux.Vars(r)["id"]
	if !ok {
		ren.JSON(w, 400, errRes(400, "no integer song ID provided"))
		return
	}

	// Verify valid integer ID
	id, err := strconv.Atoi(pID)
	if err != nil {
		ren.JSON(w, 400, errRes(400, "invalid integer song ID"))
		return
	}

	// Check for a pre-existing set of waveform values
	var values []float64
	if tmpValues, ok := waveformLRU.Get(id); ok {
		// Use existing values
		values = tmpValues.([]float64)
	} else {
		// No existing waveform
		// Attempt to load the song with matching ID
		song := &data.Song{ID: id}
		if err := song.Load(); err != nil {
			// Check for invalid ID
			if err == sql.ErrNoRows {
				ren.JSON(w, 404, errRes(404, "song ID not found"))
				return
			}

			// All other errors
			log.Println(err)
			ren.JSON(w, 500, serverErr)
			return
		}

		// Open song's backing stream
		stream, err := song.Stream()
		if err != nil {
			log.Println(err)
			ren.JSON(w, 500, serverErr)
			return
		}

		// Generate a waveform object
		wave, err := waveform.New(stream, waveform.Resolution(4))
		if err != nil {
			log.Println(err)
			ren.JSON(w, 500, serverErr)
			return
		}

		// Compute waveform values from input stream
		tmpValues, err := wave.Compute()
		if err != nil {
			// If unknown format, return JSON error
			if err == waveform.ErrFormat {
				ren.JSON(w, 501, errRes(501, "unsupported audio format"))
				return
			}

			log.Println(err)
			ren.JSON(w, 500, serverErr)
			return
		}

		// Cache values, and use for generating current and future images
		waveformLRU.Add(id, tmpValues)
		values = tmpValues
	}

	// Check for optional color parameters
	var cR, cG, cB uint8

	// Background color
	var bgColor color.Color = color.White
	if bgColorStr := r.URL.Query().Get("bg"); bgColorStr != "" {
		// Convert %23 to #
		bgColorStr, err := url.QueryUnescape(bgColorStr)
		if err == nil {
			cR, cG, cB = hexToRGB(bgColorStr)
			bgColor = color.RGBA{cR, cG, cB, 255}
		}
	}

	// Foreground color
	var fgColor color.Color = color.Black
	if fgColorStr := r.URL.Query().Get("fg"); fgColorStr != "" {
		// Convert %23 to #
		fgColorStr, err := url.QueryUnescape(fgColorStr)
		if err == nil {
			cR, cG, cB = hexToRGB(fgColorStr)
			fgColor = color.RGBA{cR, cG, cB, 255}
		}
	}

	// Alternate color; follow foreground color by default
	var altColor color.Color = fgColor
	if altColorStr := r.URL.Query().Get("alt"); altColorStr != "" {
		// Convert %23 to #
		altColorStr, err := url.QueryUnescape(altColorStr)
		if err == nil {
			cR, cG, cB = hexToRGB(altColorStr)
			altColor = color.RGBA{cR, cG, cB, 255}
		}
	}

	// If requested, resize the image to the specified width
	var sizeX, sizeY int
	if strSize := r.URL.Query().Get("size"); strSize != "" {
		// Check for dimensions in two integers
		if _, err := fmt.Sscanf(strSize, "%dx%d", &sizeX, &sizeY); err != nil {
			ren.JSON(w, 400, errRes(400, "invalid x-separated integer pair for size"))
			return
		}
	}

	// Generate waveform object with sane defaults and user settings
	wave, err := waveform.New(nil,
		waveform.BGColorFunction(waveform.SolidColor(bgColor)),
		waveform.FGColorFunction(waveform.StripeColor(fgColor, altColor)),
		waveform.Scale(5, 4),
		waveform.Sharpness(1),
		waveform.ScaleClipping(),
	)
	if err != nil {
		log.Println(err)
		ren.JSON(w, 500, serverErr)
		return
	}

	// Generate waveform image from computed values, with specified options
	img := wave.Draw(values)

	// If a resize option was set, perform it now
	if sizeX > 0 {
		// Perform image resize
		img = resize.Resize(uint(sizeX), uint(sizeY), img, resize.NearestNeighbor)
	}

	// Encode as PNG to HTTP writer
	if err := png.Encode(w, img); err != nil {
		log.Println(err)
	}
}

// hexToRGB converts a hex string to a RGB triple.
// Credit: https://code.google.com/p/gorilla/source/browse/color/hex.go?r=ef489f63418265a7249b1d53bdc358b09a4a2ea0
func hexToRGB(h string) (uint8, uint8, uint8) {
	if len(h) > 0 && h[0] == '#' {
		h = h[1:]
	}
	if len(h) == 3 {
		h = h[:1] + h[:1] + h[1:2] + h[1:2] + h[2:] + h[2:]
	}
	if len(h) == 6 {
		if rgb, err := strconv.ParseUint(string(h), 16, 32); err == nil {
			return uint8(rgb >> 16), uint8((rgb >> 8) & 0xFF), uint8(rgb & 0xFF)
		}
	}
	return 0, 0, 0
}
<|endoftext|>


<commit_before>package api

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"time"

	"github.com/FederationOfFathers/dashboard/bot"
)

func init() {
	Router.Path("/api/v0/gh/ui-rebuild").Methods("POST").HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			var payload struct {
				Ref string `json:"ref"`
			}
			if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
				logger.Warn(err.Error())
			}
			r.Body.Close()
			logger.Warn(fmt.Sprintf("%#v", payload))
			go func() {
				if payload.Ref == "refs/heads/dev" {
					bot.SendMessage("#-fof-dashboard", "rebuilding dev ui")
					now := time.Now()
					_, err := exec.Command(
						fmt.Sprintf(
							"%s/services/dashboard-dev/rebuild-ui",
							os.Getenv("HOME"),
						),
					).CombinedOutput()
					if err != nil {
						bot.SendMessage(
							"#-fof-dashboard",
							fmt.Sprintf("error rebuilding dev ui: ```%s```", err.Error()),
						)
					} else {
						bot.SendMessage(
							"#-fof-dashboard",
							fmt.Sprintf("dev ui rebuilt in %s", time.Now().Sub(now).String()),
						)
						time.Sleep(1)
					}
					err = exec.Command("/usr/bin/sv", "restart", fmt.Sprintf("%s/services/dashboard-dev/", os.Getenv("HOME"))).Run()
					if err != nil {
						logger.Warn(err.Error())
					}
				} else if payload.Ref == "refs/heads/master" {
					bot.SendMessage("#-fof-dashboard", "rebuilding production ui")
					now := time.Now()
					_, err := exec.Command(
						fmt.Sprintf(
							"%s/services/dashboard/update",
							os.Getenv("HOME"),
						),
					).CombinedOutput()
					if err != nil {
						bot.SendMessage(
							"#-fof-dashboard",
							fmt.Sprintf("error rebuilding production ui: ```%s```", err.Error()),
						)
					} else {
						bot.SendMessage(
							"#-fof-dashboard",
							fmt.Sprintf("production ui rebuilt in %s", time.Now().Sub(now).String()),
						)
					}
					err = exec.Command("/usr/bin/sv", "restart", fmt.Sprintf("%s/services/dashboard/", os.Getenv("HOME"))).Run()
					if err != nil {
						logger.Warn(err.Error())
					}
				}
			}()
		},
	)
}
<commit_msg>handle webhook for login link sending...<commit_after>package api

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"time"

	"github.com/FederationOfFathers/dashboard/bot"
)

func init() {
	Router.Path("/api/v0/slack/login").Methods("POST").HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		r.ParseForm()
		user := r.Form.Get("user_id")
		if home := os.Getenv("SERVICE_DIR"); home == "" {
			bot.SendLogin(user)
		} else {
			bot.SendDevLogin(user)
		}
	})
	Router.Path("/api/v0/gh/ui-rebuild").Methods("POST").HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			var payload struct {
				Ref string `json:"ref"`
			}
			if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
				logger.Warn(err.Error())
			}
			r.Body.Close()
			logger.Warn(fmt.Sprintf("%#v", payload))
			go func() {
				if payload.Ref == "refs/heads/dev" {
					bot.SendMessage("#-fof-dashboard", "rebuilding dev ui")
					now := time.Now()
					_, err := exec.Command(
						fmt.Sprintf(
							"%s/services/dashboard-dev/rebuild-ui",
							os.Getenv("HOME"),
						),
					).CombinedOutput()
					if err != nil {
						bot.SendMessage(
							"#-fof-dashboard",
							fmt.Sprintf("error rebuilding dev ui: ```%s```", err.Error()),
						)
					} else {
						bot.SendMessage(
							"#-fof-dashboard",
							fmt.Sprintf("dev ui rebuilt in %s", time.Now().Sub(now).String()),
						)
						time.Sleep(1)
					}
					err = exec.Command("/usr/bin/sv", "restart", fmt.Sprintf("%s/services/dashboard-dev/", os.Getenv("HOME"))).Run()
					if err != nil {
						logger.Warn(err.Error())
					}
				} else if payload.Ref == "refs/heads/master" {
					bot.SendMessage("#-fof-dashboard", "rebuilding production ui")
					now := time.Now()
					_, err := exec.Command(
						fmt.Sprintf(
							"%s/services/dashboard/update",
							os.Getenv("HOME"),
						),
					).CombinedOutput()
					if err != nil {
						bot.SendMessage(
							"#-fof-dashboard",
							fmt.Sprintf("error rebuilding production ui: ```%s```", err.Error()),
						)
					} else {
						bot.SendMessage(
							"#-fof-dashboard",
							fmt.Sprintf("production ui rebuilt in %s", time.Now().Sub(now).String()),
						)
					}
					err = exec.Command("/usr/bin/sv", "restart", fmt.Sprintf("%s/services/dashboard/", os.Getenv("HOME"))).Run()
					if err != nil {
						logger.Warn(err.Error())
					}
				}
			}()
		},
	)
}
<|endoftext|>


<commit_before>package bot

import (
	"fmt"
	"path/filepath"
	"sync"
	"time"
)

var pluginsRunning struct {
	count int
	sync.WaitGroup
	paused       bool
	shuttingDown bool
	sync.Mutex
}

// var pluginsRunning.count int
// var pluginsRunning.shuttingDown = false
// var paused = false // For Windows service pause support

// the shutdownMutex protects both the pluginsRunning.count and the pluginsRunning.shuttingDown
// flag
// var shutdownMutex sync.Mutex
// var plugRunningWaitGroup sync.WaitGroup

const keepListeningDuration = 77 * time.Second

// messageAppliesToPlugin checks the user and channel against the plugin's
// configuration to determine if the message should be evaluated. Used by
// both handleMessage and the help builtin.
func messageAppliesToPlugin(user, channel string, plugin *Plugin) bool {
	directMsg := false
	if len(channel) == 0 {
		directMsg = true
	}
	if !directMsg && plugin.DirectOnly {
		return false
	}
	if plugin.RequireAdmin {
		isAdmin := false
		robot.RLock()
		for _, adminUser := range robot.adminUsers {
			if user == adminUser {
				isAdmin = true
				break
			}
		}
		robot.RUnlock()
		if !isAdmin {
			return false
		}
	}
	if len(plugin.Users) > 0 {
		userOk := false
		for _, allowedUser := range plugin.Users {
			match, err := filepath.Match(allowedUser, user)
			if match && err == nil {
				userOk = true
			}
		}
		if !userOk {
			return false
		}
	}
	if directMsg && (plugin.AllowDirect || plugin.DirectOnly) {
		return true
	}
	if len(plugin.Channels) > 0 {
		for _, pchannel := range plugin.Channels {
			if pchannel == channel {
				return true
			}
		}
	} else {
		if plugin.AllChannels {
			return true
		}
	}
	return false
}

// checkPluginMatchers checks either command matchers (for messages directed at
// the robot), or message matchers (for ambient commands that need not be
// directed at the robot), and calls the plugin if it matches. Note: this
// function is called under a read lock on the 'b' struct.
func checkPluginMatchers(checkCommands bool, bot *Robot, messagetext string) (commandMatched bool) {
	// un-needed, but more clear
	commandMatched = false
	pluginlist.RLock()
	plugins := pluginlist.p
	pluginlist.RUnlock()
	for _, plugin := range plugins {
		Log(Trace, fmt.Sprintf("Checking message \"%s\" against plugin %s, active in %d channels (allchannels: %t)", messagetext, plugin.name, len(plugin.Channels), plugin.AllChannels))
		ok := messageAppliesToPlugin(bot.User, bot.Channel, plugin)
		if !ok {
			Log(Trace, fmt.Sprintf("Plugin %s ignoring message in channel %s, doesn't meet criteria", plugin.name, bot.Channel))
			continue
		}
		var matchers []InputMatcher
		if checkCommands {
			matchers = plugin.CommandMatchers
		} else {
			matchers = plugin.MessageMatchers
		}
		for _, matcher := range matchers {
			Log(Trace, fmt.Sprintf("Checking \"%s\" against \"%s\"", messagetext, matcher.Regex))
			matches := matcher.re.FindAllStringSubmatch(messagetext, -1)
			var matched bool
			var cmdArgs []string
			if matches != nil {
				matched = true
				cmdArgs = matches[0][1:]
				if len(matcher.Contexts) > 0 {
					// Resolve & store "it" with short-term memories
					ts := time.Now()
					shortTermMemories.Lock()
					for i, contextLabel := range matcher.Contexts {
						if contextLabel != "" {
							key := "context:" + contextLabel
							c := memoryContext{key, bot.User, bot.Channel}
							if len(cmdArgs) > i && (cmdArgs[i] == "it" || cmdArgs[i] == "") {
								s, ok := shortTermMemories.m[c]
								if ok {
									cmdArgs[i] = s.memory
									// TODO: it would probably be best to substitute the value
									// from "it" back in to the original message and re-check for
									// a match. Failing a match, matched should be set to false.
									s.timestamp = ts
									shortTermMemories.m[c] = s
								} else {
									bot.Say(fmt.Sprintf("Sorry, I don't remember which %s we were talking about", contextLabel))
									shortTermMemories.Unlock()
									return true
								}
							} else {
								s := shortTermMemory{cmdArgs[i], ts}
								shortTermMemories.m[c] = s
							}
						}
					}
					shortTermMemories.Unlock()
				}
			}
			if matched {
				commandMatched = true
				privilegesOk := true
				if len(plugin.ElevatedCommands) > 0 {
					for _, i := range plugin.ElevatedCommands {
						if matcher.Command == i {
							if robot.elevator != nil {
								// elevators have their own pluginID & name, for brain access
								pbot := &Robot{
									User:    bot.User,
									Channel: bot.Channel,
									Format:  Variable,
									// NOTE: checkPluginMatchers is called under b.lock.RLock()
									pluginID: "elevator-" + robot.elevatorProvider,
								}
								privilegesOk = robot.elevator(pbot, false)
							} else {
								privilegesOk = false
								Log(Error, "Encountered elevated command and no elevation method configured")
							}
						}
					}
				}
				if len(plugin.ElevateImmediateCommands) > 0 {
					for _, i := range plugin.ElevateImmediateCommands {
						if matcher.Command == i {
							if robot.elevator != nil {
								// elevators have their own pluginID & name, for brain access
								pbot := &Robot{
									User:    bot.User,
									Channel: bot.Channel,
									Format:  Variable,
									// NOTE: checkPluginMatchers is called under b.lock.RLock()
									pluginID: "elevator-" + robot.elevatorProvider,
								}
								privilegesOk = robot.elevator(pbot, true)
							} else {
								privilegesOk = false
								Log(Error, "Encountered elevated command and no elevation method configured")
							}
						}
					}
				}
				if privilegesOk {
					abort := false
					if plugin.name == "builtInadmin" && matcher.Command == "abort" {
						abort = true
					}
					pluginsRunning.Lock()
					if pluginsRunning.shuttingDown && !abort {
						bot.Say("Sorry, I'm shutting down and can't start any new tasks")
						pluginsRunning.Unlock()
					} else if pluginsRunning.paused && !abort {
						bot.Say("Sorry, I've been paused and can't start any new tasks")
						pluginsRunning.Unlock()
					} else {
						pluginsRunning.Unlock()
						pluginsRunning.Add(1)
						go callPlugin(bot, plugin, matcher.Command, cmdArgs...)
					}
				} else {
					Log(Error, fmt.Sprintf("Elevation failed for command \"%s\", plugin %s", matcher.Command, plugin.name))
					bot.Say(fmt.Sprintf("Sorry, the \"%s\" command requires elevated privileges", matcher.Command))
				}
			}
		}
	}
	return commandMatched
}

// handleMessage checks the message against plugin commands and full-message matches,
// then dispatches it to all applicable handlers in a separate go routine. If the robot
// was addressed directly but nothing matched, any registered CatchAll plugins are called.
// There Should Be Only One (catchall, in theory (?))
func handleMessage(isCommand bool, channel, user, messagetext string) {
	bot := &Robot{
		User:    user,
		Channel: channel,
		Format:  Variable,
	}
	defer checkPanic(bot, messagetext)
	pluginlist.RLock()
	plugins := pluginlist.p
	pluginlist.RUnlock()
	if len(channel) == 0 {
		Log(Trace, fmt.Sprintf("Bot received a direct message from %s: %s", user, messagetext))
	}
	commandMatched := false
	waitingForReply := false
	var catchAllPlugins []*Plugin
	ts := time.Now()
	lastCmdContext := memoryContext{"lastCmd", user, channel}
	lastMsgContext := memoryContext{"lastMsg", user, channel}
	var last shortTermMemory
	var ok bool
	// See if the robot got a blank message, indicating that the last message
	// was meant for it (if it was in the keepListeningDuration)
	if isCommand && messagetext == "" {
		shortTermMemories.Lock()
		last, ok = shortTermMemories.m[lastMsgContext]
		shortTermMemories.Unlock()
		if ok && ts.Sub(last.timestamp) < keepListeningDuration {
			messagetext = last.memory
			commandMatched = checkPluginMatchers(true, bot, messagetext)
		} else {
			commandMatched = true
			bot.Say("Yes?")
		}
	}
	shortTermMemories.Lock()
	last, ok = shortTermMemories.m[lastCmdContext]
	shortTermMemories.Unlock()
	if ok && !commandMatched {
		// If the robot has been spoken to recently, it will keep listening
		// for commands for a short duration
		if ts.Sub(last.timestamp) < keepListeningDuration {
			commandMatched = checkPluginMatchers(true, bot, messagetext)
		}
	}
	if !commandMatched && isCommand {
		// Even if the command doesn't match, remember the robot was spoken to
		last = shortTermMemory{messagetext, ts}
		shortTermMemories.Lock()
		shortTermMemories.m[lastCmdContext] = last
		shortTermMemories.Unlock()
		catchAllPlugins = make([]*Plugin, 0, len(plugins))
		for _, plugin := range plugins {
			if plugin.CatchAll {
				catchAllPlugins = append(catchAllPlugins, plugin)
			}
		}
		// See if a command matches (and runs)
		commandMatched = checkPluginMatchers(true, bot, messagetext)
	}
	// See if the robot was waiting on a reply
	matcher := replyMatcher{user, channel}
	Log(Trace, fmt.Sprintf("Checking replies for matcher: %q", matcher))
	replies.Lock()
	rep, waitingForReply := replies.m[matcher]
	if !waitingForReply {
		replies.Unlock()
		Log(Trace, "No matching replyWaiter")
	} else {
		delete(replies.m, matcher)
		replies.Unlock()
		if commandMatched {
			rep.replyChannel <- reply{false, true, ""}
			Log(Debug, fmt.Sprintf("User \"%s\" issued a new command while the robot was waiting for a reply in channel \"%s\"", user, channel))
		} else {
			// if the robot was waiting on a reply, we don't want to check for
			// ambient message matches - the plugin will handle it.
			commandMatched = true
			matched := false
			if rep.re.MatchString(messagetext) {
				matched = true
			}
			Log(Debug, fmt.Sprintf("Found replyWaiter for user \"%s\" in channel \"%s\", checking if message \"%s\" matches \"%s\": %t", user, channel, messagetext, rep.re.String(), matched))
			rep.replyChannel <- reply{matched, false, messagetext}
		}
	}
	// Direct commands were checked above; if a direct command didn't match,
	// and a there wasn't a reply being waited on, then we check ambient
	// MessageMatchers if it wasn't a direct command. Note that ambient
	// commands never match in a DM.
	if !commandMatched && !waitingForReply && !isCommand {
		// check for ambient message matches
		commandMatched = checkPluginMatchers(false, bot, messagetext)
	}
	if isCommand && !commandMatched { // the robot was spoken too, but nothing matched - call catchAlls
		pluginsRunning.Lock()
		if !pluginsRunning.shuttingDown {
			pluginsRunning.Unlock()
			Log(Debug, fmt.Sprintf("Unmatched command sent to robot, calling catchalls: %s", messagetext))
			for _, plugin := range catchAllPlugins {
				pluginsRunning.Add(1)
				go callPlugin(bot, plugin, "catchall", messagetext)
			}
		} else {
			// If the robot is shutting down, just ignore catch-all plugins
			pluginsRunning.Unlock()
		}
	}
	last = shortTermMemory{messagetext, ts}
	if commandMatched || isCommand {
		shortTermMemories.Lock()
		shortTermMemories.m[lastCmdContext] = last
		delete(shortTermMemories.m, lastMsgContext)
		shortTermMemories.Unlock()
	} else {
		shortTermMemories.Lock()
		shortTermMemories.m[lastMsgContext] = last
		shortTermMemories.Unlock()
	}
}
<commit_msg>Fix from gometalinter<commit_after>package bot

import (
	"fmt"
	"path/filepath"
	"sync"
	"time"
)

var pluginsRunning struct {
	count int
	sync.WaitGroup
	paused       bool
	shuttingDown bool
	sync.Mutex
}

// var pluginsRunning.count int
// var pluginsRunning.shuttingDown = false
// var paused = false // For Windows service pause support

// the shutdownMutex protects both the pluginsRunning.count and the pluginsRunning.shuttingDown
// flag
// var shutdownMutex sync.Mutex
// var plugRunningWaitGroup sync.WaitGroup

const keepListeningDuration = 77 * time.Second

// messageAppliesToPlugin checks the user and channel against the plugin's
// configuration to determine if the message should be evaluated. Used by
// both handleMessage and the help builtin.
func messageAppliesToPlugin(user, channel string, plugin *Plugin) bool {
	directMsg := false
	if len(channel) == 0 {
		directMsg = true
	}
	if !directMsg && plugin.DirectOnly {
		return false
	}
	if plugin.RequireAdmin {
		isAdmin := false
		robot.RLock()
		for _, adminUser := range robot.adminUsers {
			if user == adminUser {
				isAdmin = true
				break
			}
		}
		robot.RUnlock()
		if !isAdmin {
			return false
		}
	}
	if len(plugin.Users) > 0 {
		userOk := false
		for _, allowedUser := range plugin.Users {
			match, err := filepath.Match(allowedUser, user)
			if match && err == nil {
				userOk = true
			}
		}
		if !userOk {
			return false
		}
	}
	if directMsg && (plugin.AllowDirect || plugin.DirectOnly) {
		return true
	}
	if len(plugin.Channels) > 0 {
		for _, pchannel := range plugin.Channels {
			if pchannel == channel {
				return true
			}
		}
	} else {
		if plugin.AllChannels {
			return true
		}
	}
	return false
}

// checkPluginMatchers checks either command matchers (for messages directed at
// the robot), or message matchers (for ambient commands that need not be
// directed at the robot), and calls the plugin if it matches. Note: this
// function is called under a read lock on the 'b' struct.
func checkPluginMatchers(checkCommands bool, bot *Robot, messagetext string) (commandMatched bool) {
	// un-needed, but more clear
	commandMatched = false
	pluginlist.RLock()
	plugins := pluginlist.p
	pluginlist.RUnlock()
	for _, plugin := range plugins {
		Log(Trace, fmt.Sprintf("Checking message \"%s\" against plugin %s, active in %d channels (allchannels: %t)", messagetext, plugin.name, len(plugin.Channels), plugin.AllChannels))
		ok := messageAppliesToPlugin(bot.User, bot.Channel, plugin)
		if !ok {
			Log(Trace, fmt.Sprintf("Plugin %s ignoring message in channel %s, doesn't meet criteria", plugin.name, bot.Channel))
			continue
		}
		var matchers []InputMatcher
		if checkCommands {
			matchers = plugin.CommandMatchers
		} else {
			matchers = plugin.MessageMatchers
		}
		for _, matcher := range matchers {
			Log(Trace, fmt.Sprintf("Checking \"%s\" against \"%s\"", messagetext, matcher.Regex))
			matches := matcher.re.FindAllStringSubmatch(messagetext, -1)
			var matched bool
			var cmdArgs []string
			if matches != nil {
				matched = true
				cmdArgs = matches[0][1:]
				if len(matcher.Contexts) > 0 {
					// Resolve & store "it" with short-term memories
					ts := time.Now()
					shortTermMemories.Lock()
					for i, contextLabel := range matcher.Contexts {
						if contextLabel != "" {
							key := "context:" + contextLabel
							c := memoryContext{key, bot.User, bot.Channel}
							if len(cmdArgs) > i && (cmdArgs[i] == "it" || cmdArgs[i] == "") {
								s, ok := shortTermMemories.m[c]
								if ok {
									cmdArgs[i] = s.memory
									// TODO: it would probably be best to substitute the value
									// from "it" back in to the original message and re-check for
									// a match. Failing a match, matched should be set to false.
									s.timestamp = ts
									shortTermMemories.m[c] = s
								} else {
									bot.Say(fmt.Sprintf("Sorry, I don't remember which %s we were talking about", contextLabel))
									shortTermMemories.Unlock()
									return true
								}
							} else {
								s := shortTermMemory{cmdArgs[i], ts}
								shortTermMemories.m[c] = s
							}
						}
					}
					shortTermMemories.Unlock()
				}
			}
			if matched {
				commandMatched = true
				privilegesOk := true
				if len(plugin.ElevatedCommands) > 0 {
					for _, i := range plugin.ElevatedCommands {
						if matcher.Command == i {
							if robot.elevator != nil {
								// elevators have their own pluginID & name, for brain access
								pbot := &Robot{
									User:    bot.User,
									Channel: bot.Channel,
									Format:  Variable,
									// NOTE: checkPluginMatchers is called under b.lock.RLock()
									pluginID: "elevator-" + robot.elevatorProvider,
								}
								privilegesOk = robot.elevator(pbot, false)
							} else {
								privilegesOk = false
								Log(Error, "Encountered elevated command and no elevation method configured")
							}
						}
					}
				}
				if len(plugin.ElevateImmediateCommands) > 0 {
					for _, i := range plugin.ElevateImmediateCommands {
						if matcher.Command == i {
							if robot.elevator != nil {
								// elevators have their own pluginID & name, for brain access
								pbot := &Robot{
									User:    bot.User,
									Channel: bot.Channel,
									Format:  Variable,
									// NOTE: checkPluginMatchers is called under b.lock.RLock()
									pluginID: "elevator-" + robot.elevatorProvider,
								}
								privilegesOk = robot.elevator(pbot, true)
							} else {
								privilegesOk = false
								Log(Error, "Encountered elevated command and no elevation method configured")
							}
						}
					}
				}
				if privilegesOk {
					abort := false
					if plugin.name == "builtInadmin" && matcher.Command == "abort" {
						abort = true
					}
					pluginsRunning.Lock()
					if pluginsRunning.shuttingDown && !abort {
						bot.Say("Sorry, I'm shutting down and can't start any new tasks")
						pluginsRunning.Unlock()
					} else if pluginsRunning.paused && !abort {
						bot.Say("Sorry, I've been paused and can't start any new tasks")
						pluginsRunning.Unlock()
					} else {
						pluginsRunning.Unlock()
						pluginsRunning.Add(1)
						go callPlugin(bot, plugin, matcher.Command, cmdArgs...)
					}
				} else {
					Log(Error, fmt.Sprintf("Elevation failed for command \"%s\", plugin %s", matcher.Command, plugin.name))
					bot.Say(fmt.Sprintf("Sorry, the \"%s\" command requires elevated privileges", matcher.Command))
				}
			}
		}
	}
	return commandMatched
}

// handleMessage checks the message against plugin commands and full-message matches,
// then dispatches it to all applicable handlers in a separate go routine. If the robot
// was addressed directly but nothing matched, any registered CatchAll plugins are called.
// There Should Be Only One (catchall, in theory (?))
func handleMessage(isCommand bool, channel, user, messagetext string) {
	bot := &Robot{
		User:    user,
		Channel: channel,
		Format:  Variable,
	}
	defer checkPanic(bot, messagetext)
	pluginlist.RLock()
	plugins := pluginlist.p
	pluginlist.RUnlock()
	if len(channel) == 0 {
		Log(Trace, fmt.Sprintf("Bot received a direct message from %s: %s", user, messagetext))
	}
	commandMatched := false
	var catchAllPlugins []*Plugin
	ts := time.Now()
	lastCmdContext := memoryContext{"lastCmd", user, channel}
	lastMsgContext := memoryContext{"lastMsg", user, channel}
	var last shortTermMemory
	var ok bool
	// See if the robot got a blank message, indicating that the last message
	// was meant for it (if it was in the keepListeningDuration)
	if isCommand && messagetext == "" {
		shortTermMemories.Lock()
		last, ok = shortTermMemories.m[lastMsgContext]
		shortTermMemories.Unlock()
		if ok && ts.Sub(last.timestamp) < keepListeningDuration {
			messagetext = last.memory
			commandMatched = checkPluginMatchers(true, bot, messagetext)
		} else {
			commandMatched = true
			bot.Say("Yes?")
		}
	}
	shortTermMemories.Lock()
	last, ok = shortTermMemories.m[lastCmdContext]
	shortTermMemories.Unlock()
	if ok && !commandMatched {
		// If the robot has been spoken to recently, it will keep listening
		// for commands for a short duration
		if ts.Sub(last.timestamp) < keepListeningDuration {
			commandMatched = checkPluginMatchers(true, bot, messagetext)
		}
	}
	if !commandMatched && isCommand {
		// Even if the command doesn't match, remember the robot was spoken to
		last = shortTermMemory{messagetext, ts}
		shortTermMemories.Lock()
		shortTermMemories.m[lastCmdContext] = last
		shortTermMemories.Unlock()
		catchAllPlugins = make([]*Plugin, 0, len(plugins))
		for _, plugin := range plugins {
			if plugin.CatchAll {
				catchAllPlugins = append(catchAllPlugins, plugin)
			}
		}
		// See if a command matches (and runs)
		commandMatched = checkPluginMatchers(true, bot, messagetext)
	}
	// See if the robot was waiting on a reply
	matcher := replyMatcher{user, channel}
	Log(Trace, fmt.Sprintf("Checking replies for matcher: %q", matcher))
	replies.Lock()
	rep, waitingForReply := replies.m[matcher]
	if !waitingForReply {
		replies.Unlock()
		Log(Trace, "No matching replyWaiter")
	} else {
		delete(replies.m, matcher)
		replies.Unlock()
		if commandMatched {
			rep.replyChannel <- reply{false, true, ""}
			Log(Debug, fmt.Sprintf("User \"%s\" issued a new command while the robot was waiting for a reply in channel \"%s\"", user, channel))
		} else {
			// if the robot was waiting on a reply, we don't want to check for
			// ambient message matches - the plugin will handle it.
			commandMatched = true
			matched := false
			if rep.re.MatchString(messagetext) {
				matched = true
			}
			Log(Debug, fmt.Sprintf("Found replyWaiter for user \"%s\" in channel \"%s\", checking if message \"%s\" matches \"%s\": %t", user, channel, messagetext, rep.re.String(), matched))
			rep.replyChannel <- reply{matched, false, messagetext}
		}
	}
	// Direct commands were checked above; if a direct command didn't match,
	// and a there wasn't a reply being waited on, then we check ambient
	// MessageMatchers if it wasn't a direct command. Note that ambient
	// commands never match in a DM.
	if !commandMatched && !waitingForReply && !isCommand {
		// check for ambient message matches
		commandMatched = checkPluginMatchers(false, bot, messagetext)
	}
	if isCommand && !commandMatched { // the robot was spoken too, but nothing matched - call catchAlls
		pluginsRunning.Lock()
		if !pluginsRunning.shuttingDown {
			pluginsRunning.Unlock()
			Log(Debug, fmt.Sprintf("Unmatched command sent to robot, calling catchalls: %s", messagetext))
			for _, plugin := range catchAllPlugins {
				pluginsRunning.Add(1)
				go callPlugin(bot, plugin, "catchall", messagetext)
			}
		} else {
			// If the robot is shutting down, just ignore catch-all plugins
			pluginsRunning.Unlock()
		}
	}
	last = shortTermMemory{messagetext, ts}
	if commandMatched || isCommand {
		shortTermMemories.Lock()
		shortTermMemories.m[lastCmdContext] = last
		delete(shortTermMemories.m, lastMsgContext)
		shortTermMemories.Unlock()
	} else {
		shortTermMemories.Lock()
		shortTermMemories.m[lastMsgContext] = last
		shortTermMemories.Unlock()
	}
}
<|endoftext|>


<commit_before>package cloudflare

import (
	"context"
	"encoding/json"
	"fmt"
	"github.com/pkg/errors"
	"net/http"
	"time"
)

type TeamsProxyEndpointListResponse struct {
	Response
	ResultInfo `json:"result_info"`
	Result     []TeamsProxyEndpoint `json:"result"`
}
type TeamsProxyEndpointDetailResponse struct {
	Response
	Result TeamsProxyEndpoint `json:"result"`
}

type TeamsProxyEndpoint struct {
	ID        string     `json:"id"`
	Name      string     `json:"name"`
	IPs       []string   `json:"ips"`
	Subdomain string     `json:"subdomain"`
	CreatedAt *time.Time `json:"created_at,omitempty"`
	UpdatedAt *time.Time `json:"updated_at,omitempty"`
}

// TeamsProxyEndpoint returns a single proxy endpoints within an account.
//
// API reference: https://api.cloudflare.com/#teams-proxy-endpoints-proxy-endpoint-details
func (api *API) TeamsProxyEndpoint(ctx context.Context, accountID, proxyEndpointID string) (TeamsProxyEndpoint, error) {
	uri := fmt.Sprintf("/%s/%s/gateway/proxy_endpoints/%s", AccountRouteRoot, accountID, proxyEndpointID)

	res, err := api.makeRequestContext(ctx, http.MethodGet, uri, nil)
	if err != nil {
		return TeamsProxyEndpoint{}, err
	}

	var teamsProxyEndpointDetailResponse TeamsProxyEndpointDetailResponse
	err = json.Unmarshal(res, &teamsProxyEndpointDetailResponse)
	if err != nil {
		return TeamsProxyEndpoint{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsProxyEndpointDetailResponse.Result, nil
}

// TeamsProxyEndpoints returns all proxy endpoints within an account.
//
// API reference: https://api.cloudflare.com/#teams-proxy-endpoints-list-proxy-endpoints
func (api *API) TeamsProxyEndpoints(ctx context.Context, accountID string) ([]TeamsProxyEndpoint, ResultInfo, error) {
	uri := fmt.Sprintf("/%s/%s/gateway/proxy_endpoints", AccountRouteRoot, accountID)

	res, err := api.makeRequestContext(ctx, http.MethodGet, uri, nil)
	if err != nil {
		return []TeamsProxyEndpoint{}, ResultInfo{}, err
	}

	var teamsProxyEndpointListResponse TeamsProxyEndpointListResponse
	err = json.Unmarshal(res, &teamsProxyEndpointListResponse)
	if err != nil {
		return []TeamsProxyEndpoint{}, ResultInfo{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsProxyEndpointListResponse.Result, teamsProxyEndpointListResponse.ResultInfo, nil
}

// CreateTeamsProxyEndpoint creates a new proxy endpoint.
//
// API reference: https://api.cloudflare.com/#teams-proxy-endpoints-create-proxy-endpoint
func (api *API) CreateTeamsProxyEndpoint(ctx context.Context, accountID string, proxyEndpoint TeamsProxyEndpoint) (TeamsProxyEndpoint, error) {
	uri := fmt.Sprintf("/%s/%s/gateway/proxy_endpoints", AccountRouteRoot, accountID)

	res, err := api.makeRequestContext(ctx, http.MethodPost, uri, proxyEndpoint)
	if err != nil {
		return TeamsProxyEndpoint{}, err
	}

	var teamsProxyEndpointDetailResponse TeamsProxyEndpointDetailResponse
	err = json.Unmarshal(res, &teamsProxyEndpointDetailResponse)
	if err != nil {
		return TeamsProxyEndpoint{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsProxyEndpointDetailResponse.Result, nil
}

// UpdateTeamsProxyEndpoint updates an existing teams Proxy Endpoint.
//
// API reference: https://api.cloudflare.com/#teams-proxy-endpoints-update-proxy-endpoint
func (api *API) UpdateTeamsProxyEndpoint(ctx context.Context, accountID string, proxyEndpoint TeamsProxyEndpoint) (TeamsProxyEndpoint, error) {
	if proxyEndpoint.ID == "" {
		return TeamsProxyEndpoint{}, errors.Errorf("Proxy Endpoint ID cannot be empty")
	}

	uri := fmt.Sprintf(
		"/%s/%s/gateway/proxy_endpoints/%s",
		AccountRouteRoot,
		accountID,
		proxyEndpoint.ID,
	)

	res, err := api.makeRequestContext(ctx, http.MethodPatch, uri, proxyEndpoint)
	if err != nil {
		return TeamsProxyEndpoint{}, err
	}

	var teamsProxyEndpointDetailResponse TeamsProxyEndpointDetailResponse
	err = json.Unmarshal(res, &teamsProxyEndpointDetailResponse)
	if err != nil {
		return TeamsProxyEndpoint{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsProxyEndpointDetailResponse.Result, nil
}

// DeleteTeamsProxyEndpoint deletes a teams Proxy Endpoint.
//
// API reference: https://api.cloudflare.com/#teams-proxy-endpoints-delete-teams-proxy-endpoint
func (api *API) DeleteTeamsProxyEndpoint(ctx context.Context, accountID, proxyEndpointID string) error {
	uri := fmt.Sprintf(
		"/%s/%s/gateway/proxy_endpoints/%s",
		AccountRouteRoot,
		accountID,
		proxyEndpointID,
	)

	_, err := api.makeRequestContext(ctx, http.MethodDelete, uri, nil)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>update API documentation URLs<commit_after>package cloudflare

import (
	"context"
	"encoding/json"
	"fmt"
	"github.com/pkg/errors"
	"net/http"
	"time"
)

type TeamsProxyEndpointListResponse struct {
	Response
	ResultInfo `json:"result_info"`
	Result     []TeamsProxyEndpoint `json:"result"`
}
type TeamsProxyEndpointDetailResponse struct {
	Response
	Result TeamsProxyEndpoint `json:"result"`
}

type TeamsProxyEndpoint struct {
	ID        string     `json:"id"`
	Name      string     `json:"name"`
	IPs       []string   `json:"ips"`
	Subdomain string     `json:"subdomain"`
	CreatedAt *time.Time `json:"created_at,omitempty"`
	UpdatedAt *time.Time `json:"updated_at,omitempty"`
}

// TeamsProxyEndpoint returns a single proxy endpoints within an account.
//
// API reference: https://api.cloudflare.com/#zero-trust-gateway-proxy-endpoints-proxy-endpoint-details
func (api *API) TeamsProxyEndpoint(ctx context.Context, accountID, proxyEndpointID string) (TeamsProxyEndpoint, error) {
	uri := fmt.Sprintf("/%s/%s/gateway/proxy_endpoints/%s", AccountRouteRoot, accountID, proxyEndpointID)

	res, err := api.makeRequestContext(ctx, http.MethodGet, uri, nil)
	if err != nil {
		return TeamsProxyEndpoint{}, err
	}

	var teamsProxyEndpointDetailResponse TeamsProxyEndpointDetailResponse
	err = json.Unmarshal(res, &teamsProxyEndpointDetailResponse)
	if err != nil {
		return TeamsProxyEndpoint{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsProxyEndpointDetailResponse.Result, nil
}

// TeamsProxyEndpoints returns all proxy endpoints within an account.
//
// API reference: https://api.cloudflare.com/#zero-trust-gateway-proxy-endpoints-list-proxy-endpoints
func (api *API) TeamsProxyEndpoints(ctx context.Context, accountID string) ([]TeamsProxyEndpoint, ResultInfo, error) {
	uri := fmt.Sprintf("/%s/%s/gateway/proxy_endpoints", AccountRouteRoot, accountID)

	res, err := api.makeRequestContext(ctx, http.MethodGet, uri, nil)
	if err != nil {
		return []TeamsProxyEndpoint{}, ResultInfo{}, err
	}

	var teamsProxyEndpointListResponse TeamsProxyEndpointListResponse
	err = json.Unmarshal(res, &teamsProxyEndpointListResponse)
	if err != nil {
		return []TeamsProxyEndpoint{}, ResultInfo{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsProxyEndpointListResponse.Result, teamsProxyEndpointListResponse.ResultInfo, nil
}

// CreateTeamsProxyEndpoint creates a new proxy endpoint.
//
// API reference: https://api.cloudflare.com/#zero-trust-gateway-proxy-endpoints-create-proxy-endpoint
func (api *API) CreateTeamsProxyEndpoint(ctx context.Context, accountID string, proxyEndpoint TeamsProxyEndpoint) (TeamsProxyEndpoint, error) {
	uri := fmt.Sprintf("/%s/%s/gateway/proxy_endpoints", AccountRouteRoot, accountID)

	res, err := api.makeRequestContext(ctx, http.MethodPost, uri, proxyEndpoint)
	if err != nil {
		return TeamsProxyEndpoint{}, err
	}

	var teamsProxyEndpointDetailResponse TeamsProxyEndpointDetailResponse
	err = json.Unmarshal(res, &teamsProxyEndpointDetailResponse)
	if err != nil {
		return TeamsProxyEndpoint{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsProxyEndpointDetailResponse.Result, nil
}

// UpdateTeamsProxyEndpoint updates an existing teams Proxy Endpoint.
//
// API reference: https://api.cloudflare.com/#zero-trust-gateway-proxy-endpoints-update-proxy-endpoint
func (api *API) UpdateTeamsProxyEndpoint(ctx context.Context, accountID string, proxyEndpoint TeamsProxyEndpoint) (TeamsProxyEndpoint, error) {
	if proxyEndpoint.ID == "" {
		return TeamsProxyEndpoint{}, errors.Errorf("Proxy Endpoint ID cannot be empty")
	}

	uri := fmt.Sprintf(
		"/%s/%s/gateway/proxy_endpoints/%s",
		AccountRouteRoot,
		accountID,
		proxyEndpoint.ID,
	)

	res, err := api.makeRequestContext(ctx, http.MethodPatch, uri, proxyEndpoint)
	if err != nil {
		return TeamsProxyEndpoint{}, err
	}

	var teamsProxyEndpointDetailResponse TeamsProxyEndpointDetailResponse
	err = json.Unmarshal(res, &teamsProxyEndpointDetailResponse)
	if err != nil {
		return TeamsProxyEndpoint{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsProxyEndpointDetailResponse.Result, nil
}

// DeleteTeamsProxyEndpoint deletes a teams Proxy Endpoint.
//
// API reference: https://api.cloudflare.com/#zero-trust-gateway-proxy-endpoints-delete-proxy-endpoint
func (api *API) DeleteTeamsProxyEndpoint(ctx context.Context, accountID, proxyEndpointID string) error {
	uri := fmt.Sprintf(
		"/%s/%s/gateway/proxy_endpoints/%s",
		AccountRouteRoot,
		accountID,
		proxyEndpointID,
	)

	_, err := api.makeRequestContext(ctx, http.MethodDelete, uri, nil)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Timon Wong
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package zapsyslog

import (
	"os"
	"path"
	"strings"
	"time"

	syslog "github.com/timonwong/go-syslog"
	"github.com/timonwong/zap-syslog/internal"
	"github.com/timonwong/zap-syslog/internal/bufferpool"
	"go.uber.org/zap"
	"go.uber.org/zap/buffer"
	"go.uber.org/zap/zapcore"
)

const (
	version         = 1
	severityMask    = 0x07
	facilityMask    = 0xf8
	nilValue        = "-"
	timestampFormat = "2006-01-02T15:04:05.000000Z07:00" // RFC3339 with micro fraction seconds
	maxHostnameLen  = 255
	maxAppNameLen   = 48
)

var (
	_ zapcore.Encoder = &syslogEncoder{}
	_ jsonEncoder     = zapcore.NewJSONEncoder(zap.NewProductionEncoderConfig()).(jsonEncoder)
)

// Framing configures RFC6587 TCP transport framing.
type Framing int

// Framing.
const (
	NonTransparentFraming Framing = iota
	OctetCountingFraming
	DefaultFraming = NonTransparentFraming
)

type jsonEncoder interface {
	zapcore.Encoder
	zapcore.ArrayEncoder
}

// SyslogEncoderConfig allows users to configure the concrete encoders for zap syslog.
type SyslogEncoderConfig struct {
	zapcore.EncoderConfig

	Framing  Framing         `json:"framing" yaml:"framing"`
	Facility syslog.Priority `json:"facility" yaml:"facility"`
	Hostname string          `json:"hostname" yaml:"hostname"`
	PID      int             `json:"pid" yaml:"pid"`
	App      string          `json:"app" yaml:"app"`
}

type syslogEncoder struct {
	*SyslogEncoderConfig
	je jsonEncoder
}

func rfc5424CompliantASCIIMapper(r rune) rune {
	// PRINTUSASCII    = %d33-126
	if r < 33 || r > 126 {
		return '_'
	}
	return r
}

func toRFC5424CompliantASCIIString(s string) string {
	return strings.Map(rfc5424CompliantASCIIMapper, s)
}

// NewSyslogEncoder creates a syslogEncoder that you should not use because I'm not done with it yet.
func NewSyslogEncoder(cfg SyslogEncoderConfig) zapcore.Encoder {
	if cfg.Hostname == "" {
		hostname, _ := os.Hostname()
		cfg.Hostname = hostname
	}
	if cfg.Hostname == "" {
		cfg.Hostname = nilValue
	} else {
		hostname := toRFC5424CompliantASCIIString(cfg.Hostname)
		if len(hostname) > maxHostnameLen {
			hostname = hostname[:maxHostnameLen]
		}
		cfg.Hostname = hostname
	}

	if cfg.PID == 0 {
		cfg.PID = os.Getpid()
	}
	if cfg.App == "" {
		cfg.App = nilValue
	} else {
		app := cfg.App
		if len(app) > maxAppNameLen {
			app = path.Base(app)
		}
		if len(app) > maxAppNameLen {
			app = app[:maxAppNameLen]
		}
		app = toRFC5424CompliantASCIIString(app)
	}

	cfg.EncoderConfig.LineEnding = "\n"
	je := zapcore.NewJSONEncoder(cfg.EncoderConfig).(jsonEncoder)
	return &syslogEncoder{
		SyslogEncoderConfig: &cfg,
		je:                  je,
	}
}

func (enc *syslogEncoder) AddArray(key string, arr zapcore.ArrayMarshaler) error {
	return enc.je.AddArray(key, arr)
}

func (enc *syslogEncoder) AddObject(key string, obj zapcore.ObjectMarshaler) error {
	return enc.je.AddObject(key, obj)
}

func (enc *syslogEncoder) AddBinary(key string, val []byte)          { enc.je.AddBinary(key, val) }
func (enc *syslogEncoder) AddByteString(key string, val []byte)      { enc.je.AddByteString(key, val) }
func (enc *syslogEncoder) AddBool(key string, val bool)              { enc.je.AddBool(key, val) }
func (enc *syslogEncoder) AddComplex128(key string, val complex128)  { enc.je.AddComplex128(key, val) }
func (enc *syslogEncoder) AddDuration(key string, val time.Duration) { enc.je.AddDuration(key, val) }
func (enc *syslogEncoder) AddFloat64(key string, val float64)        { enc.je.AddFloat64(key, val) }
func (enc *syslogEncoder) AddInt64(key string, val int64)            { enc.je.AddInt64(key, val) }

func (enc *syslogEncoder) AddReflected(key string, obj interface{}) error {
	return enc.je.AddReflected(key, obj)
}

func (enc *syslogEncoder) OpenNamespace(key string)          { enc.je.OpenNamespace(key) }
func (enc *syslogEncoder) AddString(key, val string)         { enc.je.AddString(key, val) }
func (enc *syslogEncoder) AddTime(key string, val time.Time) { enc.je.AddTime(key, val) }
func (enc *syslogEncoder) AddUint64(key string, val uint64)  { enc.je.AddUint64(key, val) }

func (enc *syslogEncoder) AppendArray(arr zapcore.ArrayMarshaler) error {
	return enc.je.AppendArray(arr)
}

func (enc *syslogEncoder) AppendObject(obj zapcore.ObjectMarshaler) error {
	return enc.je.AppendObject(obj)
}

func (enc *syslogEncoder) AppendBool(val bool)              { enc.je.AppendBool(val) }
func (enc *syslogEncoder) AppendByteString(val []byte)      { enc.je.AppendByteString(val) }
func (enc *syslogEncoder) AppendComplex128(val complex128)  { enc.je.AppendComplex128(val) }
func (enc *syslogEncoder) AppendDuration(val time.Duration) { enc.je.AppendDuration(val) }
func (enc *syslogEncoder) AppendInt64(val int64)            { enc.je.AppendInt64(val) }

func (enc *syslogEncoder) AppendReflected(val interface{}) error {
	return enc.je.AppendReflected(val)
}

func (enc *syslogEncoder) AppendString(val string)            { enc.je.AppendString(val) }
func (enc *syslogEncoder) AppendTime(val time.Time)           { enc.je.AppendTime(val) }
func (enc *syslogEncoder) AppendUint64(val uint64)            { enc.je.AppendUint64(val) }
func (enc *syslogEncoder) AddComplex64(k string, v complex64) { enc.je.AddComplex64(k, v) }
func (enc *syslogEncoder) AddFloat32(k string, v float32)     { enc.je.AddFloat32(k, v) }
func (enc *syslogEncoder) AddInt(k string, v int)             { enc.je.AddInt(k, v) }
func (enc *syslogEncoder) AddInt32(k string, v int32)         { enc.je.AddInt32(k, v) }
func (enc *syslogEncoder) AddInt16(k string, v int16)         { enc.je.AddInt16(k, v) }
func (enc *syslogEncoder) AddInt8(k string, v int8)           { enc.je.AddInt8(k, v) }
func (enc *syslogEncoder) AddUint(k string, v uint)           { enc.je.AddUint(k, v) }
func (enc *syslogEncoder) AddUint32(k string, v uint32)       { enc.je.AddUint32(k, v) }
func (enc *syslogEncoder) AddUint16(k string, v uint16)       { enc.je.AddUint16(k, v) }
func (enc *syslogEncoder) AddUint8(k string, v uint8)         { enc.je.AddUint8(k, v) }
func (enc *syslogEncoder) AddUintptr(k string, v uintptr)     { enc.je.AddUintptr(k, v) }
func (enc *syslogEncoder) AppendComplex64(v complex64)        { enc.je.AppendComplex64(v) }
func (enc *syslogEncoder) AppendFloat64(v float64)            { enc.je.AppendFloat64(v) }
func (enc *syslogEncoder) AppendFloat32(v float32)            { enc.je.AppendFloat32(v) }
func (enc *syslogEncoder) AppendInt(v int)                    { enc.je.AppendInt(v) }
func (enc *syslogEncoder) AppendInt32(v int32)                { enc.je.AppendInt32(v) }
func (enc *syslogEncoder) AppendInt16(v int16)                { enc.je.AppendInt16(v) }
func (enc *syslogEncoder) AppendInt8(v int8)                  { enc.je.AppendInt8(v) }
func (enc *syslogEncoder) AppendUint(v uint)                  { enc.je.AppendUint(v) }
func (enc *syslogEncoder) AppendUint32(v uint32)              { enc.je.AppendUint32(v) }
func (enc *syslogEncoder) AppendUint16(v uint16)              { enc.je.AppendUint16(v) }
func (enc *syslogEncoder) AppendUint8(v uint8)                { enc.je.AppendUint8(v) }
func (enc *syslogEncoder) AppendUintptr(v uintptr)            { enc.je.AppendUintptr(v) }

func (enc *syslogEncoder) Clone() zapcore.Encoder {
	return enc.clone()
}

func (enc *syslogEncoder) clone() *syslogEncoder {
	clone := &syslogEncoder{
		SyslogEncoderConfig: enc.SyslogEncoderConfig,
		je:                  enc.je.Clone().(jsonEncoder),
	}
	return clone
}

func (enc *syslogEncoder) EncodeEntry(ent zapcore.Entry, fields []zapcore.Field) (*buffer.Buffer, error) {
	msg := bufferpool.Get()

	var p syslog.Priority
	switch ent.Level {
	case zapcore.FatalLevel:
		p = syslog.LOG_EMERG
	case zapcore.PanicLevel:
		p = syslog.LOG_CRIT
	case zapcore.DPanicLevel:
		p = syslog.LOG_CRIT
	case zapcore.ErrorLevel:
		p = syslog.LOG_ERR
	case zapcore.WarnLevel:
		p = syslog.LOG_WARNING
	case zapcore.InfoLevel:
		p = syslog.LOG_INFO
	case zapcore.DebugLevel:
		p = syslog.LOG_DEBUG
	}
	pr := int64((enc.Facility & facilityMask) | (p & severityMask))

	// <PRI>version
	msg.AppendByte('<')
	msg.AppendInt(pr)
	msg.AppendByte('>')
	msg.AppendInt(version)

	// SP TIMESTAMP
	ts := ent.Time.UTC().Format(timestampFormat)
	if ts == "" {
		ts = nilValue
	}
	msg.AppendByte(' ')
	msg.AppendString(ts)

	// SP HOSTNAME
	msg.AppendByte(' ')
	msg.AppendString(enc.Hostname)

	// SP APP-NAME
	msg.AppendByte(' ')
	msg.AppendString(enc.App)

	// SP PROCID
	msg.AppendByte(' ')
	msg.AppendInt(int64(enc.PID))

	// SP MSGID SP STRUCTURED-DATA (just ignore)
	msg.AppendString(" - -")

	// SP UTF8 MSG
	json, err := enc.je.EncodeEntry(ent, fields)
	if json.Len() > 0 {
		msg.AppendString(" \xef\xbb\xbf")
		bs := json.Bytes()
		if enc.Framing == OctetCountingFraming {
			// Strip trailing line feed
			bs = bs[:len(bs)-1]
		}
		msg.AppendString(internal.BytesToString(bs))
	}

	if enc.Framing != OctetCountingFraming {
		return msg, err
	}

	// SYSLOG-FRAME = MSG-LEN SP SYSLOG-MSG
	out := bufferpool.Get()
	out.AppendInt(int64(msg.Len()))
	out.AppendByte(' ')
	out.AppendString(internal.BytesToString(msg.Bytes()))
	return out, err
}
<commit_msg>Fix potential leak<commit_after>// Copyright (c) 2017 Timon Wong
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package zapsyslog

import (
	"os"
	"path"
	"strings"
	"time"

	syslog "github.com/timonwong/go-syslog"
	"github.com/timonwong/zap-syslog/internal"
	"github.com/timonwong/zap-syslog/internal/bufferpool"
	"go.uber.org/zap"
	"go.uber.org/zap/buffer"
	"go.uber.org/zap/zapcore"
)

const (
	version         = 1
	severityMask    = 0x07
	facilityMask    = 0xf8
	nilValue        = "-"
	timestampFormat = "2006-01-02T15:04:05.000000Z07:00" // RFC3339 with micro fraction seconds
	maxHostnameLen  = 255
	maxAppNameLen   = 48
)

var (
	_ zapcore.Encoder = &syslogEncoder{}
	_ jsonEncoder     = zapcore.NewJSONEncoder(zap.NewProductionEncoderConfig()).(jsonEncoder)
)

// Framing configures RFC6587 TCP transport framing.
type Framing int

// Framing.
const (
	NonTransparentFraming Framing = iota
	OctetCountingFraming
	DefaultFraming = NonTransparentFraming
)

type jsonEncoder interface {
	zapcore.Encoder
	zapcore.ArrayEncoder
}

// SyslogEncoderConfig allows users to configure the concrete encoders for zap syslog.
type SyslogEncoderConfig struct {
	zapcore.EncoderConfig

	Framing  Framing         `json:"framing" yaml:"framing"`
	Facility syslog.Priority `json:"facility" yaml:"facility"`
	Hostname string          `json:"hostname" yaml:"hostname"`
	PID      int             `json:"pid" yaml:"pid"`
	App      string          `json:"app" yaml:"app"`
}

type syslogEncoder struct {
	*SyslogEncoderConfig
	je jsonEncoder
}

func rfc5424CompliantASCIIMapper(r rune) rune {
	// PRINTUSASCII    = %d33-126
	if r < 33 || r > 126 {
		return '_'
	}
	return r
}

func toRFC5424CompliantASCIIString(s string) string {
	return strings.Map(rfc5424CompliantASCIIMapper, s)
}

// NewSyslogEncoder creates a syslogEncoder that you should not use because I'm not done with it yet.
func NewSyslogEncoder(cfg SyslogEncoderConfig) zapcore.Encoder {
	if cfg.Hostname == "" {
		hostname, _ := os.Hostname()
		cfg.Hostname = hostname
	}
	if cfg.Hostname == "" {
		cfg.Hostname = nilValue
	} else {
		hostname := toRFC5424CompliantASCIIString(cfg.Hostname)
		if len(hostname) > maxHostnameLen {
			hostname = hostname[:maxHostnameLen]
		}
		cfg.Hostname = hostname
	}

	if cfg.PID == 0 {
		cfg.PID = os.Getpid()
	}
	if cfg.App == "" {
		cfg.App = nilValue
	} else {
		app := cfg.App
		if len(app) > maxAppNameLen {
			app = path.Base(app)
		}
		if len(app) > maxAppNameLen {
			app = app[:maxAppNameLen]
		}
		app = toRFC5424CompliantASCIIString(app)
	}

	cfg.EncoderConfig.LineEnding = "\n"
	je := zapcore.NewJSONEncoder(cfg.EncoderConfig).(jsonEncoder)
	return &syslogEncoder{
		SyslogEncoderConfig: &cfg,
		je:                  je,
	}
}

func (enc *syslogEncoder) AddArray(key string, arr zapcore.ArrayMarshaler) error {
	return enc.je.AddArray(key, arr)
}

func (enc *syslogEncoder) AddObject(key string, obj zapcore.ObjectMarshaler) error {
	return enc.je.AddObject(key, obj)
}

func (enc *syslogEncoder) AddBinary(key string, val []byte)          { enc.je.AddBinary(key, val) }
func (enc *syslogEncoder) AddByteString(key string, val []byte)      { enc.je.AddByteString(key, val) }
func (enc *syslogEncoder) AddBool(key string, val bool)              { enc.je.AddBool(key, val) }
func (enc *syslogEncoder) AddComplex128(key string, val complex128)  { enc.je.AddComplex128(key, val) }
func (enc *syslogEncoder) AddDuration(key string, val time.Duration) { enc.je.AddDuration(key, val) }
func (enc *syslogEncoder) AddFloat64(key string, val float64)        { enc.je.AddFloat64(key, val) }
func (enc *syslogEncoder) AddInt64(key string, val int64)            { enc.je.AddInt64(key, val) }

func (enc *syslogEncoder) AddReflected(key string, obj interface{}) error {
	return enc.je.AddReflected(key, obj)
}

func (enc *syslogEncoder) OpenNamespace(key string)          { enc.je.OpenNamespace(key) }
func (enc *syslogEncoder) AddString(key, val string)         { enc.je.AddString(key, val) }
func (enc *syslogEncoder) AddTime(key string, val time.Time) { enc.je.AddTime(key, val) }
func (enc *syslogEncoder) AddUint64(key string, val uint64)  { enc.je.AddUint64(key, val) }

func (enc *syslogEncoder) AppendArray(arr zapcore.ArrayMarshaler) error {
	return enc.je.AppendArray(arr)
}

func (enc *syslogEncoder) AppendObject(obj zapcore.ObjectMarshaler) error {
	return enc.je.AppendObject(obj)
}

func (enc *syslogEncoder) AppendBool(val bool)              { enc.je.AppendBool(val) }
func (enc *syslogEncoder) AppendByteString(val []byte)      { enc.je.AppendByteString(val) }
func (enc *syslogEncoder) AppendComplex128(val complex128)  { enc.je.AppendComplex128(val) }
func (enc *syslogEncoder) AppendDuration(val time.Duration) { enc.je.AppendDuration(val) }
func (enc *syslogEncoder) AppendInt64(val int64)            { enc.je.AppendInt64(val) }

func (enc *syslogEncoder) AppendReflected(val interface{}) error {
	return enc.je.AppendReflected(val)
}

func (enc *syslogEncoder) AppendString(val string)            { enc.je.AppendString(val) }
func (enc *syslogEncoder) AppendTime(val time.Time)           { enc.je.AppendTime(val) }
func (enc *syslogEncoder) AppendUint64(val uint64)            { enc.je.AppendUint64(val) }
func (enc *syslogEncoder) AddComplex64(k string, v complex64) { enc.je.AddComplex64(k, v) }
func (enc *syslogEncoder) AddFloat32(k string, v float32)     { enc.je.AddFloat32(k, v) }
func (enc *syslogEncoder) AddInt(k string, v int)             { enc.je.AddInt(k, v) }
func (enc *syslogEncoder) AddInt32(k string, v int32)         { enc.je.AddInt32(k, v) }
func (enc *syslogEncoder) AddInt16(k string, v int16)         { enc.je.AddInt16(k, v) }
func (enc *syslogEncoder) AddInt8(k string, v int8)           { enc.je.AddInt8(k, v) }
func (enc *syslogEncoder) AddUint(k string, v uint)           { enc.je.AddUint(k, v) }
func (enc *syslogEncoder) AddUint32(k string, v uint32)       { enc.je.AddUint32(k, v) }
func (enc *syslogEncoder) AddUint16(k string, v uint16)       { enc.je.AddUint16(k, v) }
func (enc *syslogEncoder) AddUint8(k string, v uint8)         { enc.je.AddUint8(k, v) }
func (enc *syslogEncoder) AddUintptr(k string, v uintptr)     { enc.je.AddUintptr(k, v) }
func (enc *syslogEncoder) AppendComplex64(v complex64)        { enc.je.AppendComplex64(v) }
func (enc *syslogEncoder) AppendFloat64(v float64)            { enc.je.AppendFloat64(v) }
func (enc *syslogEncoder) AppendFloat32(v float32)            { enc.je.AppendFloat32(v) }
func (enc *syslogEncoder) AppendInt(v int)                    { enc.je.AppendInt(v) }
func (enc *syslogEncoder) AppendInt32(v int32)                { enc.je.AppendInt32(v) }
func (enc *syslogEncoder) AppendInt16(v int16)                { enc.je.AppendInt16(v) }
func (enc *syslogEncoder) AppendInt8(v int8)                  { enc.je.AppendInt8(v) }
func (enc *syslogEncoder) AppendUint(v uint)                  { enc.je.AppendUint(v) }
func (enc *syslogEncoder) AppendUint32(v uint32)              { enc.je.AppendUint32(v) }
func (enc *syslogEncoder) AppendUint16(v uint16)              { enc.je.AppendUint16(v) }
func (enc *syslogEncoder) AppendUint8(v uint8)                { enc.je.AppendUint8(v) }
func (enc *syslogEncoder) AppendUintptr(v uintptr)            { enc.je.AppendUintptr(v) }

func (enc *syslogEncoder) Clone() zapcore.Encoder {
	return enc.clone()
}

func (enc *syslogEncoder) clone() *syslogEncoder {
	clone := &syslogEncoder{
		SyslogEncoderConfig: enc.SyslogEncoderConfig,
		je:                  enc.je.Clone().(jsonEncoder),
	}
	return clone
}

func (enc *syslogEncoder) EncodeEntry(ent zapcore.Entry, fields []zapcore.Field) (*buffer.Buffer, error) {
	msg := bufferpool.Get()

	var p syslog.Priority
	switch ent.Level {
	case zapcore.FatalLevel:
		p = syslog.LOG_EMERG
	case zapcore.PanicLevel:
		p = syslog.LOG_CRIT
	case zapcore.DPanicLevel:
		p = syslog.LOG_CRIT
	case zapcore.ErrorLevel:
		p = syslog.LOG_ERR
	case zapcore.WarnLevel:
		p = syslog.LOG_WARNING
	case zapcore.InfoLevel:
		p = syslog.LOG_INFO
	case zapcore.DebugLevel:
		p = syslog.LOG_DEBUG
	}
	pr := int64((enc.Facility & facilityMask) | (p & severityMask))

	// <PRI>version
	msg.AppendByte('<')
	msg.AppendInt(pr)
	msg.AppendByte('>')
	msg.AppendInt(version)

	// SP TIMESTAMP
	ts := ent.Time.UTC().Format(timestampFormat)
	if ts == "" {
		ts = nilValue
	}
	msg.AppendByte(' ')
	msg.AppendString(ts)

	// SP HOSTNAME
	msg.AppendByte(' ')
	msg.AppendString(enc.Hostname)

	// SP APP-NAME
	msg.AppendByte(' ')
	msg.AppendString(enc.App)

	// SP PROCID
	msg.AppendByte(' ')
	msg.AppendInt(int64(enc.PID))

	// SP MSGID SP STRUCTURED-DATA (just ignore)
	msg.AppendString(" - -")

	// SP UTF8 MSG
	json, err := enc.je.EncodeEntry(ent, fields)
	if json.Len() > 0 {
		msg.AppendString(" \xef\xbb\xbf")
		bs := json.Bytes()
		if enc.Framing == OctetCountingFraming {
			// Strip trailing line feed
			bs = bs[:len(bs)-1]
		}
		msg.AppendString(internal.BytesToString(bs))
	}

	if enc.Framing != OctetCountingFraming {
		return msg, err
	}

	// SYSLOG-FRAME = MSG-LEN SP SYSLOG-MSG
	out := bufferpool.Get()
	out.AppendInt(int64(msg.Len()))
	out.AppendByte(' ')
	out.AppendString(internal.BytesToString(msg.Bytes()))
	msg.Free()
	return out, err
}
<|endoftext|>


<commit_before>package termios

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

const (
	IXON    = 0x00000400
	IXANY   = 0x00000800
	IXOFF   = 0x00001000
	CRTSCTS = 0x80000000
)

// Tcgetattr gets the current serial port settings.
func Tcgetattr(fd uintptr, argp *unix.Termios) error {
	return unix.IoctlSetTermios(int(fd), unix.TCGETS, argp)
}

// Tcsetattr sets the current serial port settings.
func Tcsetattr(fd, action uintptr, argp *unix.Termios) error {
	var request uintptr
	switch action {
	case TCSANOW:
		request = unix.TCSETS
	case TCSADRAIN:
		request = unix.TCSETSW
	case TCSAFLUSH:
		request = unix.TCSETSF
	default:
		return unix.EINVAL
	}
	return unix.IoctlSetTermios(int(fd), uint(request), argp)
}

// Tcsendbreak transmits a continuous stream of zero-valued bits for a specific
// duration, if the terminal is using asynchronous serial data transmission. If
// duration is zero, it transmits zero-valued bits for at least 0.25 seconds, and not more that 0.5 seconds.
// If duration is not zero, it sends zero-valued bits for some
// implementation-defined length of time.
func Tcsendbreak(fd uintptr, duration int) error {
	return unix.IoctlSetInt(int(fd), unix.TCSBRKP, duration)
}

// Tcdrain waits until all output written to the object referred to by fd has been transmitted.
func Tcdrain(fd uintptr) error {
	// simulate drain with TCSADRAIN
	var attr unix.Termios
	if err := Tcgetattr(fd, &attr); err != nil {
		return err
	}
	return Tcsetattr(fd, TCSADRAIN, &attr)
}

// Tcflush discards data written to the object referred to by fd but not transmitted, or data received but not read, depending on the value of selector.
func Tcflush(fd, selector uintptr) error {
	return unix.IoctlSetInt(int(fd), unix.TCFLSH, int(selector))
}

// Tiocinq returns the number of bytes in the input buffer.
func Tiocinq(fd uintptr) (int, error) {
	return unix.IoctlGetInt(int(fd), unix.TIOCINQ)
}

// Cfgetispeed returns the input baud rate stored in the termios structure.
func Cfgetispeed(attr *unix.Termios) uint32 { return attr.Ispeed }

// Cfgetospeed returns the output baud rate stored in the termios structure.
func Cfgetospeed(attr *unix.Termios) uint32 { return attr.Ospeed }
<commit_msg>fix termios_linux<commit_after>package termios

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

const (
	IXON    = 0x00000400
	IXANY   = 0x00000800
	IXOFF   = 0x00001000
	CRTSCTS = 0x80000000
)

// Tcgetattr gets the current serial port settings.
func Tcgetattr(fd uintptr) (*unix.Termios, error) {
	return unix.IoctlGetTermios(int(fd), unix.TCGETS)
}

// Tcsetattr sets the current serial port settings.
func Tcsetattr(fd, action uintptr, argp *unix.Termios) error {
	var request uintptr
	switch action {
	case TCSANOW:
		request = unix.TCSETS
	case TCSADRAIN:
		request = unix.TCSETSW
	case TCSAFLUSH:
		request = unix.TCSETSF
	default:
		return unix.EINVAL
	}
	return unix.IoctlSetTermios(int(fd), uint(request), argp)
}

// Tcsendbreak transmits a continuous stream of zero-valued bits for a specific
// duration, if the terminal is using asynchronous serial data transmission. If
// duration is zero, it transmits zero-valued bits for at least 0.25 seconds, and not more that 0.5 seconds.
// If duration is not zero, it sends zero-valued bits for some
// implementation-defined length of time.
func Tcsendbreak(fd uintptr, duration int) error {
	return unix.IoctlSetInt(int(fd), unix.TCSBRKP, duration)
}

// Tcdrain waits until all output written to the object referred to by fd has been transmitted.
func Tcdrain(fd uintptr) error {
	// simulate drain with TCSADRAIN
	attr, err := Tcgetattr(fd)
	if err != nil {
		return err
	}
	return Tcsetattr(fd, TCSADRAIN, attr)
}

// Tcflush discards data written to the object referred to by fd but not transmitted, or data received but not read, depending on the value of selector.
func Tcflush(fd, selector uintptr) error {
	return unix.IoctlSetInt(int(fd), unix.TCFLSH, int(selector))
}

// Tiocinq returns the number of bytes in the input buffer.
func Tiocinq(fd uintptr) (int, error) {
	return unix.IoctlGetInt(int(fd), unix.TIOCINQ)
}

// Cfgetispeed returns the input baud rate stored in the termios structure.
func Cfgetispeed(attr *unix.Termios) uint32 { return attr.Ispeed }

// Cfgetospeed returns the output baud rate stored in the termios structure.
func Cfgetospeed(attr *unix.Termios) uint32 { return attr.Ospeed }
<|endoftext|>


<commit_before>// errchk $G $D/$F.go

// Copyright 2009 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 main

import "fmt"

func main() {
	日本語 := 1;	// ERROR "identifier"
}
<commit_msg>change bug163 to use actual non-letter in identifier<commit_after>// errchk $G $D/$F.go

// Copyright 2009 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 main

import "fmt"

func main() {
	x⊛y := 1;	// ERROR "identifier"
}
<|endoftext|>


<commit_before>// errchk $G $D/$F.go

// Copyright 2010 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 main

type I1 interface {
	m() I2
	I2
}

type I2 interface {
	I1 // ERROR "loop|interface"
}

var i1 I1 = i2 // ERROR "missing m method|need type assertion"
var i2 I2
var i2a I2 = i1
<commit_msg>test: Match gccgo error messages.<commit_after>// errchk $G $D/$F.go

// Copyright 2010 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 main

type I1 interface {
	m() I2
	I2 // GCCGO_ERROR "loop|interface"
}

type I2 interface {
	I1 // GC_ERROR "loop|interface"
}

var i1 I1 = i2 // GC_ERROR "missing m method|need type assertion"
var i2 I2
var i2a I2 = i1
<|endoftext|>


<commit_before>package identify_test

import (
	"testing"
	"time"

	peer "github.com/ipfs/go-libp2p-peer"
	host "github.com/libp2p/go-libp2p/p2p/host"
	identify "github.com/libp2p/go-libp2p/p2p/protocol/identify"
	testutil "github.com/libp2p/go-libp2p/p2p/test/util"

	ma "github.com/jbenet/go-multiaddr"
	context "golang.org/x/net/context"
)

func subtestIDService(t *testing.T, postDialWait time.Duration) {

	// the generated networks should have the id service wired in.
	ctx := context.Background()
	h1 := testutil.GenHostSwarm(t, ctx)
	h2 := testutil.GenHostSwarm(t, ctx)

	h1p := h1.ID()
	h2p := h2.ID()

	testKnowsAddrs(t, h1, h2p, []ma.Multiaddr{}) // nothing
	testKnowsAddrs(t, h2, h1p, []ma.Multiaddr{}) // nothing

	h2pi := h2.Peerstore().PeerInfo(h2p)
	if err := h1.Connect(ctx, h2pi); err != nil {
		t.Fatal(err)
	}

	// we need to wait here if Dial returns before ID service is finished.
	if postDialWait > 0 {
		<-time.After(postDialWait)
	}

	// the IDService should be opened automatically, by the network.
	// what we should see now is that both peers know about each others listen addresses.
	t.Log("test peer1 has peer2 addrs correctly")
	testKnowsAddrs(t, h1, h2p, h2.Peerstore().Addrs(h2p)) // has them
	testHasProtocolVersions(t, h1, h2p)

	// now, this wait we do have to do. it's the wait for the Listening side
	// to be done identifying the connection.
	c := h2.Network().ConnsToPeer(h1.ID())
	if len(c) < 1 {
		t.Fatal("should have connection by now at least.")
	}
	<-h2.IDService().IdentifyWait(c[0])

	addrs := h1.Peerstore().Addrs(h1p)
	addrs = append(addrs, c[0].RemoteMultiaddr())

	// and the protocol versions.
	t.Log("test peer2 has peer1 addrs correctly")
	testKnowsAddrs(t, h2, h1p, addrs) // has them
	testHasProtocolVersions(t, h2, h1p)
}

func testKnowsAddrs(t *testing.T, h host.Host, p peer.ID, expected []ma.Multiaddr) {
	actual := h.Peerstore().Addrs(p)

	if len(actual) != len(expected) {
		t.Errorf("expected: %s", expected)
		t.Errorf("actual: %s", actual)
		t.Fatal("dont have the same addresses")
	}

	have := map[string]struct{}{}
	for _, addr := range actual {
		have[addr.String()] = struct{}{}
	}
	for _, addr := range expected {
		if _, found := have[addr.String()]; !found {
			t.Errorf("%s did not have addr for %s: %s", h.ID(), p, addr)
			// panic("ahhhhhhh")
		}
	}
}

func testHasProtocolVersions(t *testing.T, h host.Host, p peer.ID) {
	v, err := h.Peerstore().Get(p, "ProtocolVersion")
	if v == nil {
		t.Error("no protocol version")
		return
	}
	if v.(string) != identify.LibP2PVersion {
		t.Error("protocol mismatch", err)
	}
	v, err = h.Peerstore().Get(p, "AgentVersion")
	if v.(string) != identify.ClientVersion {
		t.Error("agent version mismatch", err)
	}
}

// TestIDServiceWait gives the ID service 100ms to finish after dialing
// this is becasue it used to be concurrent. Now, Dial wait till the
// id service is done.
func TestIDServiceWait(t *testing.T) {
	N := 3
	for i := 0; i < N; i++ {
		subtestIDService(t, 100*time.Millisecond)
	}
}

func TestIDServiceNoWait(t *testing.T) {
	N := 3
	for i := 0; i < N; i++ {
		subtestIDService(t, 0)
	}
}

func TestProtoMatching(t *testing.T) {
	tcp1, _ := ma.NewMultiaddr("/ip4/1.2.3.4/tcp/1234")
	tcp2, _ := ma.NewMultiaddr("/ip4/1.2.3.4/tcp/2345")
	tcp3, _ := ma.NewMultiaddr("/ip4/1.2.3.4/tcp/4567")
	utp, _ := ma.NewMultiaddr("/ip4/1.2.3.4/udp/1234/utp")

	if !identify.HasConsistentTransport(tcp1, []ma.Multiaddr{tcp2, tcp3, utp}) {
		t.Fatal("expected match")
	}

	if identify.HasConsistentTransport(utp, []ma.Multiaddr{tcp2, tcp3}) {
		t.Fatal("expected mismatch")
	}
}
<commit_msg>identify: tests verify we have public keys<commit_after>package identify_test

import (
	"testing"
	"time"

	ic "github.com/ipfs/go-libp2p-crypto"
	peer "github.com/ipfs/go-libp2p-peer"
	host "github.com/libp2p/go-libp2p/p2p/host"
	identify "github.com/libp2p/go-libp2p/p2p/protocol/identify"
	testutil "github.com/libp2p/go-libp2p/p2p/test/util"

	ma "github.com/jbenet/go-multiaddr"
	context "golang.org/x/net/context"
)

func subtestIDService(t *testing.T, postDialWait time.Duration) {

	// the generated networks should have the id service wired in.
	ctx := context.Background()
	h1 := testutil.GenHostSwarm(t, ctx)
	h2 := testutil.GenHostSwarm(t, ctx)

	h1p := h1.ID()
	h2p := h2.ID()

	testKnowsAddrs(t, h1, h2p, []ma.Multiaddr{}) // nothing
	testKnowsAddrs(t, h2, h1p, []ma.Multiaddr{}) // nothing

	h2pi := h2.Peerstore().PeerInfo(h2p)
	if err := h1.Connect(ctx, h2pi); err != nil {
		t.Fatal(err)
	}

	// we need to wait here if Dial returns before ID service is finished.
	if postDialWait > 0 {
		<-time.After(postDialWait)
	}

	// the IDService should be opened automatically, by the network.
	// what we should see now is that both peers know about each others listen addresses.
	t.Log("test peer1 has peer2 addrs correctly")
	testKnowsAddrs(t, h1, h2p, h2.Peerstore().Addrs(h2p)) // has them
	testHasProtocolVersions(t, h1, h2p)
	testHasPublicKey(t, h1, h2p, h2.Peerstore().PubKey(h2p)) // h1 should have h2's public key

	// now, this wait we do have to do. it's the wait for the Listening side
	// to be done identifying the connection.
	c := h2.Network().ConnsToPeer(h1.ID())
	if len(c) < 1 {
		t.Fatal("should have connection by now at least.")
	}
	<-h2.IDService().IdentifyWait(c[0])

	addrs := h1.Peerstore().Addrs(h1p)
	addrs = append(addrs, c[0].RemoteMultiaddr())

	// and the protocol versions.
	t.Log("test peer2 has peer1 addrs correctly")
	testKnowsAddrs(t, h2, h1p, addrs) // has them
	testHasProtocolVersions(t, h2, h1p)
	testHasPublicKey(t, h2, h1p, h1.Peerstore().PubKey(h1p)) // h1 should have h2's public key
}

func testKnowsAddrs(t *testing.T, h host.Host, p peer.ID, expected []ma.Multiaddr) {
	actual := h.Peerstore().Addrs(p)

	if len(actual) != len(expected) {
		t.Errorf("expected: %s", expected)
		t.Errorf("actual: %s", actual)
		t.Fatal("dont have the same addresses")
	}

	have := map[string]struct{}{}
	for _, addr := range actual {
		have[addr.String()] = struct{}{}
	}
	for _, addr := range expected {
		if _, found := have[addr.String()]; !found {
			t.Errorf("%s did not have addr for %s: %s", h.ID(), p, addr)
			// panic("ahhhhhhh")
		}
	}
}

func testHasProtocolVersions(t *testing.T, h host.Host, p peer.ID) {
	v, err := h.Peerstore().Get(p, "ProtocolVersion")
	if v == nil {
		t.Error("no protocol version")
		return
	}
	if v.(string) != identify.LibP2PVersion {
		t.Error("protocol mismatch", err)
	}
	v, err = h.Peerstore().Get(p, "AgentVersion")
	if v.(string) != identify.ClientVersion {
		t.Error("agent version mismatch", err)
	}
}

func testHasPublicKey(t *testing.T, h host.Host, p peer.ID, shouldBe ic.PubKey) {
	k := h.Peerstore().PubKey(p)
	if k == nil {
		t.Error("no public key")
		return
	}
	if !k.Equals(shouldBe) {
		t.Error("key mismatch")
		return
	}

	p2, err := peer.IDFromPublicKey(k)
	if err != nil {
		t.Error("could not make key")
	} else if p != p2 {
		t.Error("key does not match peerid")
	}
}

// TestIDServiceWait gives the ID service 100ms to finish after dialing
// this is becasue it used to be concurrent. Now, Dial wait till the
// id service is done.
func TestIDServiceWait(t *testing.T) {
	N := 3
	for i := 0; i < N; i++ {
		subtestIDService(t, 100*time.Millisecond)
	}
}

func TestIDServiceNoWait(t *testing.T) {
	N := 3
	for i := 0; i < N; i++ {
		subtestIDService(t, 0)
	}
}

func TestProtoMatching(t *testing.T) {
	tcp1, _ := ma.NewMultiaddr("/ip4/1.2.3.4/tcp/1234")
	tcp2, _ := ma.NewMultiaddr("/ip4/1.2.3.4/tcp/2345")
	tcp3, _ := ma.NewMultiaddr("/ip4/1.2.3.4/tcp/4567")
	utp, _ := ma.NewMultiaddr("/ip4/1.2.3.4/udp/1234/utp")

	if !identify.HasConsistentTransport(tcp1, []ma.Multiaddr{tcp2, tcp3, utp}) {
		t.Fatal("expected match")
	}

	if identify.HasConsistentTransport(utp, []ma.Multiaddr{tcp2, tcp3}) {
		t.Fatal("expected mismatch")
	}
}
<|endoftext|>


<commit_before>// Copyright 2021 PingCAP, 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 cpuprofile

import (
	"bytes"
	"context"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"runtime/pprof"
	"sync"
	"testing"
	"time"

	"github.com/google/pprof/profile"
	"github.com/pingcap/tidb/util/cpuprofile/testutil"
	"github.com/pingcap/tidb/util/testbridge"
	"github.com/stretchr/testify/require"
	"go.uber.org/goleak"
)

func TestMain(m *testing.M) {
	testbridge.SetupForCommonTest()
	// To speed up testing
	DefProfileDuration = time.Millisecond * 200
	goleak.VerifyTestMain(m)
}

func TestBasicAPI(t *testing.T) {
	err := StartCPUProfiler()
	require.NoError(t, err)
	defer StopCPUProfiler()

	err = StartCPUProfiler()
	require.Equal(t, err, errProfilerAlreadyStarted)

	// Test for close multiple times.
	StopCPUProfiler()
	StopCPUProfiler()

	globalCPUProfiler = newParallelCPUProfiler()
	StopCPUProfiler()
	err = StartCPUProfiler()
	require.NoError(t, err)
	err = StartCPUProfiler()
	require.Equal(t, err, errProfilerAlreadyStarted)
}

func TestParallelCPUProfiler(t *testing.T) {
	err := StartCPUProfiler()
	require.NoError(t, err)
	defer StopCPUProfiler()

	// Test register/unregister nil
	Register(nil)
	require.Equal(t, 0, globalCPUProfiler.consumersCount())
	Unregister(nil)
	require.Equal(t, 0, globalCPUProfiler.consumersCount())

	// Test profile error and duplicate register.
	dataCh := make(ProfileConsumer, 10)
	err = pprof.StartCPUProfile(bytes.NewBuffer(nil))
	require.NoError(t, err)

	// Test for duplicate register.
	Register(dataCh)
	Register(dataCh)
	require.Equal(t, 1, globalCPUProfiler.consumersCount())

	// Test profile error
	data := <-dataCh
	require.Equal(t, "cpu profiling already in use", data.Error.Error())
	Unregister(dataCh)
	require.Equal(t, 0, globalCPUProfiler.consumersCount())

	// shouldn't receive data from a unregistered consumer.
	data = nil
	select {
	case data = <-dataCh:
	default:
	}
	require.Nil(t, data)

	// unregister not exist consumer
	Unregister(dataCh)
	require.Equal(t, 0, globalCPUProfiler.consumersCount())

	// Test register a closed consumer
	dataCh = make(ProfileConsumer, 10)
	close(dataCh)
	Register(dataCh)
	require.Equal(t, 1, globalCPUProfiler.consumersCount())
	data, ok := <-dataCh
	require.Nil(t, data)
	require.False(t, ok)
	Unregister(dataCh)
	require.Equal(t, 0, globalCPUProfiler.consumersCount())
	pprof.StopCPUProfile()

	// Test successfully get profile data.
	dataCh = make(ProfileConsumer, 10)
	Register(dataCh)
	data = <-dataCh
	require.NoError(t, data.Error)
	profileData, err := profile.ParseData(data.Data.Bytes())
	require.NoError(t, err)
	require.NotNil(t, profileData)
	Unregister(dataCh)
	require.Equal(t, 0, globalCPUProfiler.consumersCount())

	// Test stop profiling when no consumer.
	Register(dataCh)
	for {
		// wait for parallelCPUProfiler do profiling successfully
		err = pprof.StartCPUProfile(bytes.NewBuffer(nil))
		if err != nil {
			break
		}
		pprof.StopCPUProfile()
		time.Sleep(time.Millisecond)
	}
	Unregister(dataCh)
	require.Equal(t, 0, globalCPUProfiler.consumersCount())

	// wait for parallelCPUProfiler stop profiling
	start := time.Now()
	for {
		err = pprof.StartCPUProfile(bytes.NewBuffer(nil))
		if err == nil || time.Since(start) >= time.Second*2 {
			break
		}
		time.Sleep(time.Millisecond)
	}
	require.NoError(t, err)
	pprof.StopCPUProfile()
}

func getCPUProfile(d time.Duration, w io.Writer) error {
	pc := NewCollector()
	err := pc.StartCPUProfile(w)
	if err != nil {
		return err
	}
	time.Sleep(d)
	return pc.StopCPUProfile()
}

func TestGetCPUProfile(t *testing.T) {
	err := StartCPUProfiler()
	require.NoError(t, err)
	defer StopCPUProfiler()

	// Test profile error
	err = pprof.StartCPUProfile(bytes.NewBuffer(nil))
	require.NoError(t, err)
	err = getCPUProfile(time.Millisecond*50, bytes.NewBuffer(nil))
	require.Error(t, err)
	require.Equal(t, "cpu profiling already in use", err.Error())
	pprof.StopCPUProfile()

	// test parallel get CPU profile.
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	testutil.MockCPULoad(ctx, "sql", "sql_digest", "plan_digest")
	var wg sync.WaitGroup
	for i := 0; i < 3; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			buf := bytes.NewBuffer(nil)
			err = getCPUProfile(time.Millisecond*400, buf)
			require.NoError(t, err)
			profileData, err := profile.Parse(buf)
			require.NoError(t, err)

			labelCnt := 0
			for _, s := range profileData.Sample {
				for k := range s.Label {
					require.Equal(t, "sql", k)
					labelCnt++
				}
			}
			require.True(t, labelCnt > 1)
		}()
	}
	wg.Wait()
}

func TestProfileHTTPHandler(t *testing.T) {
	err := StartCPUProfiler()
	require.NoError(t, err)
	defer StopCPUProfiler()

	// setup http server
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	router := http.NewServeMux()
	router.HandleFunc("/debug/pprof/profile", ProfileHTTPHandler)
	httpServer := &http.Server{Handler: router, WriteTimeout: time.Second * 60}
	go func() {
		if err = httpServer.Serve(listener); err != nil && err != http.ErrServerClosed {
			require.NoError(t, err)
		}
	}()
	defer func() {
		err = httpServer.Close()
		require.NoError(t, err)
	}()

	address := listener.Addr().String()

	// Test for get profile success.
	resp, err := http.Get("http://" + address + "/debug/pprof/profile?seconds=1")
	require.NoError(t, err)
	require.Equal(t, 200, resp.StatusCode)
	profileData, err := profile.Parse(resp.Body)
	require.NoError(t, err)
	require.NotNil(t, profileData)
	require.NoError(t, resp.Body.Close())

	// Test for get profile failed.
	resp, err = http.Get("http://" + address + "/debug/pprof/profile?seconds=100000")
	require.NoError(t, err)
	require.Equal(t, 400, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	require.NoError(t, err)
	require.Equal(t, "profile duration exceeds server's WriteTimeout\n", string(body))
	require.NoError(t, resp.Body.Close())
}
<commit_msg>util: fix data race in the TestGetCPUProfile (#31246)<commit_after>// Copyright 2021 PingCAP, 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 cpuprofile

import (
	"bytes"
	"context"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"runtime/pprof"
	"sync"
	"testing"
	"time"

	"github.com/google/pprof/profile"
	"github.com/pingcap/tidb/util/cpuprofile/testutil"
	"github.com/pingcap/tidb/util/testbridge"
	"github.com/stretchr/testify/require"
	"go.uber.org/goleak"
)

func TestMain(m *testing.M) {
	testbridge.SetupForCommonTest()
	// To speed up testing
	DefProfileDuration = time.Millisecond * 200
	goleak.VerifyTestMain(m)
}

func TestBasicAPI(t *testing.T) {
	err := StartCPUProfiler()
	require.NoError(t, err)
	defer StopCPUProfiler()

	err = StartCPUProfiler()
	require.Equal(t, err, errProfilerAlreadyStarted)

	// Test for close multiple times.
	StopCPUProfiler()
	StopCPUProfiler()

	globalCPUProfiler = newParallelCPUProfiler()
	StopCPUProfiler()
	err = StartCPUProfiler()
	require.NoError(t, err)
	err = StartCPUProfiler()
	require.Equal(t, err, errProfilerAlreadyStarted)
}

func TestParallelCPUProfiler(t *testing.T) {
	err := StartCPUProfiler()
	require.NoError(t, err)
	defer StopCPUProfiler()

	// Test register/unregister nil
	Register(nil)
	require.Equal(t, 0, globalCPUProfiler.consumersCount())
	Unregister(nil)
	require.Equal(t, 0, globalCPUProfiler.consumersCount())

	// Test profile error and duplicate register.
	dataCh := make(ProfileConsumer, 10)
	err = pprof.StartCPUProfile(bytes.NewBuffer(nil))
	require.NoError(t, err)

	// Test for duplicate register.
	Register(dataCh)
	Register(dataCh)
	require.Equal(t, 1, globalCPUProfiler.consumersCount())

	// Test profile error
	data := <-dataCh
	require.Equal(t, "cpu profiling already in use", data.Error.Error())
	Unregister(dataCh)
	require.Equal(t, 0, globalCPUProfiler.consumersCount())

	// shouldn't receive data from a unregistered consumer.
	data = nil
	select {
	case data = <-dataCh:
	default:
	}
	require.Nil(t, data)

	// unregister not exist consumer
	Unregister(dataCh)
	require.Equal(t, 0, globalCPUProfiler.consumersCount())

	// Test register a closed consumer
	dataCh = make(ProfileConsumer, 10)
	close(dataCh)
	Register(dataCh)
	require.Equal(t, 1, globalCPUProfiler.consumersCount())
	data, ok := <-dataCh
	require.Nil(t, data)
	require.False(t, ok)
	Unregister(dataCh)
	require.Equal(t, 0, globalCPUProfiler.consumersCount())
	pprof.StopCPUProfile()

	// Test successfully get profile data.
	dataCh = make(ProfileConsumer, 10)
	Register(dataCh)
	data = <-dataCh
	require.NoError(t, data.Error)
	profileData, err := profile.ParseData(data.Data.Bytes())
	require.NoError(t, err)
	require.NotNil(t, profileData)
	Unregister(dataCh)
	require.Equal(t, 0, globalCPUProfiler.consumersCount())

	// Test stop profiling when no consumer.
	Register(dataCh)
	for {
		// wait for parallelCPUProfiler do profiling successfully
		err = pprof.StartCPUProfile(bytes.NewBuffer(nil))
		if err != nil {
			break
		}
		pprof.StopCPUProfile()
		time.Sleep(time.Millisecond)
	}
	Unregister(dataCh)
	require.Equal(t, 0, globalCPUProfiler.consumersCount())

	// wait for parallelCPUProfiler stop profiling
	start := time.Now()
	for {
		err = pprof.StartCPUProfile(bytes.NewBuffer(nil))
		if err == nil || time.Since(start) >= time.Second*2 {
			break
		}
		time.Sleep(time.Millisecond)
	}
	require.NoError(t, err)
	pprof.StopCPUProfile()
}

func getCPUProfile(d time.Duration, w io.Writer) error {
	pc := NewCollector()
	err := pc.StartCPUProfile(w)
	if err != nil {
		return err
	}
	time.Sleep(d)
	return pc.StopCPUProfile()
}

func TestGetCPUProfile(t *testing.T) {
	err := StartCPUProfiler()
	require.NoError(t, err)
	defer StopCPUProfiler()

	// Test profile error
	err = pprof.StartCPUProfile(bytes.NewBuffer(nil))
	require.NoError(t, err)
	err = getCPUProfile(time.Millisecond*50, bytes.NewBuffer(nil))
	require.Error(t, err)
	require.Equal(t, "cpu profiling already in use", err.Error())
	pprof.StopCPUProfile()

	// test parallel get CPU profile.
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	testutil.MockCPULoad(ctx, "sql", "sql_digest", "plan_digest")
	var wg sync.WaitGroup
	for i := 0; i < 3; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			var err error
			buf := bytes.NewBuffer(nil)
			err = getCPUProfile(time.Millisecond*400, buf)
			require.NoError(t, err)
			profileData, err := profile.Parse(buf)
			require.NoError(t, err)

			labelCnt := 0
			for _, s := range profileData.Sample {
				for k := range s.Label {
					require.Equal(t, "sql", k)
					labelCnt++
				}
			}
			require.True(t, labelCnt > 1)
		}()
	}
	wg.Wait()
}

func TestProfileHTTPHandler(t *testing.T) {
	err := StartCPUProfiler()
	require.NoError(t, err)
	defer StopCPUProfiler()

	// setup http server
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	router := http.NewServeMux()
	router.HandleFunc("/debug/pprof/profile", ProfileHTTPHandler)
	httpServer := &http.Server{Handler: router, WriteTimeout: time.Second * 60}
	go func() {
		if err = httpServer.Serve(listener); err != nil && err != http.ErrServerClosed {
			require.NoError(t, err)
		}
	}()
	defer func() {
		err = httpServer.Close()
		require.NoError(t, err)
	}()

	address := listener.Addr().String()

	// Test for get profile success.
	resp, err := http.Get("http://" + address + "/debug/pprof/profile?seconds=1")
	require.NoError(t, err)
	require.Equal(t, 200, resp.StatusCode)
	profileData, err := profile.Parse(resp.Body)
	require.NoError(t, err)
	require.NotNil(t, profileData)
	require.NoError(t, resp.Body.Close())

	// Test for get profile failed.
	resp, err = http.Get("http://" + address + "/debug/pprof/profile?seconds=100000")
	require.NoError(t, err)
	require.Equal(t, 400, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	require.NoError(t, err)
	require.Equal(t, "profile duration exceeds server's WriteTimeout\n", string(body))
	require.NoError(t, resp.Body.Close())
}
<|endoftext|>