<commit_before>package youtubedl

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"

	"github.com/wader/ydls/internal/writelogger"
)

type Printer interface {
	Printf(format string, v ...interface{})
}

type nopPrinter struct{}

func (nopPrinter) Printf(format string, v ...interface{}) {}

// Error youtubedl specific error
type Error string

func (e Error) Error() string {
	return string(e)
}

// Info youtubedl info
type Info struct {
	ID         string `json:"id"`
	Type       string `json:"_type"`
	URL        string `json:"url"`
	WebpageURL string `json:"webpage_url"`
	Direct     bool   `json:"direct"`

	Artist        string  `json:"artist"`
	Uploader      string  `json:"uploader"`
	UploadDate    string  `json:"upload_date"`
	Creator       string  `json:"creator"`
	Title         string  `json:"title"`
	PlaylistTitle string  `json:"playlist_title"`
	Episode       string  `json:"episode"`
	Description   string  `json:"description"`
	Duration      float64 `json:"duration"`
	Thumbnail     string  `json:"thumbnail"`
	// not unmarshalled, populated from image thumbnail file
	ThumbnailBytes []byte `json:"-"`

	Formats   []Format              `json:"formats"`
	Subtitles map[string][]Subtitle `json:"subtitles"`

	// Playlist entries if _type is playlist
	Entries []Info `json:"entries"`

	// Info can also be a mix of Info and one Format
	Format
}

// Format youtubedl downloadable format
type Format struct {
	FormatID string  `json:"format_id"`
	Protocol string  `json:"protocol"`
	Ext      string  `json:"ext"`
	ACodec   string  `json:"acodec"`
	VCodec   string  `json:"vcodec"`
	TBR      float64 `json:"tbr"`
	ABR      float64 `json:"abr"`
	VBR      float64 `json:"vbr"`
}

type Subtitle struct {
	URL      string `json:"url"`
	Ext      string `json:"ext"`
	Language string `json:"-"`
	// not unmarshalled, populated from subtitle file
	Bytes []byte `json:"-"`
}

func (f Format) String() string {
	return fmt.Sprintf("%s:%s:%s a:%s:%f v:%s:%f %f %f",
		f.FormatID,
		f.Protocol,
		f.Ext,
		f.NormalizedACodec(),
		f.ABR,
		f.NormalizedVCodec(),
		f.VBR,
		f.TBR,
		f.NormalizedBR(),
	)
}

func (f Format) NormalizedACodec() string {
	normCodec, normCodecFound := normalizeCodecName(f.ACodec)
	if normCodecFound {
		return normCodec
	}
	normCodec, _ = guessCodecFromExt(f.Ext)
	return normCodec
}

func (f Format) NormalizedVCodec() string {
	normCodec, normCodecFound := normalizeCodecName(f.VCodec)
	if normCodecFound {
		return normCodec
	}
	_, normCodec = guessCodecFromExt(f.Ext)
	return normCodec
}

func (f Format) NormalizedBR() float64 {
	if f.TBR != 0 {
		return f.TBR
	}
	return f.ABR + f.VBR
}

// guess codec from fuzzy codec name
func normalizeCodecName(c string) (string, bool) {
	codecNameNormalizeMap := map[string]string{
		"none": "",
		"avc1": "h264",
		"mp4a": "aac",
		"mp4v": "h264",
		"h265": "hevc",
	}

	// "  NAME.something  " -> "name"
	c = strings.Trim(c, " ")
	c = strings.ToLower(c)
	p := strings.SplitN(c, ".", 2)
	c = p[0]

	if n, ok := codecNameNormalizeMap[c]; ok {
		return n, true
	}

	return c, false
}

// guess codecs based on ext
func guessCodecFromExt(ext string) (acodec string, vcodec string) {
	switch strings.ToLower(ext) {
	case "wav":
		return "wav", ""
	case "mp3":
		return "mp3", ""
	case "ogg":
		return "vorbis", ""
	case "m4a",
		"aac":
		return "aac", ""
	case "mp4",
		"m4v",
		"mov",
		"3gp":
		return "aac", "h264"
	case "webm":
		return "opus", "vp9"
	case "flv":
		return "aac", "h264"
	}
	return "", ""
}

type Options struct {
	YesPlaylist       bool // prefer playlist
	PlaylistStart     uint
	PlaylistEnd       uint
	DownloadThumbnail bool
	DownloadSubtitles bool
	DebugLog          Printer
	HTTPClient        *http.Client
}

// New downloads metadata for URL
func New(ctx context.Context, rawURL string, options Options) (result Result, err error) {
	if options.DebugLog == nil {
		options.DebugLog = nopPrinter{}
	}
	if options.HTTPClient == nil {
		options.HTTPClient = http.DefaultClient
	}

	info, rawJSON, err := infoFromURL(ctx, rawURL, options)
	if err != nil {
		return Result{}, err
	}

	rawJSONCopy := make([]byte, len(rawJSON))
	copy(rawJSONCopy, rawJSON)

	return Result{
		Info:    info,
		RawJSON: rawJSONCopy,
		Options: options,
	}, nil
}

func infoFromURL(ctx context.Context, rawURL string, options Options) (info Info, rawJSON []byte, err error) {
	cmd := exec.CommandContext(
		ctx,
		"youtube-dl",
		"--no-call-home",
		"--no-cache-dir",
		"--skip-download",
		"--restrict-filenames",
		// provide URL via stdin for security, youtube-dl has some run command args
		"--batch-file", "-",
		"-J",
	)
	if options.YesPlaylist {
		cmd.Args = append(cmd.Args, "--yes-playlist")

		if options.PlaylistStart > 0 {
			cmd.Args = append(cmd.Args,
				"--playlist-start", strconv.Itoa(int(options.PlaylistStart)),
			)
		}
		if options.PlaylistEnd > 0 {
			cmd.Args = append(cmd.Args,
				"--playlist-end", strconv.Itoa(int(options.PlaylistEnd)),
			)
		}
	} else {
		cmd.Args = append(cmd.Args,
			"--no-playlist",
			"--all-subs",
		)
	}

	tempPath, _ := ioutil.TempDir("", "ydls-youtubedl")
	defer os.RemoveAll(tempPath)

	stdoutBuf := &bytes.Buffer{}
	stderrBuf := &bytes.Buffer{}

	ydlStderr := writelogger.New(options.DebugLog, "ydl-info stderr> ")
	stderrWriter := io.MultiWriter(stderrBuf, ydlStderr)

	cmd.Dir = tempPath
	cmd.Stdout = stdoutBuf
	cmd.Stderr = stderrWriter
	cmd.Stdin = bytes.NewBufferString(rawURL + "\n")

	options.DebugLog.Printf("cmd %v", cmd.Args)
	cmdErr := cmd.Run()

	stderrLineScanner := bufio.NewScanner(stderrBuf)
	errMessage := ""
	for stderrLineScanner.Scan() {
		const errorPrefix = "ERROR: "
		line := stderrLineScanner.Text()
		if strings.HasPrefix(line, errorPrefix) {
			errMessage = line[len(errorPrefix):]
		}
	}

	if errMessage != "" {
		return Info{}, nil, Error(errMessage)
	} else if cmdErr != nil {
		return Info{}, nil, cmdErr
	}

	if infoErr := json.Unmarshal(stdoutBuf.Bytes(), &info); infoErr != nil {
		return Info{}, nil, infoErr
	}

	if options.YesPlaylist && (info.Type != "playlist" || info.Type == "multi_video") {
		return Info{}, nil, fmt.Errorf("not a playlist")
	}

	if options.DownloadThumbnail && info.Thumbnail != "" {
		resp, respErr := options.HTTPClient.Get(info.Thumbnail)
		if respErr == nil {
			buf, _ := ioutil.ReadAll(resp.Body)
			resp.Body.Close()
			info.ThumbnailBytes = buf
		}
	}

	for language, subtitles := range info.Subtitles {
		for i := range subtitles {
			subtitles[i].Language = language
		}
	}

	if options.DownloadSubtitles {
		for _, subtitles := range info.Subtitles {
			for i, subtitle := range subtitles {
				resp, respErr := options.HTTPClient.Get(subtitle.URL)
				if respErr == nil {
					buf, _ := ioutil.ReadAll(resp.Body)
					resp.Body.Close()
					subtitles[i].Bytes = buf
				}
			}
		}
	}

	return info, stdoutBuf.Bytes(), nil
}

// Result metadata for a URL
type Result struct {
	Info    Info
	RawJSON []byte  // saved raw JSON. Used later when downloading
	Options Options // options passed to New
}

// DownloadResult download result
type DownloadResult struct {
	Reader io.ReadCloser
	waitCh chan struct{}
}

// Wait for resource cleanup
func (dr *DownloadResult) Wait() {
	<-dr.waitCh
}

// Formats return all formats
// helper to take care of mixed info and format
func (result Result) Formats() []Format {
	if len(result.Info.Formats) > 0 {
		return result.Info.Formats
	}
	return []Format{result.Info.Format}
}

// Download format matched by filter
func (result Result) Download(ctx context.Context, filter string) (*DownloadResult, error) {
	debugLog := result.Options.DebugLog

	if result.Info.Type == "playlist" || result.Info.Type == "multi_video" {
		return nil, fmt.Errorf("is a playlist")
	}

	tempPath, tempErr := ioutil.TempDir("", "ydls-youtubedl")
	if tempErr != nil {
		return nil, tempErr
	}
	jsonTempPath := path.Join(tempPath, "info.json")
	if err := ioutil.WriteFile(jsonTempPath, result.RawJSON, 0600); err != nil {
		os.RemoveAll(tempPath)
		return nil, err
	}

	dr := &DownloadResult{
		waitCh: make(chan struct{}),
	}

	cmd := exec.CommandContext(
		ctx,
		"youtube-dl",
		"--no-call-home",
		"--no-cache-dir",
		"--ignore-errors",
		"--newline",
		"--restrict-filenames",
		"--load-info", jsonTempPath,
		"-o", "-",
	)
	// don't need to specify if direct as there is only one
	// also seems to be issues when using filter with generic extractor
	if !result.Info.Direct {
		cmd.Args = append(cmd.Args, "-f", filter)
	}

	cmd.Dir = tempPath
	var w io.WriteCloser
	dr.Reader, w = io.Pipe()
	cmd.Stdout = w
	cmd.Stderr = writelogger.New(debugLog, fmt.Sprintf("ydl-dl %s stderr> ", filter))

	debugLog.Printf("cmd %v", cmd.Args)
	if err := cmd.Start(); err != nil {
		os.RemoveAll(tempPath)
		return nil, err
	}

	go func() {
		cmd.Wait()
		w.Close()
		os.RemoveAll(tempPath)
		close(dr.waitCh)
	}()

	return dr, nil
}
<commit_msg>Generate some Info and Format fields from youtube-dl README<commit_after>package youtubedl

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"

	"github.com/wader/ydls/internal/writelogger"
)

type Printer interface {
	Printf(format string, v ...interface{})
}

type nopPrinter struct{}

func (nopPrinter) Printf(format string, v ...interface{}) {}

// Error youtubedl specific error
type Error string

func (e Error) Error() string {
	return string(e)
}

// Info youtubedl info
type Info struct {
	// Generated from youtube-dl README using:
	// sed -e 's/ - `\(.*\)` (\(.*\)): \(.*\)/\1 \2 `json:"\1"` \/\/ \3/' | sed -e 's/numeric/float64/' | sed -e 's/boolean/bool/' | sed -e 's/_id/ID/'  | sed -e 's/_count/Count/'| sed -e 's/_uploader/Uploader/' | sed -e 's/_key/Key/' | sed -e 's/_year/Year/' | sed -e 's/_title/Title/' | sed -e 's/_rating/Rating/'  | sed -e 's/_number/Number/'  | awk '{print toupper(substr($0, 0, 1))  substr($0, 2)}'
	ID                 string  `json:"id"`                   // Video identifier
	Title              string  `json:"title"`                // Video title
	URL                string  `json:"url"`                  // Video URL
	AltTitle           string  `json:"alt_title"`            // A secondary title of the video
	DisplayID          string  `json:"display_id"`           // An alternative identifier for the video
	Uploader           string  `json:"uploader"`             // Full name of the video uploader
	License            string  `json:"license"`              // License name the video is licensed under
	Creator            string  `json:"creator"`              // The creator of the video
	ReleaseDate        string  `json:"release_date"`         // The date (YYYYMMDD) when the video was released
	Timestamp          float64 `json:"timestamp"`            // UNIX timestamp of the moment the video became available
	UploadDate         string  `json:"upload_date"`          // Video upload date (YYYYMMDD)
	UploaderID         string  `json:"uploader_id"`          // Nickname or id of the video uploader
	Channel            string  `json:"channel"`              // Full name of the channel the video is uploaded on
	ChannelID          string  `json:"channel_id"`           // Id of the channel
	Location           string  `json:"location"`             // Physical location where the video was filmed
	Duration           float64 `json:"duration"`             // Length of the video in seconds
	ViewCount          float64 `json:"view_count"`           // How many users have watched the video on the platform
	LikeCount          float64 `json:"like_count"`           // Number of positive ratings of the video
	DislikeCount       float64 `json:"dislike_count"`        // Number of negative ratings of the video
	RepostCount        float64 `json:"repost_count"`         // Number of reposts of the video
	AverageRating      float64 `json:"average_rating"`       // Average rating give by users, the scale used depends on the webpage
	CommentCount       float64 `json:"comment_count"`        // Number of comments on the video
	AgeLimit           float64 `json:"age_limit"`            // Age restriction for the video (years)
	IsLive             bool    `json:"is_live"`              // Whether this video is a live stream or a fixed-length video
	StartTime          float64 `json:"start_time"`           // Time in seconds where the reproduction should start, as specified in the URL
	EndTime            float64 `json:"end_time"`             // Time in seconds where the reproduction should end, as specified in the URL
	Extractor          string  `json:"extractor"`            // Name of the extractor
	ExtractorKey       string  `json:"extractor_key"`        // Key name of the extractor
	Epoch              float64 `json:"epoch"`                // Unix epoch when creating the file
	Autonumber         float64 `json:"autonumber"`           // Five-digit number that will be increased with each download, starting at zero
	Playlist           string  `json:"playlist"`             // Name or id of the playlist that contains the video
	PlaylistIndex      float64 `json:"playlist_index"`       // Index of the video in the playlist padded with leading zeros according to the total length of the playlist
	PlaylistID         string  `json:"playlist_id"`          // Playlist identifier
	PlaylistTitle      string  `json:"playlist_title"`       // Playlist title
	PlaylistUploader   string  `json:"playlist_uploader"`    // Full name of the playlist uploader
	PlaylistUploaderID string  `json:"playlist_uploader_id"` // Nickname or id of the playlist uploader

	// Available for the video that belongs to some logical chapter or section:
	Chapter       string  `json:"chapter"`        // Name or title of the chapter the video belongs to
	ChapterNumber float64 `json:"chapter_number"` // Number of the chapter the video belongs to
	ChapterID     string  `json:"chapter_id"`     // Id of the chapter the video belongs to

	// Available for the video that is an episode of some series or programme:
	Series        string  `json:"series"`         // Title of the series or programme the video episode belongs to
	Season        string  `json:"season"`         // Title of the season the video episode belongs to
	SeasonNumber  float64 `json:"season_number"`  // Number of the season the video episode belongs to
	SeasonID      string  `json:"season_id"`      // Id of the season the video episode belongs to
	Episode       string  `json:"episode"`        // Title of the video episode
	EpisodeNumber float64 `json:"episode_number"` // Number of the video episode within a season
	EpisodeID     string  `json:"episode_id"`     // Id of the video episode

	// Available for the media that is a track or a part of a music album:
	Track       string  `json:"track"`        // Title of the track
	TrackNumber float64 `json:"track_number"` // Number of the track within an album or a disc
	TrackID     string  `json:"track_id"`     // Id of the track
	Artist      string  `json:"artist"`       // Artist(s) of the track
	Genre       string  `json:"genre"`        // Genre(s) of the track
	Album       string  `json:"album"`        // Title of the album the track belongs to
	AlbumType   string  `json:"album_type"`   // Type of the album
	AlbumArtist string  `json:"album_artist"` // List of all artists appeared on the album
	DiscNumber  float64 `json:"disc_number"`  // Number of the disc or other physical medium the track belongs to
	ReleaseYear float64 `json:"release_year"` // Year (YYYY) when the album was released

	Type        string `json:"_type"`
	Direct      bool   `json:"direct"`
	WebpageURL  string `json:"webpage_url"`
	Description string `json:"description"`
	Thumbnail   string `json:"thumbnail"`
	// not unmarshalled, populated from image thumbnail file
	ThumbnailBytes []byte `json:"-"`

	Formats   []Format              `json:"formats"`
	Subtitles map[string][]Subtitle `json:"subtitles"`

	// Playlist entries if _type is playlist
	Entries []Info `json:"entries"`

	// Info can also be a mix of Info and one Format
	Format
}

// Format youtubedl downloadable format
type Format struct {
	Ext            string  `json:"ext"`             // Video filename extension
	Format         string  `json:"format"`          // A human-readable description of the format
	FormatID       string  `json:"format_id"`       // Format code specified by `--format`
	FormatNote     string  `json:"format_note"`     // Additional info about the format
	Width          float64 `json:"width"`           // Width of the video
	Height         float64 `json:"height"`          // Height of the video
	Resolution     string  `json:"resolution"`      // Textual description of width and height
	TBR            float64 `json:"tbr"`             // Average bitrate of audio and video in KBit/s
	ABR            float64 `json:"abr"`             // Average audio bitrate in KBit/s
	ACodec         string  `json:"acodec"`          // Name of the audio codec in use
	ASR            float64 `json:"asr"`             // Audio sampling rate in Hertz
	VBR            float64 `json:"vbr"`             // Average video bitrate in KBit/s
	FPS            float64 `json:"fps"`             // Frame rate
	VCodec         string  `json:"vcodec"`          // Name of the video codec in use
	Container      string  `json:"container"`       // Name of the container format
	Filesize       float64 `json:"filesize"`        // The number of bytes, if known in advance
	FilesizeApprox float64 `json:"filesize_approx"` // An estimate for the number of bytes
	Protocol       string  `json:"protocol"`        // The protocol that will be used for the actual download
}

type Subtitle struct {
	URL      string `json:"url"`
	Ext      string `json:"ext"`
	Language string `json:"-"`
	// not unmarshalled, populated from subtitle file
	Bytes []byte `json:"-"`
}

func (f Format) String() string {
	return fmt.Sprintf("%s:%s:%s a:%s:%f v:%s:%f %f %f",
		f.FormatID,
		f.Protocol,
		f.Ext,
		f.NormalizedACodec(),
		f.ABR,
		f.NormalizedVCodec(),
		f.VBR,
		f.TBR,
		f.NormalizedBR(),
	)
}

func (f Format) NormalizedACodec() string {
	normCodec, normCodecFound := normalizeCodecName(f.ACodec)
	if normCodecFound {
		return normCodec
	}
	normCodec, _ = guessCodecFromExt(f.Ext)
	return normCodec
}

func (f Format) NormalizedVCodec() string {
	normCodec, normCodecFound := normalizeCodecName(f.VCodec)
	if normCodecFound {
		return normCodec
	}
	_, normCodec = guessCodecFromExt(f.Ext)
	return normCodec
}

func (f Format) NormalizedBR() float64 {
	if f.TBR != 0 {
		return f.TBR
	}
	return f.ABR + f.VBR
}

// guess codec from fuzzy codec name
func normalizeCodecName(c string) (string, bool) {
	codecNameNormalizeMap := map[string]string{
		"none": "",
		"avc1": "h264",
		"mp4a": "aac",
		"mp4v": "h264",
		"h265": "hevc",
	}

	// "  NAME.something  " -> "name"
	c = strings.Trim(c, " ")
	c = strings.ToLower(c)
	p := strings.SplitN(c, ".", 2)
	c = p[0]

	if n, ok := codecNameNormalizeMap[c]; ok {
		return n, true
	}

	return c, false
}

// guess codecs based on ext
func guessCodecFromExt(ext string) (acodec string, vcodec string) {
	switch strings.ToLower(ext) {
	case "wav":
		return "wav", ""
	case "mp3":
		return "mp3", ""
	case "ogg":
		return "vorbis", ""
	case "m4a",
		"aac":
		return "aac", ""
	case "mp4",
		"m4v",
		"mov",
		"3gp":
		return "aac", "h264"
	case "webm":
		return "opus", "vp9"
	case "flv":
		return "aac", "h264"
	}
	return "", ""
}

type Options struct {
	YesPlaylist       bool // prefer playlist
	PlaylistStart     uint
	PlaylistEnd       uint
	DownloadThumbnail bool
	DownloadSubtitles bool
	DebugLog          Printer
	HTTPClient        *http.Client
}

// New downloads metadata for URL
func New(ctx context.Context, rawURL string, options Options) (result Result, err error) {
	if options.DebugLog == nil {
		options.DebugLog = nopPrinter{}
	}
	if options.HTTPClient == nil {
		options.HTTPClient = http.DefaultClient
	}

	info, rawJSON, err := infoFromURL(ctx, rawURL, options)
	if err != nil {
		return Result{}, err
	}

	rawJSONCopy := make([]byte, len(rawJSON))
	copy(rawJSONCopy, rawJSON)

	return Result{
		Info:    info,
		RawJSON: rawJSONCopy,
		Options: options,
	}, nil
}

func infoFromURL(ctx context.Context, rawURL string, options Options) (info Info, rawJSON []byte, err error) {
	cmd := exec.CommandContext(
		ctx,
		"youtube-dl",
		"--no-call-home",
		"--no-cache-dir",
		"--skip-download",
		"--restrict-filenames",
		// provide URL via stdin for security, youtube-dl has some run command args
		"--batch-file", "-",
		"-J",
	)
	if options.YesPlaylist {
		cmd.Args = append(cmd.Args, "--yes-playlist")

		if options.PlaylistStart > 0 {
			cmd.Args = append(cmd.Args,
				"--playlist-start", strconv.Itoa(int(options.PlaylistStart)),
			)
		}
		if options.PlaylistEnd > 0 {
			cmd.Args = append(cmd.Args,
				"--playlist-end", strconv.Itoa(int(options.PlaylistEnd)),
			)
		}
	} else {
		cmd.Args = append(cmd.Args,
			"--no-playlist",
			"--all-subs",
		)
	}

	tempPath, _ := ioutil.TempDir("", "ydls-youtubedl")
	defer os.RemoveAll(tempPath)

	stdoutBuf := &bytes.Buffer{}
	stderrBuf := &bytes.Buffer{}

	ydlStderr := writelogger.New(options.DebugLog, "ydl-info stderr> ")
	stderrWriter := io.MultiWriter(stderrBuf, ydlStderr)

	cmd.Dir = tempPath
	cmd.Stdout = stdoutBuf
	cmd.Stderr = stderrWriter
	cmd.Stdin = bytes.NewBufferString(rawURL + "\n")

	options.DebugLog.Printf("cmd %v", cmd.Args)
	cmdErr := cmd.Run()

	stderrLineScanner := bufio.NewScanner(stderrBuf)
	errMessage := ""
	for stderrLineScanner.Scan() {
		const errorPrefix = "ERROR: "
		line := stderrLineScanner.Text()
		if strings.HasPrefix(line, errorPrefix) {
			errMessage = line[len(errorPrefix):]
		}
	}

	if errMessage != "" {
		return Info{}, nil, Error(errMessage)
	} else if cmdErr != nil {
		return Info{}, nil, cmdErr
	}

	if infoErr := json.Unmarshal(stdoutBuf.Bytes(), &info); infoErr != nil {
		return Info{}, nil, infoErr
	}

	if options.YesPlaylist && (info.Type != "playlist" || info.Type == "multi_video") {
		return Info{}, nil, fmt.Errorf("not a playlist")
	}

	if options.DownloadThumbnail && info.Thumbnail != "" {
		resp, respErr := options.HTTPClient.Get(info.Thumbnail)
		if respErr == nil {
			buf, _ := ioutil.ReadAll(resp.Body)
			resp.Body.Close()
			info.ThumbnailBytes = buf
		}
	}

	for language, subtitles := range info.Subtitles {
		for i := range subtitles {
			subtitles[i].Language = language
		}
	}

	if options.DownloadSubtitles {
		for _, subtitles := range info.Subtitles {
			for i, subtitle := range subtitles {
				resp, respErr := options.HTTPClient.Get(subtitle.URL)
				if respErr == nil {
					buf, _ := ioutil.ReadAll(resp.Body)
					resp.Body.Close()
					subtitles[i].Bytes = buf
				}
			}
		}
	}

	return info, stdoutBuf.Bytes(), nil
}

// Result metadata for a URL
type Result struct {
	Info    Info
	RawJSON []byte  // saved raw JSON. Used later when downloading
	Options Options // options passed to New
}

// DownloadResult download result
type DownloadResult struct {
	Reader io.ReadCloser
	waitCh chan struct{}
}

// Wait for resource cleanup
func (dr *DownloadResult) Wait() {
	<-dr.waitCh
}

// Formats return all formats
// helper to take care of mixed info and format
func (result Result) Formats() []Format {
	if len(result.Info.Formats) > 0 {
		return result.Info.Formats
	}
	return []Format{result.Info.Format}
}

// Download format matched by filter
func (result Result) Download(ctx context.Context, filter string) (*DownloadResult, error) {
	debugLog := result.Options.DebugLog

	if result.Info.Type == "playlist" || result.Info.Type == "multi_video" {
		return nil, fmt.Errorf("is a playlist")
	}

	tempPath, tempErr := ioutil.TempDir("", "ydls-youtubedl")
	if tempErr != nil {
		return nil, tempErr
	}
	jsonTempPath := path.Join(tempPath, "info.json")
	if err := ioutil.WriteFile(jsonTempPath, result.RawJSON, 0600); err != nil {
		os.RemoveAll(tempPath)
		return nil, err
	}

	dr := &DownloadResult{
		waitCh: make(chan struct{}),
	}

	cmd := exec.CommandContext(
		ctx,
		"youtube-dl",
		"--no-call-home",
		"--no-cache-dir",
		"--ignore-errors",
		"--newline",
		"--restrict-filenames",
		"--load-info", jsonTempPath,
		"-o", "-",
	)
	// don't need to specify if direct as there is only one
	// also seems to be issues when using filter with generic extractor
	if !result.Info.Direct {
		cmd.Args = append(cmd.Args, "-f", filter)
	}

	cmd.Dir = tempPath
	var w io.WriteCloser
	dr.Reader, w = io.Pipe()
	cmd.Stdout = w
	cmd.Stderr = writelogger.New(debugLog, fmt.Sprintf("ydl-dl %s stderr> ", filter))

	debugLog.Printf("cmd %v", cmd.Args)
	if err := cmd.Start(); err != nil {
		os.RemoveAll(tempPath)
		return nil, err
	}

	go func() {
		cmd.Wait()
		w.Close()
		os.RemoveAll(tempPath)
		close(dr.waitCh)
	}()

	return dr, nil
}
<|endoftext|>


<commit_before>package iodaemon_test

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path/filepath"

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

	"testing"
)

var iodaemonBinPath string
var testPrintSignalBinPath string

var tmpdir string
var socketPath string

type CompiledAssets struct {
	IoDaemonBinPath        string
	TestPrintSignalBinPath string
}

var _ = SynchronizedBeforeSuite(func() []byte {
	var err error
	assets := CompiledAssets{}
	assets.IoDaemonBinPath, err = gexec.Build("github.com/cloudfoundry-incubator/garden-linux/iodaemon/cmd/iodaemon", "-race")
	Expect(err).ToNot(HaveOccurred())

	assets.TestPrintSignalBinPath, err = gexec.Build("github.com/cloudfoundry-incubator/garden-linux/iodaemon/test_print_signals")
	Expect(err).ToNot(HaveOccurred())

	marshalledAssets, err := json.Marshal(assets)
	Expect(err).ToNot(HaveOccurred())
	return marshalledAssets
}, func(marshalledAssets []byte) {
	assets := CompiledAssets{}
	err := json.Unmarshal(marshalledAssets, &assets)
	Expect(err).ToNot(HaveOccurred())
	iodaemonBinPath = assets.IoDaemonBinPath
	testPrintSignalBinPath = assets.TestPrintSignalBinPath
})

var _ = SynchronizedAfterSuite(func() {
	//noop
}, func() {
	gexec.CleanupBuildArtifacts()
})

var _ = BeforeEach(func() {
	var err error

	tmpdir, err = ioutil.TempDir("", "socket-dir")
	Expect(err).ToNot(HaveOccurred())

	socketPath = filepath.Join(tmpdir, "iodaemon.sock")
})

var _ = AfterEach(func() {
	os.RemoveAll(tmpdir)
})

func TestIodaemon(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Iodaemon Suite")
}
<commit_msg>Reduce iodaemon test run time by ~100 seconds<commit_after>package iodaemon_test

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path/filepath"

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

	"testing"
)

var iodaemonBinPath string
var testPrintSignalBinPath string

var tmpdir string
var socketPath string

type CompiledAssets struct {
	IoDaemonBinPath        string
	TestPrintSignalBinPath string
}

var _ = SynchronizedBeforeSuite(func() []byte {
	var err error
	assets := CompiledAssets{}
	assets.IoDaemonBinPath, err = gexec.Build("github.com/cloudfoundry-incubator/garden-linux/iodaemon/cmd/iodaemon")
	Expect(err).ToNot(HaveOccurred())

	assets.TestPrintSignalBinPath, err = gexec.Build("github.com/cloudfoundry-incubator/garden-linux/iodaemon/test_print_signals")
	Expect(err).ToNot(HaveOccurred())

	marshalledAssets, err := json.Marshal(assets)
	Expect(err).ToNot(HaveOccurred())
	return marshalledAssets
}, func(marshalledAssets []byte) {
	assets := CompiledAssets{}
	err := json.Unmarshal(marshalledAssets, &assets)
	Expect(err).ToNot(HaveOccurred())
	iodaemonBinPath = assets.IoDaemonBinPath
	testPrintSignalBinPath = assets.TestPrintSignalBinPath
})

var _ = SynchronizedAfterSuite(func() {
	//noop
}, func() {
	gexec.CleanupBuildArtifacts()
})

var _ = BeforeEach(func() {
	var err error

	tmpdir, err = ioutil.TempDir("", "socket-dir")
	Expect(err).ToNot(HaveOccurred())

	socketPath = filepath.Join(tmpdir, "iodaemon.sock")
})

var _ = AfterEach(func() {
	os.RemoveAll(tmpdir)
})

func TestIodaemon(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Iodaemon Suite")
}
<|endoftext|>


<commit_before>package gcssink

import (
	"context"
	"fmt"
	"log"
	"os"

	"cloud.google.com/go/storage"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/replication/source"
	"github.com/chrislusf/seaweedfs/weed/util"
	"github.com/chrislusf/seaweedfs/weed/filer2"
	"github.com/chrislusf/seaweedfs/weed/replication/sink"
	"google.golang.org/api/option"
)

type GcsSink struct {
	client      *storage.Client
	bucket      string
	dir         string
	filerSource *source.FilerSource
}

func init() {
	sink.Sinks = append(sink.Sinks, &GcsSink{})
}

func (g *GcsSink) GetName() string {
	return "google_cloud_storage"
}

func (g *GcsSink) GetSinkToDirectory() string {
	return g.dir
}

func (g *GcsSink) Initialize(configuration util.Configuration) error {
	return g.initialize(
		configuration.GetString("google_application_credentials"),
		configuration.GetString("bucket"),
		configuration.GetString("directory"),
	)
}

func (g *GcsSink) SetSourceFiler(s *source.FilerSource) {
	g.filerSource = s
}

func (g *GcsSink) initialize(google_application_credentials, bucketName, dir string) (error) {
	g.bucket = bucketName
	g.dir = dir

	ctx := context.Background()
	// Creates a client.
	if google_application_credentials == "" {
		var found bool
		google_application_credentials, found = os.LookupEnv("GOOGLE_APPLICATION_CREDENTIALS")
		if !found {
			log.Fatalf("need to specific GOOGLE_APPLICATION_CREDENTIALS env variable or google_application_credentials in replication.toml")
		}
	}
	client, err := storage.NewClient(ctx, option.WithCredentialsFile(google_application_credentials))
	if err != nil {
		log.Fatalf("Failed to create client: %v", err)
	}

	g.client = client

	return nil
}

func (g *GcsSink) DeleteEntry(key string, isDirectory, deleteIncludeChunks bool) error {

	if isDirectory {
		key = key + "/"
	}

	if err := g.client.Bucket(g.bucket).Object(key).Delete(context.Background()); err != nil {
		return fmt.Errorf("gcs delete %s %s", g.bucket, key)
	}

	return nil

}

func (g *GcsSink) CreateEntry(key string, entry *filer_pb.Entry) error {

	if entry.IsDirectory {
		return nil
	}

	totalSize := filer2.TotalSize(entry.Chunks)
	chunkViews := filer2.ViewFromChunks(entry.Chunks, 0, int(totalSize))

	ctx := context.Background()

	wc := g.client.Bucket(g.bucket).Object(key).NewWriter(ctx)

	for _, chunk := range chunkViews {

		fileUrl, err := g.filerSource.LookupFileId(chunk.FileId)
		if err != nil {
			return err
		}

		_, err = util.ReadUrlAsStream(fileUrl, chunk.Offset, int(chunk.Size), func(data []byte) {
			wc.Write(data)
		})

		if err != nil {
			return err
		}

	}

	if err := wc.Close(); err != nil {
		return err
	}

	return nil

}

func (g *GcsSink) UpdateEntry(key string, oldEntry, newEntry *filer_pb.Entry, deleteIncludeChunks bool) (foundExistingEntry bool, err error) {
	// TODO improve efficiency
	return false, nil
}
<commit_msg>adjust error format<commit_after>package gcssink

import (
	"context"
	"fmt"
	"log"
	"os"

	"cloud.google.com/go/storage"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/replication/source"
	"github.com/chrislusf/seaweedfs/weed/util"
	"github.com/chrislusf/seaweedfs/weed/filer2"
	"github.com/chrislusf/seaweedfs/weed/replication/sink"
	"google.golang.org/api/option"
)

type GcsSink struct {
	client      *storage.Client
	bucket      string
	dir         string
	filerSource *source.FilerSource
}

func init() {
	sink.Sinks = append(sink.Sinks, &GcsSink{})
}

func (g *GcsSink) GetName() string {
	return "google_cloud_storage"
}

func (g *GcsSink) GetSinkToDirectory() string {
	return g.dir
}

func (g *GcsSink) Initialize(configuration util.Configuration) error {
	return g.initialize(
		configuration.GetString("google_application_credentials"),
		configuration.GetString("bucket"),
		configuration.GetString("directory"),
	)
}

func (g *GcsSink) SetSourceFiler(s *source.FilerSource) {
	g.filerSource = s
}

func (g *GcsSink) initialize(google_application_credentials, bucketName, dir string) (error) {
	g.bucket = bucketName
	g.dir = dir

	ctx := context.Background()
	// Creates a client.
	if google_application_credentials == "" {
		var found bool
		google_application_credentials, found = os.LookupEnv("GOOGLE_APPLICATION_CREDENTIALS")
		if !found {
			log.Fatalf("need to specific GOOGLE_APPLICATION_CREDENTIALS env variable or google_application_credentials in replication.toml")
		}
	}
	client, err := storage.NewClient(ctx, option.WithCredentialsFile(google_application_credentials))
	if err != nil {
		log.Fatalf("Failed to create client: %v", err)
	}

	g.client = client

	return nil
}

func (g *GcsSink) DeleteEntry(key string, isDirectory, deleteIncludeChunks bool) error {

	if isDirectory {
		key = key + "/"
	}

	if err := g.client.Bucket(g.bucket).Object(key).Delete(context.Background()); err != nil {
		return fmt.Errorf("gcs delete %s%s: %v", g.bucket, key, err)
	}

	return nil

}

func (g *GcsSink) CreateEntry(key string, entry *filer_pb.Entry) error {

	if entry.IsDirectory {
		return nil
	}

	totalSize := filer2.TotalSize(entry.Chunks)
	chunkViews := filer2.ViewFromChunks(entry.Chunks, 0, int(totalSize))

	ctx := context.Background()

	wc := g.client.Bucket(g.bucket).Object(key).NewWriter(ctx)

	for _, chunk := range chunkViews {

		fileUrl, err := g.filerSource.LookupFileId(chunk.FileId)
		if err != nil {
			return err
		}

		_, err = util.ReadUrlAsStream(fileUrl, chunk.Offset, int(chunk.Size), func(data []byte) {
			wc.Write(data)
		})

		if err != nil {
			return err
		}

	}

	if err := wc.Close(); err != nil {
		return err
	}

	return nil

}

func (g *GcsSink) UpdateEntry(key string, oldEntry, newEntry *filer_pb.Entry, deleteIncludeChunks bool) (foundExistingEntry bool, err error) {
	// TODO improve efficiency
	return false, nil
}
<|endoftext|>


<commit_before>package weed_server

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"mime"
	"net/http"
	"net/url"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/images"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	xhttp "github.com/chrislusf/seaweedfs/weed/s3api/http"
	"github.com/chrislusf/seaweedfs/weed/stats"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func (fs *FilerServer) GetOrHeadHandler(w http.ResponseWriter, r *http.Request) {

	path := r.URL.Path
	isForDirectory := strings.HasSuffix(path, "/")
	if isForDirectory && len(path) > 1 {
		path = path[:len(path)-1]
	}

	entry, err := fs.filer.FindEntry(context.Background(), util.FullPath(path))
	if err != nil {
		if path == "/" {
			fs.listDirectoryHandler(w, r)
			return
		}
		if err == filer_pb.ErrNotFound {
			glog.V(1).Infof("Not found %s: %v", path, err)
			stats.FilerRequestCounter.WithLabelValues("read.notfound").Inc()
			w.WriteHeader(http.StatusNotFound)
		} else {
			glog.Errorf("Internal %s: %v", path, err)
			stats.FilerRequestCounter.WithLabelValues("read.internalerror").Inc()
			w.WriteHeader(http.StatusInternalServerError)
		}
		return
	}

	if entry.IsDirectory() {
		if fs.option.DisableDirListing {
			w.WriteHeader(http.StatusMethodNotAllowed)
			return
		}
		fs.listDirectoryHandler(w, r)
		return
	}

	if isForDirectory {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	// set etag
	etag := filer.ETagEntry(entry)
	if ifm := r.Header.Get("If-Match"); ifm != "" && (ifm != "\""+etag+"\"" && ifm != etag) {
		w.WriteHeader(http.StatusPreconditionFailed)
		return
	}

	w.Header().Set("Accept-Ranges", "bytes")

	// mime type
	mimeType := entry.Attr.Mime
	if mimeType == "" {
		if ext := filepath.Ext(entry.Name()); ext != "" {
			mimeType = mime.TypeByExtension(ext)
		}
	}
	if mimeType != "" {
		w.Header().Set("Content-Type", mimeType)
	}

	// if modified since
	if !entry.Attr.Mtime.IsZero() {
		w.Header().Set("Last-Modified", entry.Attr.Mtime.UTC().Format(http.TimeFormat))
		if r.Header.Get("If-Modified-Since") != "" {
			if t, parseError := time.Parse(http.TimeFormat, r.Header.Get("If-Modified-Since")); parseError == nil {
				if !t.Before(entry.Attr.Mtime) {
					w.WriteHeader(http.StatusNotModified)
					return
				}
			}
		}
	}

	// print out the header from extended properties
	for k, v := range entry.Extended {
		if !strings.HasPrefix(k, "xattr-") {
			// "xattr-" prefix is set in filesys.XATTR_PREFIX
			w.Header().Set(k, string(v))
		}
	}

	//Seaweed custom header are not visible to Vue or javascript
	seaweedHeaders := []string{}
	for header := range w.Header() {
		if strings.HasPrefix(header, "Seaweed-") {
			seaweedHeaders = append(seaweedHeaders, header)
		}
	}
	seaweedHeaders = append(seaweedHeaders, "Content-Disposition")
	w.Header().Set("Access-Control-Expose-Headers", strings.Join(seaweedHeaders, ","))

	//set tag count
	if r.Method == "GET" {
		tagCount := 0
		for k := range entry.Extended {
			if strings.HasPrefix(k, xhttp.AmzObjectTagging+"-") {
				tagCount++
			}
		}
		if tagCount > 0 {
			w.Header().Set(xhttp.AmzTagCount, strconv.Itoa(tagCount))
		}
	}

	if inm := r.Header.Get("If-None-Match"); inm == "\""+etag+"\"" {
		w.WriteHeader(http.StatusNotModified)
		return
	}
	setEtag(w, etag)

	filename := entry.Name()
	filename = url.QueryEscape(filename)
	adjustHeaderContentDisposition(w, r, filename)

	totalSize := int64(entry.Size())

	if r.Method == "HEAD" {
		w.Header().Set("Content-Length", strconv.FormatInt(totalSize, 10))
		return
	}

	if rangeReq := r.Header.Get("Range"); rangeReq == "" {
		ext := filepath.Ext(filename)
		if len(ext) > 0 {
			ext = strings.ToLower(ext)
		}
		width, height, mode, shouldResize := shouldResizeImages(ext, r)
		if shouldResize {
			data, err := filer.ReadAll(fs.filer.MasterClient, entry.Chunks)
			if err != nil {
				glog.Errorf("failed to read %s: %v", path, err)
				w.WriteHeader(http.StatusNotModified)
				return
			}
			rs, _, _ := images.Resized(ext, bytes.NewReader(data), width, height, mode)
			io.Copy(w, rs)
			return
		}
	}

	processRangeRequest(r, w, totalSize, mimeType, func(writer io.Writer, offset int64, size int64) error {
		if offset+size <= int64(len(entry.Content)) {
			_, err := writer.Write(entry.Content[offset : offset+size])
			if err != nil {
				glog.Errorf("failed to write entry content: %v", err)
			}
			return err
		}
		chunks := entry.Chunks
		if entry.IsInRemoteOnly() {
			dir, name := entry.FullPath.DirAndName()
			if resp, err := fs.DownloadToLocal(context.Background(), &filer_pb.DownloadToLocalRequest{
				Directory: dir,
				Name:      name,
			}); err != nil {
				return fmt.Errorf("cache %s: %v", entry.FullPath, err)
			} else {
				chunks = resp.Entry.Chunks
			}
		}

		err = filer.StreamContent(fs.filer.MasterClient, writer, chunks, offset, size)
		if err != nil {
			glog.Errorf("failed to stream content %s: %v", r.URL, err)
		}
		return err
	})
}
<commit_msg>add cache error log<commit_after>package weed_server

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"mime"
	"net/http"
	"net/url"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/images"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	xhttp "github.com/chrislusf/seaweedfs/weed/s3api/http"
	"github.com/chrislusf/seaweedfs/weed/stats"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func (fs *FilerServer) GetOrHeadHandler(w http.ResponseWriter, r *http.Request) {

	path := r.URL.Path
	isForDirectory := strings.HasSuffix(path, "/")
	if isForDirectory && len(path) > 1 {
		path = path[:len(path)-1]
	}

	entry, err := fs.filer.FindEntry(context.Background(), util.FullPath(path))
	if err != nil {
		if path == "/" {
			fs.listDirectoryHandler(w, r)
			return
		}
		if err == filer_pb.ErrNotFound {
			glog.V(1).Infof("Not found %s: %v", path, err)
			stats.FilerRequestCounter.WithLabelValues("read.notfound").Inc()
			w.WriteHeader(http.StatusNotFound)
		} else {
			glog.Errorf("Internal %s: %v", path, err)
			stats.FilerRequestCounter.WithLabelValues("read.internalerror").Inc()
			w.WriteHeader(http.StatusInternalServerError)
		}
		return
	}

	if entry.IsDirectory() {
		if fs.option.DisableDirListing {
			w.WriteHeader(http.StatusMethodNotAllowed)
			return
		}
		fs.listDirectoryHandler(w, r)
		return
	}

	if isForDirectory {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	// set etag
	etag := filer.ETagEntry(entry)
	if ifm := r.Header.Get("If-Match"); ifm != "" && (ifm != "\""+etag+"\"" && ifm != etag) {
		w.WriteHeader(http.StatusPreconditionFailed)
		return
	}

	w.Header().Set("Accept-Ranges", "bytes")

	// mime type
	mimeType := entry.Attr.Mime
	if mimeType == "" {
		if ext := filepath.Ext(entry.Name()); ext != "" {
			mimeType = mime.TypeByExtension(ext)
		}
	}
	if mimeType != "" {
		w.Header().Set("Content-Type", mimeType)
	}

	// if modified since
	if !entry.Attr.Mtime.IsZero() {
		w.Header().Set("Last-Modified", entry.Attr.Mtime.UTC().Format(http.TimeFormat))
		if r.Header.Get("If-Modified-Since") != "" {
			if t, parseError := time.Parse(http.TimeFormat, r.Header.Get("If-Modified-Since")); parseError == nil {
				if !t.Before(entry.Attr.Mtime) {
					w.WriteHeader(http.StatusNotModified)
					return
				}
			}
		}
	}

	// print out the header from extended properties
	for k, v := range entry.Extended {
		if !strings.HasPrefix(k, "xattr-") {
			// "xattr-" prefix is set in filesys.XATTR_PREFIX
			w.Header().Set(k, string(v))
		}
	}

	//Seaweed custom header are not visible to Vue or javascript
	seaweedHeaders := []string{}
	for header := range w.Header() {
		if strings.HasPrefix(header, "Seaweed-") {
			seaweedHeaders = append(seaweedHeaders, header)
		}
	}
	seaweedHeaders = append(seaweedHeaders, "Content-Disposition")
	w.Header().Set("Access-Control-Expose-Headers", strings.Join(seaweedHeaders, ","))

	//set tag count
	if r.Method == "GET" {
		tagCount := 0
		for k := range entry.Extended {
			if strings.HasPrefix(k, xhttp.AmzObjectTagging+"-") {
				tagCount++
			}
		}
		if tagCount > 0 {
			w.Header().Set(xhttp.AmzTagCount, strconv.Itoa(tagCount))
		}
	}

	if inm := r.Header.Get("If-None-Match"); inm == "\""+etag+"\"" {
		w.WriteHeader(http.StatusNotModified)
		return
	}
	setEtag(w, etag)

	filename := entry.Name()
	filename = url.QueryEscape(filename)
	adjustHeaderContentDisposition(w, r, filename)

	totalSize := int64(entry.Size())

	if r.Method == "HEAD" {
		w.Header().Set("Content-Length", strconv.FormatInt(totalSize, 10))
		return
	}

	if rangeReq := r.Header.Get("Range"); rangeReq == "" {
		ext := filepath.Ext(filename)
		if len(ext) > 0 {
			ext = strings.ToLower(ext)
		}
		width, height, mode, shouldResize := shouldResizeImages(ext, r)
		if shouldResize {
			data, err := filer.ReadAll(fs.filer.MasterClient, entry.Chunks)
			if err != nil {
				glog.Errorf("failed to read %s: %v", path, err)
				w.WriteHeader(http.StatusNotModified)
				return
			}
			rs, _, _ := images.Resized(ext, bytes.NewReader(data), width, height, mode)
			io.Copy(w, rs)
			return
		}
	}

	processRangeRequest(r, w, totalSize, mimeType, func(writer io.Writer, offset int64, size int64) error {
		if offset+size <= int64(len(entry.Content)) {
			_, err := writer.Write(entry.Content[offset : offset+size])
			if err != nil {
				glog.Errorf("failed to write entry content: %v", err)
			}
			return err
		}
		chunks := entry.Chunks
		if entry.IsInRemoteOnly() {
			dir, name := entry.FullPath.DirAndName()
			if resp, err := fs.DownloadToLocal(context.Background(), &filer_pb.DownloadToLocalRequest{
				Directory: dir,
				Name:      name,
			}); err != nil {
				glog.Errorf("cache %s: %v", entry.FullPath, err)
				return fmt.Errorf("cache %s: %v", entry.FullPath, err)
			} else {
				chunks = resp.Entry.Chunks
			}
		}

		err = filer.StreamContent(fs.filer.MasterClient, writer, chunks, offset, size)
		if err != nil {
			glog.Errorf("failed to stream content %s: %v", r.URL, err)
		}
		return err
	})
}
<|endoftext|>


<commit_before>package go_action_sdk

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestDefaultResponse_TellResponseReturnsExpectedOutput(t *testing.T) {
	builder := NewResponseBuilder()
	cToken := "42"
	message := "This Is A Test"
	response := builder.TellResponse(message, &cToken)
	assert.False(t, response.ExpectUserResponse)
	assert.Nil(t, response.ExpectedInputs)
	assert.Equal(t, &cToken, response.ConversationToken)
	assert.Equal(t, &message, response.FinalResponse_.SpeechResponse_.TextToSpeech)
	assert.Nil(t, response.FinalResponse_.SpeechResponse_.SSML)
}

func TestDefaultResponse_TellResponseSSMLReturnsExpectedOutput(t *testing.T) {
	builder := NewResponseBuilder()
	cToken := "42"
	message := "<s>This Is A Test</s>"
	response := builder.TellResponseSSML(message, &cToken)
	assert.False(t, response.ExpectUserResponse)
	assert.Nil(t, response.ExpectedInputs)
	assert.Equal(t, &cToken, response.ConversationToken)
	assert.Equal(t, &message, response.FinalResponse_.SpeechResponse_.SSML)
	assert.Nil(t, response.FinalResponse_.SpeechResponse_.TextToSpeech)
}

func TestDefaultResponse_AskResponseReturnsExpectedOutput(t *testing.T) {
	// nil noInputPrompt check
	builder := NewResponseBuilder()
	cToken := "42"
	message := "This Is A Test"
	response := builder.AskResponse(message, &cToken, nil)
	assert.True(t, response.ExpectUserResponse)
	assert.Nil(t, response.FinalResponse_)
	assert.NotNil(t, response.ExpectedInputs)
	assert.Len(t, response.ExpectedInputs, 1)
	assert.Len(t, response.ExpectedInputs[0].PossibleIntents, 1)
	assert.Equal(t, "assistant.intent.action.TEXT", response.ExpectedInputs[0].PossibleIntents[0].Intent)
	assert.Nil(t, response.ExpectedInputs[0].PossibleIntents[0].InputValueSpec_)
	assert.Len(t, response.ExpectedInputs[0].InputPrompt_.InitialPrompts, 1)
	assert.Nil(t, response.ExpectedInputs[0].InputPrompt_.InitialPrompts[0].SSML)
	assert.Equal(t, &message, response.ExpectedInputs[0].InputPrompt_.InitialPrompts[0].TextToSpeech)
	assert.Nil(t, response.ExpectedInputs[0].InputPrompt_.NoInputPrompts)
}

func TestDefaultResponse_AskResponseSSMLReturnsExpectedOutput(t *testing.T) {
	// nil noInputPrompt check
	builder := NewResponseBuilder()
	cToken := "42"
	message := "This Is A Test"
	response := builder.AskResponseSSML(message, &cToken, nil)
	assert.True(t, response.ExpectUserResponse)
	assert.Nil(t, response.FinalResponse_)
	assert.NotNil(t, response.ExpectedInputs)
	assert.Len(t, response.ExpectedInputs, 1)
	assert.Len(t, response.ExpectedInputs[0].PossibleIntents, 1)
	assert.Equal(t, "assistant.intent.action.TEXT", response.ExpectedInputs[0].PossibleIntents[0].Intent)
	assert.Nil(t, response.ExpectedInputs[0].PossibleIntents[0].InputValueSpec_)
	assert.Len(t, response.ExpectedInputs[0].InputPrompt_.InitialPrompts, 1)
	assert.Nil(t, response.ExpectedInputs[0].InputPrompt_.InitialPrompts[0].TextToSpeech)
	assert.Equal(t, &message, response.ExpectedInputs[0].InputPrompt_.InitialPrompts[0].SSML)
	assert.Nil(t, response.ExpectedInputs[0].InputPrompt_.NoInputPrompts)
}
<commit_msg>Last unit test cases added<commit_after>package go_action_sdk

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestDefaultResponse_TellResponseReturnsExpectedOutput(t *testing.T) {
	builder := NewResponseBuilder()
	cToken := "42"
	message := "This Is A Test"
	response := builder.TellResponse(message, &cToken)
	assert.False(t, response.ExpectUserResponse)
	assert.Nil(t, response.ExpectedInputs)
	assert.Equal(t, &cToken, response.ConversationToken)
	assert.Equal(t, &message, response.FinalResponse_.SpeechResponse_.TextToSpeech)
	assert.Nil(t, response.FinalResponse_.SpeechResponse_.SSML)
}

func TestDefaultResponse_TellResponseSSMLReturnsExpectedOutput(t *testing.T) {
	builder := NewResponseBuilder()
	cToken := "42"
	message := "<s>This Is A Test</s>"
	response := builder.TellResponseSSML(message, &cToken)
	assert.False(t, response.ExpectUserResponse)
	assert.Nil(t, response.ExpectedInputs)
	assert.Equal(t, &cToken, response.ConversationToken)
	assert.Equal(t, &message, response.FinalResponse_.SpeechResponse_.SSML)
	assert.Nil(t, response.FinalResponse_.SpeechResponse_.TextToSpeech)
}

func TestDefaultResponse_AskResponseReturnsExpectedOutput(t *testing.T) {
	// nil noInputPrompt check
	builder := NewResponseBuilder()
	cToken := "42"
	message := "This Is A Test"
	response := builder.AskResponse(message, &cToken, nil)
	assert.True(t, response.ExpectUserResponse)
	assert.Nil(t, response.FinalResponse_)
	assert.NotNil(t, response.ExpectedInputs)
	assert.Len(t, response.ExpectedInputs, 1)
	assert.Len(t, response.ExpectedInputs[0].PossibleIntents, 1)
	assert.Equal(t, "assistant.intent.action.TEXT", response.ExpectedInputs[0].PossibleIntents[0].Intent)
	assert.Nil(t, response.ExpectedInputs[0].PossibleIntents[0].InputValueSpec_)
	assert.Len(t, response.ExpectedInputs[0].InputPrompt_.InitialPrompts, 1)
	assert.Nil(t, response.ExpectedInputs[0].InputPrompt_.InitialPrompts[0].SSML)
	assert.Equal(t, &message, response.ExpectedInputs[0].InputPrompt_.InitialPrompts[0].TextToSpeech)
	assert.Nil(t, response.ExpectedInputs[0].InputPrompt_.NoInputPrompts)

	// actual noInputPrompt check
	cToken = "42"
	message = "This Is A Test"
	noInputPromptStr := "Nevermind"
	noInputPromptSlice := make([]string, 1)
	noInputPromptSlice[0] = noInputPromptStr
	response = builder.AskResponse(message, &cToken, noInputPromptSlice)
	assert.True(t, response.ExpectUserResponse)
	assert.Nil(t, response.FinalResponse_)
	assert.NotNil(t, response.ExpectedInputs)
	assert.Len(t, response.ExpectedInputs, 1)
	assert.Len(t, response.ExpectedInputs[0].PossibleIntents, 1)
	assert.Equal(t, "assistant.intent.action.TEXT", response.ExpectedInputs[0].PossibleIntents[0].Intent)
	assert.Nil(t, response.ExpectedInputs[0].PossibleIntents[0].InputValueSpec_)
	assert.Len(t, response.ExpectedInputs[0].InputPrompt_.InitialPrompts, 1)
	assert.Nil(t, response.ExpectedInputs[0].InputPrompt_.InitialPrompts[0].SSML)
	assert.Equal(t, &message, response.ExpectedInputs[0].InputPrompt_.InitialPrompts[0].TextToSpeech)
	assert.Len(t, response.ExpectedInputs[0].InputPrompt_.NoInputPrompts, 1)
	assert.Equal(t, &noInputPromptStr, response.ExpectedInputs[0].InputPrompt_.NoInputPrompts[0].TextToSpeech)
	assert.Nil(t, response.ExpectedInputs[0].InputPrompt_.NoInputPrompts[0].SSML)
}

func TestDefaultResponse_AskResponseSSMLReturnsExpectedOutput(t *testing.T) {
	// nil noInputPrompt check
	builder := NewResponseBuilder()
	cToken := "42"
	message := "This Is A Test"
	response := builder.AskResponseSSML(message, &cToken, nil)
	assert.True(t, response.ExpectUserResponse)
	assert.Nil(t, response.FinalResponse_)
	assert.NotNil(t, response.ExpectedInputs)
	assert.Len(t, response.ExpectedInputs, 1)
	assert.Len(t, response.ExpectedInputs[0].PossibleIntents, 1)
	assert.Equal(t, "assistant.intent.action.TEXT", response.ExpectedInputs[0].PossibleIntents[0].Intent)
	assert.Nil(t, response.ExpectedInputs[0].PossibleIntents[0].InputValueSpec_)
	assert.Len(t, response.ExpectedInputs[0].InputPrompt_.InitialPrompts, 1)
	assert.Nil(t, response.ExpectedInputs[0].InputPrompt_.InitialPrompts[0].TextToSpeech)
	assert.Equal(t, &message, response.ExpectedInputs[0].InputPrompt_.InitialPrompts[0].SSML)
	assert.Nil(t, response.ExpectedInputs[0].InputPrompt_.NoInputPrompts)

	// actual noInputPrompt check
	cToken = "42"
	message = "This Is A Test"
	noInputPromptStr := "Nevermind"
	noInputPromptSlice := make([]string, 1)
	noInputPromptSlice[0] = noInputPromptStr
	response = builder.AskResponseSSML(message, &cToken, noInputPromptSlice)
	assert.True(t, response.ExpectUserResponse)
	assert.Nil(t, response.FinalResponse_)
	assert.NotNil(t, response.ExpectedInputs)
	assert.Len(t, response.ExpectedInputs, 1)
	assert.Len(t, response.ExpectedInputs[0].PossibleIntents, 1)
	assert.Equal(t, "assistant.intent.action.TEXT", response.ExpectedInputs[0].PossibleIntents[0].Intent)
	assert.Nil(t, response.ExpectedInputs[0].PossibleIntents[0].InputValueSpec_)
	assert.Len(t, response.ExpectedInputs[0].InputPrompt_.InitialPrompts, 1)
	assert.Nil(t, response.ExpectedInputs[0].InputPrompt_.InitialPrompts[0].TextToSpeech)
	assert.Equal(t, &message, response.ExpectedInputs[0].InputPrompt_.InitialPrompts[0].SSML)
	assert.Len(t, response.ExpectedInputs[0].InputPrompt_.NoInputPrompts, 1)
	assert.Equal(t, &noInputPromptStr, response.ExpectedInputs[0].InputPrompt_.NoInputPrompts[0].SSML)
	assert.Nil(t, response.ExpectedInputs[0].InputPrompt_.NoInputPrompts[0].TextToSpeech)
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"flag"
	"fmt"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/golang/glog"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	v3 "google.golang.org/api/monitoring/v3"

	"github.com/GoogleCloudPlatform/k8s-stackdriver/prometheus-to-sd/config"
	"github.com/GoogleCloudPlatform/k8s-stackdriver/prometheus-to-sd/flags"
	"github.com/GoogleCloudPlatform/k8s-stackdriver/prometheus-to-sd/translator"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var (
	metricsPrefix = flag.String("stackdriver-prefix", "container.googleapis.com/master",
		"Prefix that is appended to every metric. Could be rewritten by metricsPrefix in per "+
			"component configuration.")
	autoWhitelistMetrics = flag.Bool("auto-whitelist-metrics", false,
		"If component has no whitelisted metrics, prometheus-to-sd will fetch them from Stackdriver.")
	metricDescriptorsResolution = flag.Duration("metric-descriptors-resolution", 10*time.Minute,
		"The resolution at which prometheus-to-sd will scrape metric descriptors from Stackdriver.")
	apioverride = flag.String("api-override", "",
		"The stackdriver API endpoint to override the default one used (which is prod).")
	source = flags.Uris{}
	podId  = flag.String("pod-id", "machine",
		"Name of the pod in which monitored component is running.")
	namespaceId = flag.String("namespace-id", "",
		"Namespace name of the pod in which monitored component is running.")
	zoneOverride = flag.String("zone-override", "",
		"Name of the zone to override the default one (in which component is running).")
	monitoredResourceTypes = flag.String("monitored-resource-types", "gke_container",
		"The monitored resource types to use, either the legacy 'gke_container', or the new 'k8s'")
	omitComponentName = flag.Bool("omit-component-name", true,
		"If metric name starts with the component name then this substring is removed to keep metric name shorter.")
	debugPort      = flag.Uint("port", 6061, "Port on which debug information is exposed.")
	dynamicSources = flags.Uris{}
	scrapeInterval = flag.Duration("scrape-interval", 60*time.Second,
		"The interval between metric scrapes. If there are multiple scrapes between two exports, the last present value is exported, even when missing from last scraping.")
	exportInterval = flag.Duration("export-interval", 60*time.Second,
		"The interval between metric exports. Can't be lower than --scrape-interval.")
	downcaseMetricNames = flag.Bool("downcase-metric-names", false,
		"If enabled, will downcase all metric names.")
	delayedShutdownTimeout = flag.Duration("delayed-shutdown-timeout", 120*time.Second,
		"Time to wait for the shutdown after receiving SIGTERM. 0 value means shutdown immediately, negative value results in ignoring signal." +
		" Default value is 120 seconds.")
)

func main() {
	flag.Set("logtostderr", "true")
	flag.Var(&source, "source", "source(s) to watch in [component-name]:http://host:port/path?whitelisted=a,b,c&podIdLabel=d&namespaceIdLabel=e&containerNameLabel=f&metricsPrefix=prefix format. Can be specified multiple times")
	flag.Var(&dynamicSources, "dynamic-source",
		`dynamic source(s) to watch in format: "[component-name]:http://:port/path?whitelisted=metric1,metric2&podIdLabel=label1&namespaceIdLabel=label2&containerNameLabel=label3&metricsPrefix=prefix". Dynamic sources are components (on the same node) discovered dynamically using the kubernetes api.`,
	)

	defer glog.Flush()
	flag.Parse()

	if *delayedShutdownTimeout < 0 {
		signal.Ignore(syscall.SIGTERM)
	} else {
		sigTermChannel := make(chan os.Signal, 1)
		signal.Notify(sigTermChannel, syscall.SIGTERM)

		go func() {
			<- sigTermChannel
			glog.Infof("SIGTERM has been received, Waiting %s before the shutdown.", delayedShutdownTimeout.String())

			time.Sleep(*delayedShutdownTimeout)
			glog.Info("Shutting down after receiving SIGTERM.")
			os.Exit(0)
		} ()
	}

	gceConf, err := config.GetGceConfig(*zoneOverride, *monitoredResourceTypes)
	if err != nil {
		glog.Fatalf("Failed to get GCE config: %v", err)
	}
	glog.Infof("GCE config: %+v", gceConf)

	sourceConfigs := getSourceConfigs(*metricsPrefix, gceConf)
	glog.Infof("Built the following source configs: %v", sourceConfigs)

	go func() {
		http.Handle("/metrics", promhttp.Handler())
		glog.Error(http.ListenAndServe(fmt.Sprintf(":%d", *debugPort), nil))
	}()

	client := oauth2.NewClient(context.Background(), google.ComputeTokenSource(""))
	stackdriverService, err := v3.New(client)
	if *apioverride != "" {
		stackdriverService.BasePath = *apioverride
	}
	if err != nil {
		glog.Fatalf("Failed to create Stackdriver client: %v", err)
	}
	glog.V(4).Infof("Successfully created Stackdriver client")

	if len(sourceConfigs) == 0 {
		glog.Fatalf("No sources defined. Please specify at least one --source flag.")
	}

	if *scrapeInterval > *exportInterval {
		glog.Fatalf("--scrape-interval cannot be bigger than --export-interval")
	}

	for _, sourceConfig := range sourceConfigs {
		glog.V(4).Infof("Starting goroutine for %+v", sourceConfig)

		// Pass sourceConfig as a parameter to avoid using the last sourceConfig by all goroutines.
		go readAndPushDataToStackdriver(stackdriverService, gceConf, sourceConfig)
	}

	// As worker goroutines work forever, block main thread as well.
	<-make(chan int)
}

func getSourceConfigs(defaultMetricsPrefix string, gceConfig *config.GceConfig) []*config.SourceConfig {
	glog.Info("Taking source configs from flags")
	staticSourceConfigs := config.SourceConfigsFromFlags(source, podId, namespaceId, defaultMetricsPrefix)
	glog.Info("Taking source configs from kubernetes api server")
	dynamicSourceConfigs, err := config.SourceConfigsFromDynamicSources(gceConfig, []flags.Uri(dynamicSources))
	if err != nil {
		glog.Fatalf(err.Error())
	}
	return append(staticSourceConfigs, dynamicSourceConfigs...)
}

func readAndPushDataToStackdriver(stackdriverService *v3.Service, gceConf *config.GceConfig, sourceConfig *config.SourceConfig) {
	glog.Infof("Running prometheus-to-sd, monitored target is %s %v:%v", sourceConfig.Component, sourceConfig.Host, sourceConfig.Port)
	commonConfig := &config.CommonConfig{
		GceConfig:           gceConf,
		SourceConfig:        sourceConfig,
		OmitComponentName:   *omitComponentName,
		DowncaseMetricNames: *downcaseMetricNames,
	}
	metricDescriptorCache := translator.NewMetricDescriptorCache(stackdriverService, commonConfig)
	signal := time.After(0)
	useWhitelistedMetricsAutodiscovery := *autoWhitelistMetrics && len(sourceConfig.Whitelisted) == 0
	timeSeriesBuilder := translator.NewTimeSeriesBuilder(commonConfig, metricDescriptorCache)
	exportTicker := time.Tick(*exportInterval)

	for range time.Tick(*scrapeInterval) {
		// Possibly exporting as a first thing, since errors down the
		// road will jump to next iteration of the loop.
		select {
		case <-exportTicker:
			ts, err := timeSeriesBuilder.Build()
			if err != nil {
				glog.Errorf("Could not build time series for component %v: %v", sourceConfig.Component, err)
			} else {
				translator.SendToStackdriver(stackdriverService, commonConfig, ts)
			}
		default:
		}

		// Mark cache at the beginning of each iteration as stale. Cache is considered refreshed only if during
		// current iteration there was successful call to Refresh function.
		metricDescriptorCache.MarkStale()
		glog.V(4).Infof("Scraping metrics of component %v", sourceConfig.Component)
		select {
		case <-signal:
			glog.V(4).Infof("Updating metrics cache for component %v", sourceConfig.Component)
			metricDescriptorCache.Refresh()
			if useWhitelistedMetricsAutodiscovery {
				sourceConfig.UpdateWhitelistedMetrics(metricDescriptorCache.GetMetricNames())
				glog.V(2).Infof("Autodiscovered whitelisted metrics for component %v: %v", sourceConfig.Component, sourceConfig.Whitelisted)
			}
			signal = time.After(*metricDescriptorsResolution)
		default:
		}
		if useWhitelistedMetricsAutodiscovery && len(sourceConfig.Whitelisted) == 0 {
			glog.V(4).Infof("Skipping %v component as there are no metric to expose.", sourceConfig.Component)
			continue
		}
		metrics, err := translator.GetPrometheusMetrics(sourceConfig)
		if err != nil {
			glog.V(2).Infof("Error while getting Prometheus metrics %v for component %v", err, sourceConfig.Component)
			continue
		}
		timeSeriesBuilder.Update(metrics, time.Now())
	}
}
<commit_msg>reformatted prometheus-to-sd/main.go file<commit_after>/*
Copyright 2017 Google Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"flag"
	"fmt"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/golang/glog"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	v3 "google.golang.org/api/monitoring/v3"

	"github.com/GoogleCloudPlatform/k8s-stackdriver/prometheus-to-sd/config"
	"github.com/GoogleCloudPlatform/k8s-stackdriver/prometheus-to-sd/flags"
	"github.com/GoogleCloudPlatform/k8s-stackdriver/prometheus-to-sd/translator"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var (
	metricsPrefix = flag.String("stackdriver-prefix", "container.googleapis.com/master",
		"Prefix that is appended to every metric. Could be rewritten by metricsPrefix in per "+
				"component configuration.")
	autoWhitelistMetrics = flag.Bool("auto-whitelist-metrics", false,
		"If component has no whitelisted metrics, prometheus-to-sd will fetch them from Stackdriver.")
	metricDescriptorsResolution = flag.Duration("metric-descriptors-resolution", 10*time.Minute,
		"The resolution at which prometheus-to-sd will scrape metric descriptors from Stackdriver.")
	apioverride = flag.String("api-override", "",
		"The stackdriver API endpoint to override the default one used (which is prod).")
	source = flags.Uris{}
	podId  = flag.String("pod-id", "machine",
		"Name of the pod in which monitored component is running.")
	namespaceId = flag.String("namespace-id", "",
		"Namespace name of the pod in which monitored component is running.")
	zoneOverride = flag.String("zone-override", "",
		"Name of the zone to override the default one (in which component is running).")
	monitoredResourceTypes = flag.String("monitored-resource-types", "gke_container",
		"The monitored resource types to use, either the legacy 'gke_container', or the new 'k8s'")
	omitComponentName = flag.Bool("omit-component-name", true,
		"If metric name starts with the component name then this substring is removed to keep metric name shorter.")
	debugPort      = flag.Uint("port", 6061, "Port on which debug information is exposed.")
	dynamicSources = flags.Uris{}
	scrapeInterval = flag.Duration("scrape-interval", 60*time.Second,
		"The interval between metric scrapes. If there are multiple scrapes between two exports, the last present value is exported, even when missing from last scraping.")
	exportInterval = flag.Duration("export-interval", 60*time.Second,
		"The interval between metric exports. Can't be lower than --scrape-interval.")
	downcaseMetricNames = flag.Bool("downcase-metric-names", false,
		"If enabled, will downcase all metric names.")
	delayedShutdownTimeout = flag.Duration("delayed-shutdown-timeout", 120*time.Second,
		"Time to wait for the shutdown after receiving SIGTERM. 0 value means shutdown immediately, negative value results in ignoring signal."+
				" Default value is 120 seconds.")
)

func main() {
	flag.Set("logtostderr", "true")
	flag.Var(&source, "source", "source(s) to watch in [component-name]:http://host:port/path?whitelisted=a,b,c&podIdLabel=d&namespaceIdLabel=e&containerNameLabel=f&metricsPrefix=prefix format. Can be specified multiple times")
	flag.Var(&dynamicSources, "dynamic-source",
		`dynamic source(s) to watch in format: "[component-name]:http://:port/path?whitelisted=metric1,metric2&podIdLabel=label1&namespaceIdLabel=label2&containerNameLabel=label3&metricsPrefix=prefix". Dynamic sources are components (on the same node) discovered dynamically using the kubernetes api.`,
	)

	defer glog.Flush()
	flag.Parse()

	if *delayedShutdownTimeout < 0 {
		signal.Ignore(syscall.SIGTERM)
	} else {
		sigTermChannel := make(chan os.Signal, 1)
		signal.Notify(sigTermChannel, syscall.SIGTERM)

		go func() {
			<-sigTermChannel
			glog.Infof("SIGTERM has been received, Waiting %s before the shutdown.", delayedShutdownTimeout.String())

			time.Sleep(*delayedShutdownTimeout)
			glog.Info("Shutting down after receiving SIGTERM.")
			os.Exit(0)
		}()
	}

	gceConf, err := config.GetGceConfig(*zoneOverride, *monitoredResourceTypes)
	if err != nil {
		glog.Fatalf("Failed to get GCE config: %v", err)
	}
	glog.Infof("GCE config: %+v", gceConf)

	sourceConfigs := getSourceConfigs(*metricsPrefix, gceConf)
	glog.Infof("Built the following source configs: %v", sourceConfigs)

	go func() {
		http.Handle("/metrics", promhttp.Handler())
		glog.Error(http.ListenAndServe(fmt.Sprintf(":%d", *debugPort), nil))
	}()

	client := oauth2.NewClient(context.Background(), google.ComputeTokenSource(""))
	stackdriverService, err := v3.New(client)
	if *apioverride != "" {
		stackdriverService.BasePath = *apioverride
	}
	if err != nil {
		glog.Fatalf("Failed to create Stackdriver client: %v", err)
	}
	glog.V(4).Infof("Successfully created Stackdriver client")

	if len(sourceConfigs) == 0 {
		glog.Fatalf("No sources defined. Please specify at least one --source flag.")
	}

	if *scrapeInterval > *exportInterval {
		glog.Fatalf("--scrape-interval cannot be bigger than --export-interval")
	}

	for _, sourceConfig := range sourceConfigs {
		glog.V(4).Infof("Starting goroutine for %+v", sourceConfig)

		// Pass sourceConfig as a parameter to avoid using the last sourceConfig by all goroutines.
		go readAndPushDataToStackdriver(stackdriverService, gceConf, sourceConfig)
	}

	// As worker goroutines work forever, block main thread as well.
	<-make(chan int)
}

func getSourceConfigs(defaultMetricsPrefix string, gceConfig *config.GceConfig) []*config.SourceConfig {
	glog.Info("Taking source configs from flags")
	staticSourceConfigs := config.SourceConfigsFromFlags(source, podId, namespaceId, defaultMetricsPrefix)
	glog.Info("Taking source configs from kubernetes api server")
	dynamicSourceConfigs, err := config.SourceConfigsFromDynamicSources(gceConfig, []flags.Uri(dynamicSources))
	if err != nil {
		glog.Fatalf(err.Error())
	}
	return append(staticSourceConfigs, dynamicSourceConfigs...)
}

func readAndPushDataToStackdriver(stackdriverService *v3.Service, gceConf *config.GceConfig, sourceConfig *config.SourceConfig) {
	glog.Infof("Running prometheus-to-sd, monitored target is %s %v:%v", sourceConfig.Component, sourceConfig.Host, sourceConfig.Port)
	commonConfig := &config.CommonConfig{
		GceConfig:           gceConf,
		SourceConfig:        sourceConfig,
		OmitComponentName:   *omitComponentName,
		DowncaseMetricNames: *downcaseMetricNames,
	}
	metricDescriptorCache := translator.NewMetricDescriptorCache(stackdriverService, commonConfig)
	signal := time.After(0)
	useWhitelistedMetricsAutodiscovery := *autoWhitelistMetrics && len(sourceConfig.Whitelisted) == 0
	timeSeriesBuilder := translator.NewTimeSeriesBuilder(commonConfig, metricDescriptorCache)
	exportTicker := time.Tick(*exportInterval)

	for range time.Tick(*scrapeInterval) {
		// Possibly exporting as a first thing, since errors down the
		// road will jump to next iteration of the loop.
		select {
		case <-exportTicker:
			ts, err := timeSeriesBuilder.Build()
			if err != nil {
				glog.Errorf("Could not build time series for component %v: %v", sourceConfig.Component, err)
			} else {
				translator.SendToStackdriver(stackdriverService, commonConfig, ts)
			}
		default:
		}

		// Mark cache at the beginning of each iteration as stale. Cache is considered refreshed only if during
		// current iteration there was successful call to Refresh function.
		metricDescriptorCache.MarkStale()
		glog.V(4).Infof("Scraping metrics of component %v", sourceConfig.Component)
		select {
		case <-signal:
			glog.V(4).Infof("Updating metrics cache for component %v", sourceConfig.Component)
			metricDescriptorCache.Refresh()
			if useWhitelistedMetricsAutodiscovery {
				sourceConfig.UpdateWhitelistedMetrics(metricDescriptorCache.GetMetricNames())
				glog.V(2).Infof("Autodiscovered whitelisted metrics for component %v: %v", sourceConfig.Component, sourceConfig.Whitelisted)
			}
			signal = time.After(*metricDescriptorsResolution)
		default:
		}
		if useWhitelistedMetricsAutodiscovery && len(sourceConfig.Whitelisted) == 0 {
			glog.V(4).Infof("Skipping %v component as there are no metric to expose.", sourceConfig.Component)
			continue
		}
		metrics, err := translator.GetPrometheusMetrics(sourceConfig)
		if err != nil {
			glog.V(2).Infof("Error while getting Prometheus metrics %v for component %v", err, sourceConfig.Component)
			continue
		}
		timeSeriesBuilder.Update(metrics, time.Now())
	}
}
<|endoftext|>


<commit_before>package app

import (
	"github.com/gin-gonic/gin"
	"net/http"
)

func Creator(c *gin.Context) {
	if r, ok := c.Get("resource"); ok {
		ri := r.(ResourceIface)

		c.Bind(ri.Form())
		ri.Save()

		c.Header("Location", ri.Url())
		// c.JSON(http.StatusCreated, "")
		c.Data(http.StatusCreated, gin.MIMEJSON, nil)
	}
}

func Getter(c *gin.Context)     {}
func Updater(c *gin.Context)    {}
func ListGetter(c *gin.Context) {}
func Destroyer(c *gin.Context)  {}
<commit_msg>remove comment from Creator<commit_after>package app

import (
	"github.com/gin-gonic/gin"
	"net/http"
)

func Creator(c *gin.Context) {
	if r, ok := c.Get("resource"); ok {
		ri := r.(ResourceIface)

		c.Bind(ri.Form())
		ri.Save()

		c.Header("Location", ri.Url())
		c.Data(http.StatusCreated, gin.MIMEJSON, nil)
	}
}

func Getter(c *gin.Context)     {}
func Updater(c *gin.Context)    {}
func ListGetter(c *gin.Context) {}
func Destroyer(c *gin.Context)  {}
<|endoftext|>


<commit_before>// 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.

// Vet is a simple checker for static errors in Go source code.
// See doc.go for more information.
package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

var verbose = flag.Bool("v", false, "verbose")
var exitCode = 0

// Flags to control which checks to perform
var (
	vetAll             = flag.Bool("all", true, "check everything; disabled if any explicit check is requested")
	vetAtomic          = flag.Bool("atomic", false, "check for common mistaken usages of the sync/atomic package")
	vetBuildTags       = flag.Bool("buildtags", false, "check that +build tags are valid")
	vetMethods         = flag.Bool("methods", false, "check that canonically named methods are canonically defined")
	vetPrintf          = flag.Bool("printf", false, "check printf-like invocations")
	vetStructTags      = flag.Bool("structtags", false, "check that struct field tags have canonical format")
	vetRangeLoops      = flag.Bool("rangeloops", false, "check that range loop variables are used correctly")
	vetUntaggedLiteral = flag.Bool("composites", false, "check that composite literals used type-tagged elements")
)

// setExit sets the value for os.Exit when it is called, later.  It
// remembers the highest value.
func setExit(err int) {
	if err > exitCode {
		exitCode = err
	}
}

// Usage is a replacement usage function for the flags package.
func Usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	flag.PrintDefaults()
	os.Exit(2)
}

// File is a wrapper for the state of a file used in the parser.
// The parse tree walkers are all methods of this type.
type File struct {
	fset *token.FileSet
	file *ast.File
	b    bytes.Buffer // for use by methods
}

func main() {
	flag.Usage = Usage
	flag.Parse()

	// If a check is named explicitly, turn off the 'all' flag.
	if *vetMethods || *vetPrintf || *vetStructTags || *vetUntaggedLiteral || *vetRangeLoops {
		*vetAll = false
	}

	if *printfuncs != "" {
		for _, name := range strings.Split(*printfuncs, ",") {
			if len(name) == 0 {
				flag.Usage()
			}
			skip := 0
			if colon := strings.LastIndex(name, ":"); colon > 0 {
				var err error
				skip, err = strconv.Atoi(name[colon+1:])
				if err != nil {
					errorf(`illegal format for "Func:N" argument %q; %s`, name, err)
				}
				name = name[:colon]
			}
			name = strings.ToLower(name)
			if name[len(name)-1] == 'f' {
				printfList[name] = skip
			} else {
				printList[name] = skip
			}
		}
	}

	if flag.NArg() == 0 {
		doFile("stdin", os.Stdin)
	} else {
		for _, name := range flag.Args() {
			// Is it a directory?
			if fi, err := os.Stat(name); err == nil && fi.IsDir() {
				walkDir(name)
			} else {
				doFile(name, nil)
			}
		}
	}
	os.Exit(exitCode)
}

// doFile analyzes one file.  If the reader is nil, the source code is read from the
// named file.
func doFile(name string, reader io.Reader) {
	if reader == nil {
		f, err := os.Open(name)
		if err != nil {
			errorf("%s: %s", name, err)
			return
		}
		defer f.Close()
		reader = f
	}
	data, err := ioutil.ReadAll(reader)
	if err != nil {
		errorf("%s: %s", name, err)
		return
	}
	checkBuildTag(name, data)
	fs := token.NewFileSet()
	parsedFile, err := parser.ParseFile(fs, name, bytes.NewReader(data), 0)
	if err != nil {
		errorf("%s: %s", name, err)
		return
	}
	file := &File{fset: fs, file: parsedFile}
	file.walkFile(name, parsedFile)
}

func visit(path string, f os.FileInfo, err error) error {
	if err != nil {
		errorf("walk error: %s", err)
		return nil
	}
	if !f.IsDir() && strings.HasSuffix(path, ".go") {
		doFile(path, nil)
	}
	return nil
}

// walkDir recursively walks the tree looking for .go files.
func walkDir(root string) {
	filepath.Walk(root, visit)
}

// error formats the error to standard error, adding program
// identification and a newline
func errorf(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "vet: "+format+"\n", args...)
	setExit(2)
}

// Println is fmt.Println guarded by -v.
func Println(args ...interface{}) {
	if !*verbose {
		return
	}
	fmt.Println(args...)
}

// Printf is fmt.Printf guarded by -v.
func Printf(format string, args ...interface{}) {
	if !*verbose {
		return
	}
	fmt.Printf(format+"\n", args...)
}

// Bad reports an error and sets the exit code..
func (f *File) Bad(pos token.Pos, args ...interface{}) {
	f.Warn(pos, args...)
	setExit(1)
}

// Badf reports a formatted error and sets the exit code.
func (f *File) Badf(pos token.Pos, format string, args ...interface{}) {
	f.Warnf(pos, format, args...)
	setExit(1)
}

// Warn reports an error but does not set the exit code.
func (f *File) Warn(pos token.Pos, args ...interface{}) {
	loc := f.fset.Position(pos).String() + ": "
	fmt.Fprint(os.Stderr, loc+fmt.Sprintln(args...))
}

// Warnf reports a formatted error but does not set the exit code.
func (f *File) Warnf(pos token.Pos, format string, args ...interface{}) {
	loc := f.fset.Position(pos).String() + ": "
	fmt.Fprintf(os.Stderr, loc+format+"\n", args...)
}

// walkFile walks the file's tree.
func (f *File) walkFile(name string, file *ast.File) {
	Println("Checking file", name)
	ast.Walk(f, file)
}

// Visit implements the ast.Visitor interface.
func (f *File) Visit(node ast.Node) ast.Visitor {
	switch n := node.(type) {
	case *ast.AssignStmt:
		f.walkAssignStmt(n)
	case *ast.CallExpr:
		f.walkCallExpr(n)
	case *ast.CompositeLit:
		f.walkCompositeLit(n)
	case *ast.Field:
		f.walkFieldTag(n)
	case *ast.FuncDecl:
		f.walkMethodDecl(n)
	case *ast.InterfaceType:
		f.walkInterfaceType(n)
	case *ast.RangeStmt:
		f.walkRangeStmt(n)
	}
	return f
}

// walkCall walks an assignment statement
func (f *File) walkAssignStmt(stmt *ast.AssignStmt) {
	f.checkAtomicAssignment(stmt)
}

// walkCall walks a call expression.
func (f *File) walkCall(call *ast.CallExpr, name string) {
	f.checkFmtPrintfCall(call, name)
}

// walkCallExpr walks a call expression.
func (f *File) walkCallExpr(call *ast.CallExpr) {
	switch x := call.Fun.(type) {
	case *ast.Ident:
		f.walkCall(call, x.Name)
	case *ast.SelectorExpr:
		f.walkCall(call, x.Sel.Name)
	}
}

// walkCompositeLit walks a composite literal.
func (f *File) walkCompositeLit(c *ast.CompositeLit) {
	f.checkUntaggedLiteral(c)
}

// walkFieldTag walks a struct field tag.
func (f *File) walkFieldTag(field *ast.Field) {
	if field.Tag == nil {
		return
	}
	f.checkCanonicalFieldTag(field)
}

// walkMethodDecl walks the method's signature.
func (f *File) walkMethod(id *ast.Ident, t *ast.FuncType) {
	f.checkCanonicalMethod(id, t)
}

// walkMethodDecl walks the method signature in the declaration.
func (f *File) walkMethodDecl(d *ast.FuncDecl) {
	if d.Recv == nil {
		// not a method
		return
	}
	f.walkMethod(d.Name, d.Type)
}

// walkInterfaceType walks the method signatures of an interface.
func (f *File) walkInterfaceType(t *ast.InterfaceType) {
	for _, field := range t.Methods.List {
		for _, id := range field.Names {
			f.walkMethod(id, field.Type.(*ast.FuncType))
		}
	}
}

// walkRangeStmt walks a range statement.
func (f *File) walkRangeStmt(n *ast.RangeStmt) {
	checkRangeLoop(f, n)
}

// goFmt returns a string representation of the expression
func (f *File) gofmt(x ast.Expr) string {
	f.b.Reset()
	printer.Fprint(&f.b, f.fset, x)
	return f.b.String()
}
<commit_msg>cmd/vet: add missing -all logic<commit_after>// 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.

// Vet is a simple checker for static errors in Go source code.
// See doc.go for more information.
package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

var verbose = flag.Bool("v", false, "verbose")
var exitCode = 0

// Flags to control which checks to perform.
// NOTE: Add new flags to the if statement at the top of func main too.
var (
	vetAll             = flag.Bool("all", true, "check everything; disabled if any explicit check is requested")
	vetAtomic          = flag.Bool("atomic", false, "check for common mistaken usages of the sync/atomic package")
	vetBuildTags       = flag.Bool("buildtags", false, "check that +build tags are valid")
	vetMethods         = flag.Bool("methods", false, "check that canonically named methods are canonically defined")
	vetPrintf          = flag.Bool("printf", false, "check printf-like invocations")
	vetStructTags      = flag.Bool("structtags", false, "check that struct field tags have canonical format")
	vetRangeLoops      = flag.Bool("rangeloops", false, "check that range loop variables are used correctly")
	vetUntaggedLiteral = flag.Bool("composites", false, "check that composite literals used type-tagged elements")
)

// setExit sets the value for os.Exit when it is called, later.  It
// remembers the highest value.
func setExit(err int) {
	if err > exitCode {
		exitCode = err
	}
}

// Usage is a replacement usage function for the flags package.
func Usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	flag.PrintDefaults()
	os.Exit(2)
}

// File is a wrapper for the state of a file used in the parser.
// The parse tree walkers are all methods of this type.
type File struct {
	fset *token.FileSet
	file *ast.File
	b    bytes.Buffer // for use by methods
}

func main() {
	flag.Usage = Usage
	flag.Parse()

	// If a check is named explicitly, turn off the 'all' flag.
	if *vetAtomic || *vetBuildTags || *vetMethods || *vetPrintf || *vetStructTags || *vetRangeLoops || *vetUntaggedLiteral {
		*vetAll = false
	}

	if *printfuncs != "" {
		for _, name := range strings.Split(*printfuncs, ",") {
			if len(name) == 0 {
				flag.Usage()
			}
			skip := 0
			if colon := strings.LastIndex(name, ":"); colon > 0 {
				var err error
				skip, err = strconv.Atoi(name[colon+1:])
				if err != nil {
					errorf(`illegal format for "Func:N" argument %q; %s`, name, err)
				}
				name = name[:colon]
			}
			name = strings.ToLower(name)
			if name[len(name)-1] == 'f' {
				printfList[name] = skip
			} else {
				printList[name] = skip
			}
		}
	}

	if flag.NArg() == 0 {
		doFile("stdin", os.Stdin)
	} else {
		for _, name := range flag.Args() {
			// Is it a directory?
			if fi, err := os.Stat(name); err == nil && fi.IsDir() {
				walkDir(name)
			} else {
				doFile(name, nil)
			}
		}
	}
	os.Exit(exitCode)
}

// doFile analyzes one file.  If the reader is nil, the source code is read from the
// named file.
func doFile(name string, reader io.Reader) {
	if reader == nil {
		f, err := os.Open(name)
		if err != nil {
			errorf("%s: %s", name, err)
			return
		}
		defer f.Close()
		reader = f
	}
	data, err := ioutil.ReadAll(reader)
	if err != nil {
		errorf("%s: %s", name, err)
		return
	}
	checkBuildTag(name, data)
	fs := token.NewFileSet()
	parsedFile, err := parser.ParseFile(fs, name, bytes.NewReader(data), 0)
	if err != nil {
		errorf("%s: %s", name, err)
		return
	}
	file := &File{fset: fs, file: parsedFile}
	file.walkFile(name, parsedFile)
}

func visit(path string, f os.FileInfo, err error) error {
	if err != nil {
		errorf("walk error: %s", err)
		return nil
	}
	if !f.IsDir() && strings.HasSuffix(path, ".go") {
		doFile(path, nil)
	}
	return nil
}

// walkDir recursively walks the tree looking for .go files.
func walkDir(root string) {
	filepath.Walk(root, visit)
}

// error formats the error to standard error, adding program
// identification and a newline
func errorf(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "vet: "+format+"\n", args...)
	setExit(2)
}

// Println is fmt.Println guarded by -v.
func Println(args ...interface{}) {
	if !*verbose {
		return
	}
	fmt.Println(args...)
}

// Printf is fmt.Printf guarded by -v.
func Printf(format string, args ...interface{}) {
	if !*verbose {
		return
	}
	fmt.Printf(format+"\n", args...)
}

// Bad reports an error and sets the exit code..
func (f *File) Bad(pos token.Pos, args ...interface{}) {
	f.Warn(pos, args...)
	setExit(1)
}

// Badf reports a formatted error and sets the exit code.
func (f *File) Badf(pos token.Pos, format string, args ...interface{}) {
	f.Warnf(pos, format, args...)
	setExit(1)
}

// Warn reports an error but does not set the exit code.
func (f *File) Warn(pos token.Pos, args ...interface{}) {
	loc := f.fset.Position(pos).String() + ": "
	fmt.Fprint(os.Stderr, loc+fmt.Sprintln(args...))
}

// Warnf reports a formatted error but does not set the exit code.
func (f *File) Warnf(pos token.Pos, format string, args ...interface{}) {
	loc := f.fset.Position(pos).String() + ": "
	fmt.Fprintf(os.Stderr, loc+format+"\n", args...)
}

// walkFile walks the file's tree.
func (f *File) walkFile(name string, file *ast.File) {
	Println("Checking file", name)
	ast.Walk(f, file)
}

// Visit implements the ast.Visitor interface.
func (f *File) Visit(node ast.Node) ast.Visitor {
	switch n := node.(type) {
	case *ast.AssignStmt:
		f.walkAssignStmt(n)
	case *ast.CallExpr:
		f.walkCallExpr(n)
	case *ast.CompositeLit:
		f.walkCompositeLit(n)
	case *ast.Field:
		f.walkFieldTag(n)
	case *ast.FuncDecl:
		f.walkMethodDecl(n)
	case *ast.InterfaceType:
		f.walkInterfaceType(n)
	case *ast.RangeStmt:
		f.walkRangeStmt(n)
	}
	return f
}

// walkCall walks an assignment statement
func (f *File) walkAssignStmt(stmt *ast.AssignStmt) {
	f.checkAtomicAssignment(stmt)
}

// walkCall walks a call expression.
func (f *File) walkCall(call *ast.CallExpr, name string) {
	f.checkFmtPrintfCall(call, name)
}

// walkCallExpr walks a call expression.
func (f *File) walkCallExpr(call *ast.CallExpr) {
	switch x := call.Fun.(type) {
	case *ast.Ident:
		f.walkCall(call, x.Name)
	case *ast.SelectorExpr:
		f.walkCall(call, x.Sel.Name)
	}
}

// walkCompositeLit walks a composite literal.
func (f *File) walkCompositeLit(c *ast.CompositeLit) {
	f.checkUntaggedLiteral(c)
}

// walkFieldTag walks a struct field tag.
func (f *File) walkFieldTag(field *ast.Field) {
	if field.Tag == nil {
		return
	}
	f.checkCanonicalFieldTag(field)
}

// walkMethodDecl walks the method's signature.
func (f *File) walkMethod(id *ast.Ident, t *ast.FuncType) {
	f.checkCanonicalMethod(id, t)
}

// walkMethodDecl walks the method signature in the declaration.
func (f *File) walkMethodDecl(d *ast.FuncDecl) {
	if d.Recv == nil {
		// not a method
		return
	}
	f.walkMethod(d.Name, d.Type)
}

// walkInterfaceType walks the method signatures of an interface.
func (f *File) walkInterfaceType(t *ast.InterfaceType) {
	for _, field := range t.Methods.List {
		for _, id := range field.Names {
			f.walkMethod(id, field.Type.(*ast.FuncType))
		}
	}
}

// walkRangeStmt walks a range statement.
func (f *File) walkRangeStmt(n *ast.RangeStmt) {
	checkRangeLoop(f, n)
}

// goFmt returns a string representation of the expression
func (f *File) gofmt(x ast.Expr) string {
	f.b.Reset()
	printer.Fprint(&f.b, f.fset, x)
	return f.b.String()
}
<|endoftext|>


<commit_before>// This file is subject to a 1-clause BSD license.
// Its Datas can be found in the enclosed LICENSE file.

package proto

import (
	"io"
	"strconv"
	"strings"
)

// Message is a parsed incoming message.
type Message struct {
	SenderName string
	SenderMask string
	Receiver   string
	Data       string
	Command    uint16
}

// InChannel returns true if this message came from a channel context
// instead of a user or service.
func (m *Message) InChannel() bool {
	if len(m.Receiver) == 0 {
		return false
	}

	c := m.Receiver[0]
	return c == '#' && c == '&' && c == '!' && c == '+'
}

// parseMessage parses a message from the given data.
func parseMessage(data string) (m *Message, err error) {
	if len(data) == 0 {
		return nil, io.EOF
	}

	m = new(Message)
	m.Command = Unknown
	m.Data = data

	list := strings.Split(data, " ")

	switch list[0] {
	case "PING":
		m.Command = CmdPing
		m.Data = list[1][1:]
		return

	case "ERROR":
		m.Command = CmdError
		m.Data = list[1][1:]
		return
	}

	if len(list) < 3 {
		return
	}

	m.SenderMask = list[0][1:]
	m.Command = findType(list[1])
	m.Receiver = list[2]
	m.Data = strings.Join(list[3:], " ")

	idx := strings.Index(m.SenderMask, "!")
	if idx > -1 {
		m.SenderName = m.SenderMask[:idx]
		m.SenderMask = m.SenderMask[idx+1:]
	}

	if len(m.Data) > 0 && m.Data[0] == ':' {
		m.Data = m.Data[1:]
	}

	if !m.InChannel() {
		m.Receiver = m.SenderName
	}

	return
}

// findType attempts to parse a protocol ID from the input string.
// These come as 3-digit numbers or a string. For example: "001" or "NOTICE"
func findType(v string) uint16 {
	n, err := strconv.ParseUint(v, 10, 16)
	if err == nil {
		return uint16(n)
	}

	v = strings.ToUpper(v)

	switch v {
	case "ADMIN":
		return CmdAdmin
	case "AWAY":
		return CmdAway
	case "CONNECT":
		return CmdConnect
	case "DIE":
		return CmdDie
	case "ERROR":
		return CmdError
	case "INFO":
		return CmdInfo
	case "INVITE":
		return CmdInvite
	case "ISON":
		return CmdIsOn
	case "JOIN":
		return CmdJoin
	case "KICK":
		return CmdKick
	case "KILL":
		return CmdKill
	case "LINKS":
		return CmdLinks
	case "LIST":
		return CmdList
	case "LUSERS":
		return CmdLUsers
	case "MODE":
		return CmdMode
	case "MOTD":
		return CmdMOTD
	case "NAMES":
		return CmdNames
	case "NICK":
		return CmdNick
	case "NJOIN":
		return CmdNJoin
	case "NOTICE":
		return CmdNotice
	case "OPER":
		return CmdOper
	case "PART":
		return CmdPart
	case "PASS":
		return CmdPass
	case "PING":
		return CmdPing
	case "PONG":
		return CmdPong
	case "PRIVMSG":
		return CmdPrivMsg
	case "QUIT":
		return CmdQuit
	case "REHASH":
		return CmdRehash
	case "RESTART":
		return CmdRestart
	case "SERVER":
		return CmdServer
	case "SERVICE":
		return CmdService
	case "SERVLIST":
		return CmdServList
	case "SQUERY":
		return CmdSQuery
	case "SQUIRT":
		return CmdSquirt
	case "SQUIT":
		return CmdSQuit
	case "STATS":
		return CmdStats
	case "SUMMON":
		return CmdSummon
	case "TIME":
		return CmdTime
	case "TOPIC":
		return CmdTopic
	case "TRACE":
		return CmdTrace
	case "USER":
		return CmdUser
	case "USERHOST":
		return CmdUserHost
	case "USERS":
		return CmdUsers
	case "VERSION":
		return CmdVersion
	case "WALLOPS":
		return CmdWAllOps
	case "WHO":
		return CmdWho
	case "WHOIS":
		return CmdWhoIs
	case "WHOWAS":
		return CmdWhoWas
	}

	return Unknown
}
<commit_msg>Fixes Message.InChannel.<commit_after>// This file is subject to a 1-clause BSD license.
// Its Datas can be found in the enclosed LICENSE file.

package proto

import (
	"io"
	"strconv"
	"strings"
)

// Message is a parsed incoming message.
type Message struct {
	SenderName string
	SenderMask string
	Receiver   string
	Data       string
	Command    uint16
}

// InChannel returns true if this message came from a channel context
// instead of a user or service.
func (m *Message) InChannel() bool {
	if len(m.Receiver) == 0 {
		return false
	}

	c := m.Receiver[0]
	return c == '#' || c == '&' || c == '!' || c == '+'
}

// parseMessage parses a message from the given data.
func parseMessage(data string) (m *Message, err error) {
	if len(data) == 0 {
		return nil, io.EOF
	}

	m = new(Message)
	m.Command = Unknown
	m.Data = data

	list := strings.Split(data, " ")

	switch list[0] {
	case "PING":
		m.Command = CmdPing
		m.Data = list[1][1:]
		return

	case "ERROR":
		m.Command = CmdError
		m.Data = list[1][1:]
		return
	}

	if len(list) < 3 {
		return
	}

	m.SenderMask = list[0][1:]
	m.Command = findType(list[1])
	m.Receiver = list[2]
	m.Data = strings.Join(list[3:], " ")

	idx := strings.Index(m.SenderMask, "!")
	if idx > -1 {
		m.SenderName = m.SenderMask[:idx]
		m.SenderMask = m.SenderMask[idx+1:]
	}

	if len(m.Data) > 0 && m.Data[0] == ':' {
		m.Data = m.Data[1:]
	}

	if !m.InChannel() {
		m.Receiver = m.SenderName
	}

	return
}

// findType attempts to parse a protocol ID from the input string.
// These come as 3-digit numbers or a string. For example: "001" or "NOTICE"
func findType(v string) uint16 {
	n, err := strconv.ParseUint(v, 10, 16)
	if err == nil {
		return uint16(n)
	}

	v = strings.ToUpper(v)

	switch v {
	case "ADMIN":
		return CmdAdmin
	case "AWAY":
		return CmdAway
	case "CONNECT":
		return CmdConnect
	case "DIE":
		return CmdDie
	case "ERROR":
		return CmdError
	case "INFO":
		return CmdInfo
	case "INVITE":
		return CmdInvite
	case "ISON":
		return CmdIsOn
	case "JOIN":
		return CmdJoin
	case "KICK":
		return CmdKick
	case "KILL":
		return CmdKill
	case "LINKS":
		return CmdLinks
	case "LIST":
		return CmdList
	case "LUSERS":
		return CmdLUsers
	case "MODE":
		return CmdMode
	case "MOTD":
		return CmdMOTD
	case "NAMES":
		return CmdNames
	case "NICK":
		return CmdNick
	case "NJOIN":
		return CmdNJoin
	case "NOTICE":
		return CmdNotice
	case "OPER":
		return CmdOper
	case "PART":
		return CmdPart
	case "PASS":
		return CmdPass
	case "PING":
		return CmdPing
	case "PONG":
		return CmdPong
	case "PRIVMSG":
		return CmdPrivMsg
	case "QUIT":
		return CmdQuit
	case "REHASH":
		return CmdRehash
	case "RESTART":
		return CmdRestart
	case "SERVER":
		return CmdServer
	case "SERVICE":
		return CmdService
	case "SERVLIST":
		return CmdServList
	case "SQUERY":
		return CmdSQuery
	case "SQUIRT":
		return CmdSquirt
	case "SQUIT":
		return CmdSQuit
	case "STATS":
		return CmdStats
	case "SUMMON":
		return CmdSummon
	case "TIME":
		return CmdTime
	case "TOPIC":
		return CmdTopic
	case "TRACE":
		return CmdTrace
	case "USER":
		return CmdUser
	case "USERHOST":
		return CmdUserHost
	case "USERS":
		return CmdUsers
	case "VERSION":
		return CmdVersion
	case "WALLOPS":
		return CmdWAllOps
	case "WHO":
		return CmdWho
	case "WHOIS":
		return CmdWhoIs
	case "WHOWAS":
		return CmdWhoWas
	}

	return Unknown
}
<|endoftext|>


<commit_before>package protocol
import ("log"
    "bytes"
    "text/template"
    "github.com/amdonov/xmlsig"
    "encoding/base64"
    "bufio"
    "net/http"
    "encoding/xml")

func NewPOSTResponseMarshaller(signer xmlsig.Signer) ResponseMarshaller {
    generator := &postResponseMarshaller{signer:signer}
    generator.template = template.New("postResponse")
    generator.template.Parse(`<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.1//EN"
"http://www.w3.org/TR/xhtml11/DTD/xhtml11.dtd">

Note: Since your browser does not support JavaScript,
you must press the Continue button once to proceed.

<form action="{{ .AssertionConsumerServiceURL }}"
method="post">

<input type="hidden" name="RelayState"
value="{{ .RelayState }}"/>
<input type="hidden" name="SAMLResponse"
value="{{ .SAMLResponse }}"/>

<input type="submit" value="Continue"/>

</form>

`)
    return generator
}

type postResponseMarshaller struct {
    template *template.Template
    signer xmlsig.Signer
}

func (gen *postResponseMarshaller) Marshal(writer http.ResponseWriter, request *http.Request,
response *Response, authRequest *AuthnRequest, relayState string) {
    // Don't need to change the response. Go ahead and sign it
    signature, err := gen.signer.Sign(response.Assertion)
    if err!=nil {
        log.Println(err)
        return
    }
    response.Assertion.Signature = signature
    var xmlbuff bytes.Buffer
    memWriter := bufio.NewWriter(&xmlbuff)
    writer.Write([]byte(xml.Header))
    encoder := xml.NewEncoder(memWriter)
    encoder.Encode(response)
    memWriter.Flush()

    samlMessage := base64.StdEncoding.EncodeToString(xmlbuff.Bytes())
    postResponse := POSTResponse{relayState, samlMessage, authRequest.AssertionConsumerServiceURL}
    gen.template.Execute(writer, postResponse)
}

type POSTResponse struct {
    RelayState string
    SAMLResponse string
    AssertionConsumerServiceURL string
}
<commit_msg>Was writing XML declaration to the wrong writer<commit_after>package protocol
import ("log"
    "bytes"
    "text/template"
    "github.com/amdonov/xmlsig"
    "encoding/base64"
    "bufio"
    "net/http"
    "encoding/xml")

func NewPOSTResponseMarshaller(signer xmlsig.Signer) ResponseMarshaller {
    generator := &postResponseMarshaller{signer:signer}
    generator.template = template.New("postResponse")
    generator.template.Parse(`<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.1//EN"
"http://www.w3.org/TR/xhtml11/DTD/xhtml11.dtd">

Note: Since your browser does not support JavaScript,
you must press the Continue button once to proceed.

<form action="{{ .AssertionConsumerServiceURL }}"
method="post">

<input type="hidden" name="RelayState"
value="{{ .RelayState }}"/>
<input type="hidden" name="SAMLResponse"
value="{{ .SAMLResponse }}"/>

<input type="submit" value="Continue"/>

</form>

`)
    return generator
}

type postResponseMarshaller struct {
    template *template.Template
    signer xmlsig.Signer
}

func (gen *postResponseMarshaller) Marshal(writer http.ResponseWriter, request *http.Request,
response *Response, authRequest *AuthnRequest, relayState string) {
    // Don't need to change the response. Go ahead and sign it
    signature, err := gen.signer.Sign(response.Assertion)
    if err!=nil {
        log.Println(err)
        return
    }
    response.Assertion.Signature = signature
    var xmlbuff bytes.Buffer
    memWriter := bufio.NewWriter(&xmlbuff)
    memWriter.Write([]byte(xml.Header))
    encoder := xml.NewEncoder(memWriter)
    encoder.Encode(response)
    memWriter.Flush()

    samlMessage := base64.StdEncoding.EncodeToString(xmlbuff.Bytes())
    postResponse := POSTResponse{relayState, samlMessage, authRequest.AssertionConsumerServiceURL}
    gen.template.Execute(writer, postResponse)
}

type POSTResponse struct {
    RelayState string
    SAMLResponse string
    AssertionConsumerServiceURL string
}
<|endoftext|>