<commit_before>package ytdl

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/rs/zerolog/log"
)

const youtubeBaseURL = "https://www.youtube.com/watch"
const youtubeEmbeddedBaseURL = "https://www.youtube.com/embed/"
const youtubeVideoEURL = "https://youtube.googleapis.com/v/"
const youtubeVideoInfoURL = "https://www.youtube.com/get_video_info"
const youtubeDateFormat = "2006-01-02"

// VideoInfo contains the info a youtube video
type VideoInfo struct {
	ID              string     // The video ID
	Title           string     // The video title
	Description     string     // The video description
	DatePublished   time.Time  // The date the video was published
	Formats         FormatList // Formats the video is available in
	DASHManifestURL string     // URI of the DASH manifest file
	HLSManifestURL  string     // URI of the HLS manifest file
	Keywords        []string   // List of keywords associated with the video
	Uploader        string     // Author of the video
	Song            string
	Artist          string
	Album           string
	Writers         string
	Duration        time.Duration // Duration of the video
	htmlPlayerFile  string
}

// GetVideoInfo fetches info from a url string, url object, or a url string
func GetVideoInfo(cx context.Context, value interface{}) (*VideoInfo, error) {
	return DefaultClient.GetVideoInfo(cx, value)
}

// GetVideoInfo fetches info from a url string, url object, or a url string
func (c *Client) GetVideoInfo(cx context.Context, value interface{}) (*VideoInfo, error) {
	switch t := value.(type) {
	case *url.URL:
		videoID := extractVideoID(t)
		if len(videoID) == 0 {
			return nil, fmt.Errorf("invalid youtube URL, no video id")
		}
		return c.GetVideoInfoFromID(cx, videoID)
	case string:
		if strings.HasPrefix(t, "https://") {
			uri, err := url.ParseRequestURI(t)
			if err != nil {
				return nil, err
			}
			return c.GetVideoInfo(cx, uri)
		}
		return c.GetVideoInfoFromID(cx, t)
	default:
		return nil, fmt.Errorf("Identifier type must be a string, *url.URL, or []byte")
	}
}

// GetVideoInfoFromShortURL fetches video info from a short youtube url
func extractVideoID(u *url.URL) string {
	switch u.Host {
	case "www.youtube.com", "youtube.com", "m.youtube.com":
		if u.Path == "/watch" {
			return u.Query().Get("v")
		}
		if strings.HasPrefix(u.Path, "/embed/") {
			return u.Path[7:]
		}
	case "youtu.be":
		if len(u.Path) > 1 {
			return u.Path[1:]
		}
	}
	return ""
}

// GetVideoInfoFromID fetches video info from a youtube video id
func (c *Client) GetVideoInfoFromID(cx context.Context, id string) (*VideoInfo, error) {
	body, err := c.httpGetAndCheckResponseReadBody(cx, youtubeBaseURL+"?v="+id)

	if err != nil {
		return nil, err
	}
	return c.getVideoInfoFromHTML(cx, id, body)
}

// GetDownloadURL gets the download url for a format
func (c *Client) GetDownloadURL(cx context.Context, info *VideoInfo, format *Format) (*url.URL, error) {
	return c.getDownloadURL(cx, format, info.htmlPlayerFile)
}

// GetThumbnailURL returns a url for the thumbnail image
// with the given quality
func (info *VideoInfo) GetThumbnailURL(quality ThumbnailQuality) *url.URL {
	u, _ := url.Parse(fmt.Sprintf("http://img.youtube.com/vi/%s/%s.jpg",
		info.ID, quality))
	return u
}

// Download is a convenience method to download a format to an io.Writer
func (c *Client) Download(cx context.Context, info *VideoInfo, format *Format, dest io.Writer) error {
	u, err := c.GetDownloadURL(cx, info, format)
	if err != nil {
		return err
	}

	resp, err := c.httpGetAndCheckResponse(cx, u.String())
	if err != nil {
		return err
	}

	defer resp.Body.Close()
	_, err = io.Copy(dest, resp.Body)
	return err
}

var (
	regexpPlayerConfig          = regexp.MustCompile("ytplayer.config = (.*?);ytplayer.web_player_context_config")
	regexpInitialData           = regexp.MustCompile(`\["ytInitialData"\] = (.+);`)
	regexpInitialPlayerResponse = regexp.MustCompile(`\["ytInitialPlayerResponse"\] = (.+);`)
)

func (c *Client) getVideoInfoFromHTML(cx context.Context, id string, html []byte) (*VideoInfo, error) {
	info := &VideoInfo{}

	if matches := regexpInitialData.FindSubmatch(html); len(matches) > 0 {
		data := initialData{}

		if err := json.Unmarshal(matches[1], &data); err != nil {
			return nil, err
		}

		contents := data.Contents.TwoColumnWatchNextResults.Results.Results.Contents

		if len(contents) >= 2 {
			infoRenderer := contents[1].VideoSecondaryInfoRenderer

			info.Description = infoRenderer.Description.String()
			rows := infoRenderer.MetadataRowContainer.MetadataRowContainerRenderer.Rows

			info.Artist = rows.Get("Artist")
			info.Album = rows.Get("Album")
			info.Song = rows.Get("Song")
			info.Writers = rows.Get("Writers")
		}
	}

	info.ID = id

	var jsonConfig playerConfig

	// match json in javascript
	if matches := regexpPlayerConfig.FindSubmatch(html); len(matches) > 1 {
		err := json.Unmarshal(matches[1], &jsonConfig)
		if err != nil {
			return nil, err
		}
		//glog.Errorf("jsonCOnfig %s\n\n", jsonConfig.Args.PlayerResponse)
	} else {
		log.Debug().Msg("Unable to extract json from default url, trying embedded url")

		info, err := c.getVideoInfoFromEmbedded(cx, id)
		if err != nil {
			return nil, err
		}
		query := url.Values{
			"video_id": []string{id},
			"eurl":     []string{youtubeVideoEURL + id},
		}

		if sts, ok := info["sts"].(float64); ok {
			query.Add("sts", strconv.Itoa(int(sts)))
		}
		body, err := c.httpGetAndCheckResponseReadBody(cx, youtubeVideoInfoURL+"?"+query.Encode())
		if err != nil {
			return nil, fmt.Errorf("Unable to read video info: %w", err)
		}

		query, err = url.ParseQuery(string(body))
		if err != nil {
			return nil, fmt.Errorf("Unable to parse video info data: %w", err)
		}

		for k, v := range query {
			switch k {
			case "errorcode":
				jsonConfig.Args.Errorcode = v[0]
			case "reason":
				jsonConfig.Args.Reason = v[0]
			case "status":
				jsonConfig.Args.Status = v[0]
			case "player_response":
				jsonConfig.Args.PlayerResponse = v[0]
			case "url_encoded_fmt_stream_map":
				jsonConfig.Args.URLEncodedFmtStreamMap = v[0]
			case "adaptive_fmts":
				jsonConfig.Args.AdaptiveFmts = v[0]
			case "dashmpd":
				jsonConfig.Args.Dashmpd = v[0]
			default:
				// log.Debug().Msgf("unknown query param: %v", k)
			}
		}
	}

	inf := jsonConfig.Args
	if inf.Status == "fail" {
		return nil, fmt.Errorf("Error %s:%s", inf.Errorcode, inf.Reason)
	}

	var formats FormatList
	c.addFormatsByQueryStrings(&formats, strings.NewReader(inf.URLEncodedFmtStreamMap), false)
	if inf.AdaptiveFmts != "" {
		c.addFormatsByQueryStrings(&formats, strings.NewReader(inf.AdaptiveFmts), true)
	}
	if inf.PlayerResponse != "" {
		response := &playerResponse{}

		if err := json.Unmarshal([]byte(inf.PlayerResponse), &response); err != nil {
			return nil, fmt.Errorf("Couldn't parse player response: %w", err)
		}
		info.DASHManifestURL = response.StreamingData.DashManifestUrl
		info.HLSManifestURL = response.StreamingData.HlsManifestUrl
		if response.PlayabilityStatus.Status != "OK" {
			return nil, fmt.Errorf("Unavailable because: %s", response.PlayabilityStatus.Reason)
		}

		c.addFormatsByInfos(&formats, response.StreamingData.Formats, false)
		c.addFormatsByInfos(&formats, response.StreamingData.AdaptiveFormats, true)
		if seconds := response.VideoDetails.LengthSeconds; seconds != "" {
			val, err := strconv.Atoi(seconds)
			if err == nil {
				info.Duration = time.Duration(val) * time.Second
			}
		}

		if date, err := time.Parse(youtubeDateFormat, response.Microformat.Renderer.PublishDate); err == nil {
			info.DatePublished = date
		} else {
			log.Debug().Msgf("Unable to parse date published %v", err)
		}

		info.Title = response.VideoDetails.Title
		info.Uploader = response.VideoDetails.Author
	} else {
		log.Debug().Msg("Unable to extract player response JSON")
	}

	info.htmlPlayerFile = jsonConfig.Assets.JS

	if len(formats) == 0 {
		log.Debug().Msgf("No formats found")
	}
	if dashManifest := inf.Dashmpd; dashManifest != "" {
		tokens, err := c.getSigTokens(cx, info.htmlPlayerFile)
		if err != nil {
			return nil, fmt.Errorf("Unable to extract signature tokens: %w", err)
		}
		regex := regexp.MustCompile("\\/s\\/([a-fA-F0-9\\.]+)")
		regexSub := regexp.MustCompile("([a-fA-F0-9\\.]+)")
		info.DASHManifestURL = regex.ReplaceAllStringFunc(dashManifest, func(str string) string {
			return "/signature/" + decipherTokens(tokens, regexSub.FindString(str))
		})
		dashFormats, err := c.getDashManifest(cx, info.DASHManifestURL)
		if err != nil {
			return nil, fmt.Errorf("Unable to extract dash manifest: %w", err)
		}
		for _, dashFormat := range dashFormats {
			added := false
			for j, format := range formats {
				if dashFormat.Itag == format.Itag {
					formats[j] = dashFormat
					added = true
					break
				}
			}
			if !added {
				formats = append(formats, dashFormat)
			}
		}
	}
	info.Formats = formats
	return info, nil
}

func (c *Client) addFormatsByInfos(formats *FormatList, infos []formatInfo, adaptive bool) {
	for _, info := range infos {
		if err := formats.addByInfo(info, adaptive); err != nil {
			log.Debug().Err(err)
		}
	}
}

func (c *Client) addFormatsByQueryStrings(formats *FormatList, rd io.Reader, adaptive bool) {
	r := bufio.NewReader(rd)
	for {
		line, err := r.ReadString(',')
		if err == io.EOF {
			break
		}
		if err := formats.addByQueryString(line[:len(line)-1], adaptive); err != nil {
			log.Debug().Err(err)
		}
	}
}

func (c *Client) getVideoInfoFromEmbedded(cx context.Context, id string) (map[string]interface{}, error) {
	var jsonConfig map[string]interface{}

	html, err := c.httpGetAndCheckResponseReadBody(cx, youtubeEmbeddedBaseURL+id)

	if err != nil {
		return nil, fmt.Errorf("Embedded url request returned %w", err)
	}

	//	re = regexp.MustCompile("\"sts\"\\s*:\\s*(\\d+)")
	re := regexp.MustCompile("yt.setConfig\\({'PLAYER_CONFIG': (.*?)}\\);")

	matches := re.FindSubmatch(html)
	if len(matches) < 2 {
		return nil, fmt.Errorf("Error extracting sts from embedded url response")
	}
	dec := json.NewDecoder(bytes.NewBuffer(matches[1]))
	err = dec.Decode(&jsonConfig)
	if err != nil {
		return nil, fmt.Errorf("Unable to extract json from embedded url: %w", err)
	}

	return jsonConfig, nil
}

func (c *Client) getDashManifest(cx context.Context, urlString string) (formats []*Format, err error) {

	resp, err := c.httpGetAndCheckResponse(cx, urlString)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	dec := xml.NewDecoder(resp.Body)
	var token xml.Token
	for ; err == nil; token, err = dec.Token() {
		if el, ok := token.(xml.StartElement); ok && el.Name.Local == "Representation" {
			var rep representation
			err = dec.DecodeElement(&rep, &el)
			if err != nil {
				break
			}
			if itag := getItag(rep.Itag); itag != nil {
				format := &Format{
					Itag:     *itag,
					url:      rep.URL,
					FromDASH: true,
				}
				if rep.Height != 0 {
					format.Itag.Resolution = strconv.Itoa(rep.Height) + "p"
				}
				formats = append(formats, format)
			} else {
				log.Debug().Msgf("No metadata found for itag: %v, skipping...", rep.Itag)
			}
		}
	}
	if err != io.EOF {
		return nil, err
	}
	return formats, nil
}
<commit_msg>Update ytplayer config regex to be less brittle to changes<commit_after>package ytdl

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/rs/zerolog/log"
)

const youtubeBaseURL = "https://www.youtube.com/watch"
const youtubeEmbeddedBaseURL = "https://www.youtube.com/embed/"
const youtubeVideoEURL = "https://youtube.googleapis.com/v/"
const youtubeVideoInfoURL = "https://www.youtube.com/get_video_info"
const youtubeDateFormat = "2006-01-02"

// VideoInfo contains the info a youtube video
type VideoInfo struct {
	ID              string     // The video ID
	Title           string     // The video title
	Description     string     // The video description
	DatePublished   time.Time  // The date the video was published
	Formats         FormatList // Formats the video is available in
	DASHManifestURL string     // URI of the DASH manifest file
	HLSManifestURL  string     // URI of the HLS manifest file
	Keywords        []string   // List of keywords associated with the video
	Uploader        string     // Author of the video
	Song            string
	Artist          string
	Album           string
	Writers         string
	Duration        time.Duration // Duration of the video
	htmlPlayerFile  string
}

// GetVideoInfo fetches info from a url string, url object, or a url string
func GetVideoInfo(cx context.Context, value interface{}) (*VideoInfo, error) {
	return DefaultClient.GetVideoInfo(cx, value)
}

// GetVideoInfo fetches info from a url string, url object, or a url string
func (c *Client) GetVideoInfo(cx context.Context, value interface{}) (*VideoInfo, error) {
	switch t := value.(type) {
	case *url.URL:
		videoID := extractVideoID(t)
		if len(videoID) == 0 {
			return nil, fmt.Errorf("invalid youtube URL, no video id")
		}
		return c.GetVideoInfoFromID(cx, videoID)
	case string:
		if strings.HasPrefix(t, "https://") {
			uri, err := url.ParseRequestURI(t)
			if err != nil {
				return nil, err
			}
			return c.GetVideoInfo(cx, uri)
		}
		return c.GetVideoInfoFromID(cx, t)
	default:
		return nil, fmt.Errorf("Identifier type must be a string, *url.URL, or []byte")
	}
}

// GetVideoInfoFromShortURL fetches video info from a short youtube url
func extractVideoID(u *url.URL) string {
	switch u.Host {
	case "www.youtube.com", "youtube.com", "m.youtube.com":
		if u.Path == "/watch" {
			return u.Query().Get("v")
		}
		if strings.HasPrefix(u.Path, "/embed/") {
			return u.Path[7:]
		}
	case "youtu.be":
		if len(u.Path) > 1 {
			return u.Path[1:]
		}
	}
	return ""
}

// GetVideoInfoFromID fetches video info from a youtube video id
func (c *Client) GetVideoInfoFromID(cx context.Context, id string) (*VideoInfo, error) {
	body, err := c.httpGetAndCheckResponseReadBody(cx, youtubeBaseURL+"?v="+id)

	if err != nil {
		return nil, err
	}
	return c.getVideoInfoFromHTML(cx, id, body)
}

// GetDownloadURL gets the download url for a format
func (c *Client) GetDownloadURL(cx context.Context, info *VideoInfo, format *Format) (*url.URL, error) {
	return c.getDownloadURL(cx, format, info.htmlPlayerFile)
}

// GetThumbnailURL returns a url for the thumbnail image
// with the given quality
func (info *VideoInfo) GetThumbnailURL(quality ThumbnailQuality) *url.URL {
	u, _ := url.Parse(fmt.Sprintf("http://img.youtube.com/vi/%s/%s.jpg",
		info.ID, quality))
	return u
}

// Download is a convenience method to download a format to an io.Writer
func (c *Client) Download(cx context.Context, info *VideoInfo, format *Format, dest io.Writer) error {
	u, err := c.GetDownloadURL(cx, info, format)
	if err != nil {
		return err
	}

	resp, err := c.httpGetAndCheckResponse(cx, u.String())
	if err != nil {
		return err
	}

	defer resp.Body.Close()
	_, err = io.Copy(dest, resp.Body)
	return err
}

var (
	regexpPlayerConfig          = regexp.MustCompile("ytplayer.config = (.*?);ytplayer.")
	regexpInitialData           = regexp.MustCompile(`\["ytInitialData"\] = (.+);`)
	regexpInitialPlayerResponse = regexp.MustCompile(`\["ytInitialPlayerResponse"\] = (.+);`)
)

func (c *Client) getVideoInfoFromHTML(cx context.Context, id string, html []byte) (*VideoInfo, error) {
	info := &VideoInfo{}

	if matches := regexpInitialData.FindSubmatch(html); len(matches) > 0 {
		data := initialData{}

		if err := json.Unmarshal(matches[1], &data); err != nil {
			return nil, err
		}

		contents := data.Contents.TwoColumnWatchNextResults.Results.Results.Contents

		if len(contents) >= 2 {
			infoRenderer := contents[1].VideoSecondaryInfoRenderer

			info.Description = infoRenderer.Description.String()
			rows := infoRenderer.MetadataRowContainer.MetadataRowContainerRenderer.Rows

			info.Artist = rows.Get("Artist")
			info.Album = rows.Get("Album")
			info.Song = rows.Get("Song")
			info.Writers = rows.Get("Writers")
		}
	}

	info.ID = id

	var jsonConfig playerConfig

	// match json in javascript
	if matches := regexpPlayerConfig.FindSubmatch(html); len(matches) > 1 {
		err := json.Unmarshal(matches[1], &jsonConfig)
		if err != nil {
			return nil, err
		}
		//glog.Errorf("jsonCOnfig %s\n\n", jsonConfig.Args.PlayerResponse)
	} else {
		log.Debug().Msg("Unable to extract json from default url, trying embedded url")

		info, err := c.getVideoInfoFromEmbedded(cx, id)
		if err != nil {
			return nil, err
		}
		query := url.Values{
			"video_id": []string{id},
			"eurl":     []string{youtubeVideoEURL + id},
		}

		if sts, ok := info["sts"].(float64); ok {
			query.Add("sts", strconv.Itoa(int(sts)))
		}
		body, err := c.httpGetAndCheckResponseReadBody(cx, youtubeVideoInfoURL+"?"+query.Encode())
		if err != nil {
			return nil, fmt.Errorf("Unable to read video info: %w", err)
		}

		query, err = url.ParseQuery(string(body))
		if err != nil {
			return nil, fmt.Errorf("Unable to parse video info data: %w", err)
		}

		for k, v := range query {
			switch k {
			case "errorcode":
				jsonConfig.Args.Errorcode = v[0]
			case "reason":
				jsonConfig.Args.Reason = v[0]
			case "status":
				jsonConfig.Args.Status = v[0]
			case "player_response":
				jsonConfig.Args.PlayerResponse = v[0]
			case "url_encoded_fmt_stream_map":
				jsonConfig.Args.URLEncodedFmtStreamMap = v[0]
			case "adaptive_fmts":
				jsonConfig.Args.AdaptiveFmts = v[0]
			case "dashmpd":
				jsonConfig.Args.Dashmpd = v[0]
			default:
				// log.Debug().Msgf("unknown query param: %v", k)
			}
		}
	}

	inf := jsonConfig.Args
	if inf.Status == "fail" {
		return nil, fmt.Errorf("Error %s:%s", inf.Errorcode, inf.Reason)
	}

	var formats FormatList
	c.addFormatsByQueryStrings(&formats, strings.NewReader(inf.URLEncodedFmtStreamMap), false)
	if inf.AdaptiveFmts != "" {
		c.addFormatsByQueryStrings(&formats, strings.NewReader(inf.AdaptiveFmts), true)
	}
	if inf.PlayerResponse != "" {
		response := &playerResponse{}

		if err := json.Unmarshal([]byte(inf.PlayerResponse), &response); err != nil {
			return nil, fmt.Errorf("Couldn't parse player response: %w", err)
		}
		info.DASHManifestURL = response.StreamingData.DashManifestUrl
		info.HLSManifestURL = response.StreamingData.HlsManifestUrl
		if response.PlayabilityStatus.Status != "OK" {
			return nil, fmt.Errorf("Unavailable because: %s", response.PlayabilityStatus.Reason)
		}

		c.addFormatsByInfos(&formats, response.StreamingData.Formats, false)
		c.addFormatsByInfos(&formats, response.StreamingData.AdaptiveFormats, true)
		if seconds := response.VideoDetails.LengthSeconds; seconds != "" {
			val, err := strconv.Atoi(seconds)
			if err == nil {
				info.Duration = time.Duration(val) * time.Second
			}
		}

		if date, err := time.Parse(youtubeDateFormat, response.Microformat.Renderer.PublishDate); err == nil {
			info.DatePublished = date
		} else {
			log.Debug().Msgf("Unable to parse date published %v", err)
		}

		info.Title = response.VideoDetails.Title
		info.Uploader = response.VideoDetails.Author
	} else {
		log.Debug().Msg("Unable to extract player response JSON")
	}

	info.htmlPlayerFile = jsonConfig.Assets.JS

	if len(formats) == 0 {
		log.Debug().Msgf("No formats found")
	}
	if dashManifest := inf.Dashmpd; dashManifest != "" {
		tokens, err := c.getSigTokens(cx, info.htmlPlayerFile)
		if err != nil {
			return nil, fmt.Errorf("Unable to extract signature tokens: %w", err)
		}
		regex := regexp.MustCompile("\\/s\\/([a-fA-F0-9\\.]+)")
		regexSub := regexp.MustCompile("([a-fA-F0-9\\.]+)")
		info.DASHManifestURL = regex.ReplaceAllStringFunc(dashManifest, func(str string) string {
			return "/signature/" + decipherTokens(tokens, regexSub.FindString(str))
		})
		dashFormats, err := c.getDashManifest(cx, info.DASHManifestURL)
		if err != nil {
			return nil, fmt.Errorf("Unable to extract dash manifest: %w", err)
		}
		for _, dashFormat := range dashFormats {
			added := false
			for j, format := range formats {
				if dashFormat.Itag == format.Itag {
					formats[j] = dashFormat
					added = true
					break
				}
			}
			if !added {
				formats = append(formats, dashFormat)
			}
		}
	}
	info.Formats = formats
	return info, nil
}

func (c *Client) addFormatsByInfos(formats *FormatList, infos []formatInfo, adaptive bool) {
	for _, info := range infos {
		if err := formats.addByInfo(info, adaptive); err != nil {
			log.Debug().Err(err)
		}
	}
}

func (c *Client) addFormatsByQueryStrings(formats *FormatList, rd io.Reader, adaptive bool) {
	r := bufio.NewReader(rd)
	for {
		line, err := r.ReadString(',')
		if err == io.EOF {
			break
		}
		if err := formats.addByQueryString(line[:len(line)-1], adaptive); err != nil {
			log.Debug().Err(err)
		}
	}
}

func (c *Client) getVideoInfoFromEmbedded(cx context.Context, id string) (map[string]interface{}, error) {
	var jsonConfig map[string]interface{}

	html, err := c.httpGetAndCheckResponseReadBody(cx, youtubeEmbeddedBaseURL+id)

	if err != nil {
		return nil, fmt.Errorf("Embedded url request returned %w", err)
	}

	//	re = regexp.MustCompile("\"sts\"\\s*:\\s*(\\d+)")
	re := regexp.MustCompile("yt.setConfig\\({'PLAYER_CONFIG': (.*?)}\\);")

	matches := re.FindSubmatch(html)
	if len(matches) < 2 {
		return nil, fmt.Errorf("Error extracting sts from embedded url response")
	}
	dec := json.NewDecoder(bytes.NewBuffer(matches[1]))
	err = dec.Decode(&jsonConfig)
	if err != nil {
		return nil, fmt.Errorf("Unable to extract json from embedded url: %w", err)
	}

	return jsonConfig, nil
}

func (c *Client) getDashManifest(cx context.Context, urlString string) (formats []*Format, err error) {

	resp, err := c.httpGetAndCheckResponse(cx, urlString)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	dec := xml.NewDecoder(resp.Body)
	var token xml.Token
	for ; err == nil; token, err = dec.Token() {
		if el, ok := token.(xml.StartElement); ok && el.Name.Local == "Representation" {
			var rep representation
			err = dec.DecodeElement(&rep, &el)
			if err != nil {
				break
			}
			if itag := getItag(rep.Itag); itag != nil {
				format := &Format{
					Itag:     *itag,
					url:      rep.URL,
					FromDASH: true,
				}
				if rep.Height != 0 {
					format.Itag.Resolution = strconv.Itoa(rep.Height) + "p"
				}
				formats = append(formats, format)
			} else {
				log.Debug().Msgf("No metadata found for itag: %v, skipping...", rep.Itag)
			}
		}
	}
	if err != io.EOF {
		return nil, err
	}
	return formats, nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"crypto/rsa"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"time"

	"github.com/docker/docker/pkg/term"
	log "github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh"
	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/compute/v1"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/storage/v1"
)

const pollingInterval = 500 * time.Millisecond
const timeout = 300

// GCPClient contains state required for communication with GCP
type GCPClient struct {
	client      *http.Client
	compute     *compute.Service
	storage     *storage.Service
	projectName string
	fileName    string
	privKey     *rsa.PrivateKey
}

// NewGCPClient creates a new GCP client
func NewGCPClient(keys, projectName string) (*GCPClient, error) {
	log.Debugf("Connecting to GCP")
	ctx := context.Background()
	var client *GCPClient
	if keys != "" {
		log.Debugf("Using Keys %s", keys)
		f, err := os.Open(keys)
		if err != nil {
			return nil, err
		}

		jsonKey, err := ioutil.ReadAll(f)
		if err != nil {
			return nil, err
		}

		config, err := google.JWTConfigFromJSON(jsonKey,
			storage.DevstorageReadWriteScope,
			compute.ComputeScope,
		)
		if err != nil {
			return nil, err
		}

		client = &GCPClient{
			client:      config.Client(ctx),
			projectName: projectName,
		}
	} else {
		log.Debugf("Using Application Default crednetials")
		gc, err := google.DefaultClient(
			ctx,
			storage.DevstorageReadWriteScope,
			compute.ComputeScope,
		)
		if err != nil {
			return nil, err
		}
		client = &GCPClient{
			client:      gc,
			projectName: projectName,
		}
	}

	var err error
	client.compute, err = compute.New(client.client)
	if err != nil {
		return nil, err
	}

	client.storage, err = storage.New(client.client)
	if err != nil {
		return nil, err
	}

	log.Debugf("Generating SSH Keypair")
	client.privKey, err = rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, err
	}

	return client, nil
}

// UploadFile uploads a file to Google Storage
func (g GCPClient) UploadFile(src, dst, bucketName string, public bool) error {
	log.Infof("Uploading file %s to Google Storage as %s", src, dst)
	f, err := os.Open(src)
	if err != nil {
		return err
	}
	defer f.Close()

	objectCall := g.storage.Objects.Insert(bucketName, &storage.Object{Name: dst}).Media(f)

	if public {
		objectCall.PredefinedAcl("publicRead")
	}

	_, err = objectCall.Do()
	if err != nil {
		return err
	}
	log.Infof("Upload Complete!")
	fmt.Println("gs://" + bucketName + "/" + dst)
	return nil
}

// CreateImage creates a GCP image using the a source from Google Storage
func (g GCPClient) CreateImage(name, storageURL, family string, replace bool) error {
	if replace {
		if err := g.DeleteImage(name); err != nil {
			return err
		}
	}

	log.Infof("Creating image: %s", name)
	imgObj := &compute.Image{
		RawDisk: &compute.ImageRawDisk{
			Source: storageURL,
		},
		Name: name,
	}

	if family != "" {
		imgObj.Family = family
	}

	op, err := g.compute.Images.Insert(g.projectName, imgObj).Do()
	if err != nil {
		return err
	}

	if err := g.pollOperationStatus(op.Name); err != nil {
		return err
	}
	log.Infof("Image %s created", name)
	return nil
}

// DeleteImage deletes and image
func (g GCPClient) DeleteImage(name string) error {
	var notFound bool
	op, err := g.compute.Images.Delete(g.projectName, name).Do()
	if err != nil {
		if _, ok := err.(*googleapi.Error); !ok {
			return err
		}
		if err.(*googleapi.Error).Code != 404 {
			return err
		}
		notFound = true
	}
	if !notFound {
		log.Infof("Deleting existing image...")
		if err := g.pollOperationStatus(op.Name); err != nil {
			return err
		}
		log.Infof("Image %s deleted", name)
	}
	return nil
}

// CreateInstance creates and starts an instance on GCP
func (g GCPClient) CreateInstance(name, image, zone, machineType string, disks Disks, replace bool) error {
	if replace {
		if err := g.DeleteInstance(name, zone, true); err != nil {
			return err
		}
	}

	log.Infof("Creating instance %s from image %s (type: %s in %s)", name, image, machineType, zone)

	enabled := new(string)
	*enabled = "1"

	k, err := ssh.NewPublicKey(g.privKey.Public())
	if err != nil {
		return err
	}
	sshKey := new(string)
	*sshKey = fmt.Sprintf("moby:%s moby", string(ssh.MarshalAuthorizedKey(k)))

	instanceDisks := []*compute.AttachedDisk{
		{
			AutoDelete: true,
			Boot:       true,
			InitializeParams: &compute.AttachedDiskInitializeParams{
				SourceImage: fmt.Sprintf("global/images/%s", image),
			},
		},
	}

	for i, disk := range disks {
		var diskName string
		if disk.Path != "" {
			diskName = disk.Path
		} else {
			diskName = fmt.Sprintf("%s-disk-%d", name, i)
		}
		var diskSizeGb int64
		if disk.Size == 0 {
			diskSizeGb = int64(1)
		} else {
			diskSizeGb = int64(convertMBtoGB(disk.Size))
		}
		diskOp, err := g.compute.Disks.Insert(g.projectName, zone, &compute.Disk{Name: diskName, SizeGb: diskSizeGb}).Do()
		if err != nil {
			return err
		}
		if err := g.pollZoneOperationStatus(diskOp.Name, zone); err != nil {
			return err
		}
		instanceDisks = append(instanceDisks, &compute.AttachedDisk{
			AutoDelete: true,
			Boot:       false,
			Source:     fmt.Sprintf("zones/%s/disks/%s", zone, diskName),
		})
	}

	instanceObj := &compute.Instance{
		MachineType: fmt.Sprintf("zones/%s/machineTypes/%s", zone, machineType),
		Name:        name,
		Disks:       instanceDisks,
		NetworkInterfaces: []*compute.NetworkInterface{
			{
				Network: "global/networks/default",
				AccessConfigs: []*compute.AccessConfig{
					{
						Type: "ONE_TO_ONE_NAT",
					},
				},
			},
		},
		Metadata: &compute.Metadata{
			Items: []*compute.MetadataItems{
				{
					Key:   "serial-port-enable",
					Value: enabled,
				},
				{
					Key:   "ssh-keys",
					Value: sshKey,
				},
			},
		},
	}

	// Don't wait for operation to complete!
	// A headstart is needed as by the time we've polled for this event to be
	// completed, the instance may have already terminated
	_, err = g.compute.Instances.Insert(g.projectName, zone, instanceObj).Do()
	if err != nil {
		return err
	}
	log.Infof("Instance created")
	return nil
}

// DeleteInstance removes an instance
func (g GCPClient) DeleteInstance(instance, zone string, wait bool) error {
	var notFound bool
	op, err := g.compute.Instances.Delete(g.projectName, zone, instance).Do()
	if err != nil {
		if _, ok := err.(*googleapi.Error); !ok {
			return err
		}
		if err.(*googleapi.Error).Code != 404 {
			return err
		}
		notFound = true
	}
	if !notFound && wait {
		log.Infof("Deleting existing instance...")
		if err := g.pollZoneOperationStatus(op.Name, zone); err != nil {
			return err
		}
		log.Infof("Instance %s deleted", instance)
	}
	return nil
}

// GetInstanceSerialOutput streams the serial output of an instance
func (g GCPClient) GetInstanceSerialOutput(instance, zone string) error {
	log.Infof("Getting serial port output for instance %s", instance)
	var next int64
	for {
		res, err := g.compute.Instances.GetSerialPortOutput(g.projectName, zone, instance).Start(next).Do()
		if err != nil {
			if err.(*googleapi.Error).Code == 400 {
				// Instance may not be ready yet...
				time.Sleep(pollingInterval)
				continue
			}
			if err.(*googleapi.Error).Code == 503 {
				// Timeout received when the instance has terminated
				break
			}
			return err
		}
		fmt.Printf(res.Contents)
		next = res.Next
		// When the instance has been stopped, Start and Next will both be 0
		if res.Start > 0 && next == 0 {
			break
		}
	}
	return nil
}

// ConnectToInstanceSerialPort uses SSH to connect to the serial port of the instance
func (g GCPClient) ConnectToInstanceSerialPort(instance, zone string) error {
	log.Infof("Connecting to serial port of instance %s", instance)
	gPubKeyURL := "https://cloud-certs.storage.googleapis.com/google-cloud-serialport-host-key.pub"
	resp, err := http.Get(gPubKeyURL)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	gPubKey, _, _, _, err := ssh.ParseAuthorizedKey(body)
	if err != nil {
		return err
	}

	signer, err := ssh.NewSignerFromKey(g.privKey)
	if err != nil {
		return err
	}
	config := &ssh.ClientConfig{
		User: fmt.Sprintf("%s.%s.%s.moby", g.projectName, zone, instance),
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
		HostKeyCallback: ssh.FixedHostKey(gPubKey),
		Timeout:         5 * time.Second,
	}

	var conn *ssh.Client
	// Retry connection as VM may not be ready yet
	for i := 0; i < timeout; i++ {
		conn, err = ssh.Dial("tcp", "ssh-serialport.googleapis.com:9600", config)
		if err != nil {
			time.Sleep(pollingInterval)
			continue
		}
		break
	}
	if conn == nil {
		return fmt.Errorf(err.Error())
	}
	defer conn.Close()

	session, err := conn.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()

	stdin, err := session.StdinPipe()
	if err != nil {
		return fmt.Errorf("Unable to setup stdin for session: %v", err)
	}
	go io.Copy(stdin, os.Stdin)

	stdout, err := session.StdoutPipe()
	if err != nil {
		return fmt.Errorf("Unable to setup stdout for session: %v", err)
	}
	go io.Copy(os.Stdout, stdout)

	stderr, err := session.StderrPipe()
	if err != nil {
		return fmt.Errorf("Unable to setup stderr for session: %v", err)
	}
	go io.Copy(os.Stderr, stderr)
	/*
		c := make(chan os.Signal, 1)
		exit := make(chan bool, 1)
		signal.Notify(c)
		go func(exit <-chan bool, c <-chan os.Signal) {
			select {
			case <-exit:
				return
			case s := <-c:
				switch s {
				// CTRL+C
				case os.Interrupt:
					session.Signal(ssh.SIGINT)
				// CTRL+\
				case os.Kill:
					session.Signal(ssh.SIGQUIT)
				default:
					log.Debugf("Received signal %s but not forwarding to ssh", s)
				}
			}
		}(exit, c)
	*/
	var termWidth, termHeight int
	fd := os.Stdin.Fd()

	if term.IsTerminal(fd) {
		oldState, err := term.MakeRaw(fd)
		if err != nil {
			return err
		}

		defer term.RestoreTerminal(fd, oldState)

		winsize, err := term.GetWinsize(fd)
		if err != nil {
			termWidth = 80
			termHeight = 24
		} else {
			termWidth = int(winsize.Width)
			termHeight = int(winsize.Height)
		}
	}

	if err = session.RequestPty("xterm", termHeight, termWidth, ssh.TerminalModes{
		ssh.ECHO: 1,
	}); err != nil {
		return err
	}

	if err = session.Shell(); err != nil {
		return err
	}

	err = session.Wait()
	//exit <- true
	if err != nil {
		return err
	}
	return nil
}

func (g *GCPClient) pollOperationStatus(operationName string) error {
	for i := 0; i < timeout; i++ {
		operation, err := g.compute.GlobalOperations.Get(g.projectName, operationName).Do()
		if err != nil {
			return fmt.Errorf("error fetching operation status: %v", err)
		}
		if operation.Error != nil {
			return fmt.Errorf("error running operation: %v", operation.Error)
		}
		if operation.Status == "DONE" {
			return nil
		}
		time.Sleep(pollingInterval)
	}
	return fmt.Errorf("timeout waiting for operation to finish")

}
func (g *GCPClient) pollZoneOperationStatus(operationName, zone string) error {
	for i := 0; i < timeout; i++ {
		operation, err := g.compute.ZoneOperations.Get(g.projectName, zone, operationName).Do()
		if err != nil {
			return fmt.Errorf("error fetching operation status: %v", err)
		}
		if operation.Error != nil {
			return fmt.Errorf("error running operation: %v", operation.Error)
		}
		if operation.Status == "DONE" {
			return nil
		}
		time.Sleep(pollingInterval)
	}
	return fmt.Errorf("timeout waiting for operation to finish")
}
<commit_msg>Fix GCP error printing<commit_after>package main

import (
	"crypto/rand"
	"crypto/rsa"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"time"

	"github.com/docker/docker/pkg/term"
	log "github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh"
	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/compute/v1"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/storage/v1"
)

const pollingInterval = 500 * time.Millisecond
const timeout = 300

// GCPClient contains state required for communication with GCP
type GCPClient struct {
	client      *http.Client
	compute     *compute.Service
	storage     *storage.Service
	projectName string
	fileName    string
	privKey     *rsa.PrivateKey
}

// NewGCPClient creates a new GCP client
func NewGCPClient(keys, projectName string) (*GCPClient, error) {
	log.Debugf("Connecting to GCP")
	ctx := context.Background()
	var client *GCPClient
	if keys != "" {
		log.Debugf("Using Keys %s", keys)
		f, err := os.Open(keys)
		if err != nil {
			return nil, err
		}

		jsonKey, err := ioutil.ReadAll(f)
		if err != nil {
			return nil, err
		}

		config, err := google.JWTConfigFromJSON(jsonKey,
			storage.DevstorageReadWriteScope,
			compute.ComputeScope,
		)
		if err != nil {
			return nil, err
		}

		client = &GCPClient{
			client:      config.Client(ctx),
			projectName: projectName,
		}
	} else {
		log.Debugf("Using Application Default crednetials")
		gc, err := google.DefaultClient(
			ctx,
			storage.DevstorageReadWriteScope,
			compute.ComputeScope,
		)
		if err != nil {
			return nil, err
		}
		client = &GCPClient{
			client:      gc,
			projectName: projectName,
		}
	}

	var err error
	client.compute, err = compute.New(client.client)
	if err != nil {
		return nil, err
	}

	client.storage, err = storage.New(client.client)
	if err != nil {
		return nil, err
	}

	log.Debugf("Generating SSH Keypair")
	client.privKey, err = rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, err
	}

	return client, nil
}

// UploadFile uploads a file to Google Storage
func (g GCPClient) UploadFile(src, dst, bucketName string, public bool) error {
	log.Infof("Uploading file %s to Google Storage as %s", src, dst)
	f, err := os.Open(src)
	if err != nil {
		return err
	}
	defer f.Close()

	objectCall := g.storage.Objects.Insert(bucketName, &storage.Object{Name: dst}).Media(f)

	if public {
		objectCall.PredefinedAcl("publicRead")
	}

	_, err = objectCall.Do()
	if err != nil {
		return err
	}
	log.Infof("Upload Complete!")
	fmt.Println("gs://" + bucketName + "/" + dst)
	return nil
}

// CreateImage creates a GCP image using the a source from Google Storage
func (g GCPClient) CreateImage(name, storageURL, family string, replace bool) error {
	if replace {
		if err := g.DeleteImage(name); err != nil {
			return err
		}
	}

	log.Infof("Creating image: %s", name)
	imgObj := &compute.Image{
		RawDisk: &compute.ImageRawDisk{
			Source: storageURL,
		},
		Name: name,
	}

	if family != "" {
		imgObj.Family = family
	}

	op, err := g.compute.Images.Insert(g.projectName, imgObj).Do()
	if err != nil {
		return err
	}

	if err := g.pollOperationStatus(op.Name); err != nil {
		return err
	}
	log.Infof("Image %s created", name)
	return nil
}

// DeleteImage deletes and image
func (g GCPClient) DeleteImage(name string) error {
	var notFound bool
	op, err := g.compute.Images.Delete(g.projectName, name).Do()
	if err != nil {
		if _, ok := err.(*googleapi.Error); !ok {
			return err
		}
		if err.(*googleapi.Error).Code != 404 {
			return err
		}
		notFound = true
	}
	if !notFound {
		log.Infof("Deleting existing image...")
		if err := g.pollOperationStatus(op.Name); err != nil {
			return err
		}
		log.Infof("Image %s deleted", name)
	}
	return nil
}

// CreateInstance creates and starts an instance on GCP
func (g GCPClient) CreateInstance(name, image, zone, machineType string, disks Disks, replace bool) error {
	if replace {
		if err := g.DeleteInstance(name, zone, true); err != nil {
			return err
		}
	}

	log.Infof("Creating instance %s from image %s (type: %s in %s)", name, image, machineType, zone)

	enabled := new(string)
	*enabled = "1"

	k, err := ssh.NewPublicKey(g.privKey.Public())
	if err != nil {
		return err
	}
	sshKey := new(string)
	*sshKey = fmt.Sprintf("moby:%s moby", string(ssh.MarshalAuthorizedKey(k)))

	instanceDisks := []*compute.AttachedDisk{
		{
			AutoDelete: true,
			Boot:       true,
			InitializeParams: &compute.AttachedDiskInitializeParams{
				SourceImage: fmt.Sprintf("global/images/%s", image),
			},
		},
	}

	for i, disk := range disks {
		var diskName string
		if disk.Path != "" {
			diskName = disk.Path
		} else {
			diskName = fmt.Sprintf("%s-disk-%d", name, i)
		}
		var diskSizeGb int64
		if disk.Size == 0 {
			diskSizeGb = int64(1)
		} else {
			diskSizeGb = int64(convertMBtoGB(disk.Size))
		}
		diskOp, err := g.compute.Disks.Insert(g.projectName, zone, &compute.Disk{Name: diskName, SizeGb: diskSizeGb}).Do()
		if err != nil {
			return err
		}
		if err := g.pollZoneOperationStatus(diskOp.Name, zone); err != nil {
			return err
		}
		instanceDisks = append(instanceDisks, &compute.AttachedDisk{
			AutoDelete: true,
			Boot:       false,
			Source:     fmt.Sprintf("zones/%s/disks/%s", zone, diskName),
		})
	}

	instanceObj := &compute.Instance{
		MachineType: fmt.Sprintf("zones/%s/machineTypes/%s", zone, machineType),
		Name:        name,
		Disks:       instanceDisks,
		NetworkInterfaces: []*compute.NetworkInterface{
			{
				Network: "global/networks/default",
				AccessConfigs: []*compute.AccessConfig{
					{
						Type: "ONE_TO_ONE_NAT",
					},
				},
			},
		},
		Metadata: &compute.Metadata{
			Items: []*compute.MetadataItems{
				{
					Key:   "serial-port-enable",
					Value: enabled,
				},
				{
					Key:   "ssh-keys",
					Value: sshKey,
				},
			},
		},
	}

	// Don't wait for operation to complete!
	// A headstart is needed as by the time we've polled for this event to be
	// completed, the instance may have already terminated
	_, err = g.compute.Instances.Insert(g.projectName, zone, instanceObj).Do()
	if err != nil {
		return err
	}
	log.Infof("Instance created")
	return nil
}

// DeleteInstance removes an instance
func (g GCPClient) DeleteInstance(instance, zone string, wait bool) error {
	var notFound bool
	op, err := g.compute.Instances.Delete(g.projectName, zone, instance).Do()
	if err != nil {
		if _, ok := err.(*googleapi.Error); !ok {
			return err
		}
		if err.(*googleapi.Error).Code != 404 {
			return err
		}
		notFound = true
	}
	if !notFound && wait {
		log.Infof("Deleting existing instance...")
		if err := g.pollZoneOperationStatus(op.Name, zone); err != nil {
			return err
		}
		log.Infof("Instance %s deleted", instance)
	}
	return nil
}

// GetInstanceSerialOutput streams the serial output of an instance
func (g GCPClient) GetInstanceSerialOutput(instance, zone string) error {
	log.Infof("Getting serial port output for instance %s", instance)
	var next int64
	for {
		res, err := g.compute.Instances.GetSerialPortOutput(g.projectName, zone, instance).Start(next).Do()
		if err != nil {
			if err.(*googleapi.Error).Code == 400 {
				// Instance may not be ready yet...
				time.Sleep(pollingInterval)
				continue
			}
			if err.(*googleapi.Error).Code == 503 {
				// Timeout received when the instance has terminated
				break
			}
			return err
		}
		fmt.Printf(res.Contents)
		next = res.Next
		// When the instance has been stopped, Start and Next will both be 0
		if res.Start > 0 && next == 0 {
			break
		}
	}
	return nil
}

// ConnectToInstanceSerialPort uses SSH to connect to the serial port of the instance
func (g GCPClient) ConnectToInstanceSerialPort(instance, zone string) error {
	log.Infof("Connecting to serial port of instance %s", instance)
	gPubKeyURL := "https://cloud-certs.storage.googleapis.com/google-cloud-serialport-host-key.pub"
	resp, err := http.Get(gPubKeyURL)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	gPubKey, _, _, _, err := ssh.ParseAuthorizedKey(body)
	if err != nil {
		return err
	}

	signer, err := ssh.NewSignerFromKey(g.privKey)
	if err != nil {
		return err
	}
	config := &ssh.ClientConfig{
		User: fmt.Sprintf("%s.%s.%s.moby", g.projectName, zone, instance),
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
		HostKeyCallback: ssh.FixedHostKey(gPubKey),
		Timeout:         5 * time.Second,
	}

	var conn *ssh.Client
	// Retry connection as VM may not be ready yet
	for i := 0; i < timeout; i++ {
		conn, err = ssh.Dial("tcp", "ssh-serialport.googleapis.com:9600", config)
		if err != nil {
			time.Sleep(pollingInterval)
			continue
		}
		break
	}
	if conn == nil {
		return fmt.Errorf(err.Error())
	}
	defer conn.Close()

	session, err := conn.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()

	stdin, err := session.StdinPipe()
	if err != nil {
		return fmt.Errorf("Unable to setup stdin for session: %v", err)
	}
	go io.Copy(stdin, os.Stdin)

	stdout, err := session.StdoutPipe()
	if err != nil {
		return fmt.Errorf("Unable to setup stdout for session: %v", err)
	}
	go io.Copy(os.Stdout, stdout)

	stderr, err := session.StderrPipe()
	if err != nil {
		return fmt.Errorf("Unable to setup stderr for session: %v", err)
	}
	go io.Copy(os.Stderr, stderr)
	/*
		c := make(chan os.Signal, 1)
		exit := make(chan bool, 1)
		signal.Notify(c)
		go func(exit <-chan bool, c <-chan os.Signal) {
			select {
			case <-exit:
				return
			case s := <-c:
				switch s {
				// CTRL+C
				case os.Interrupt:
					session.Signal(ssh.SIGINT)
				// CTRL+\
				case os.Kill:
					session.Signal(ssh.SIGQUIT)
				default:
					log.Debugf("Received signal %s but not forwarding to ssh", s)
				}
			}
		}(exit, c)
	*/
	var termWidth, termHeight int
	fd := os.Stdin.Fd()

	if term.IsTerminal(fd) {
		oldState, err := term.MakeRaw(fd)
		if err != nil {
			return err
		}

		defer term.RestoreTerminal(fd, oldState)

		winsize, err := term.GetWinsize(fd)
		if err != nil {
			termWidth = 80
			termHeight = 24
		} else {
			termWidth = int(winsize.Width)
			termHeight = int(winsize.Height)
		}
	}

	if err = session.RequestPty("xterm", termHeight, termWidth, ssh.TerminalModes{
		ssh.ECHO: 1,
	}); err != nil {
		return err
	}

	if err = session.Shell(); err != nil {
		return err
	}

	err = session.Wait()
	//exit <- true
	if err != nil {
		return err
	}
	return nil
}

func (g *GCPClient) pollOperationStatus(operationName string) error {
	for i := 0; i < timeout; i++ {
		operation, err := g.compute.GlobalOperations.Get(g.projectName, operationName).Do()
		if err != nil {
			return fmt.Errorf("error fetching operation status: %v", err)
		}
		if operation.Error != nil {
			return fmt.Errorf("error running operation: %s", operation.Error)
		}
		if operation.Status == "DONE" {
			return nil
		}
		time.Sleep(pollingInterval)
	}
	return fmt.Errorf("timeout waiting for operation to finish")

}
func (g *GCPClient) pollZoneOperationStatus(operationName, zone string) error {
	for i := 0; i < timeout; i++ {
		operation, err := g.compute.ZoneOperations.Get(g.projectName, zone, operationName).Do()
		if err != nil {
			return fmt.Errorf("error fetching operation status: %v", err)
		}
		if operation.Error != nil {
			return fmt.Errorf("error running operation: %s", operation.Error)
		}
		if operation.Status == "DONE" {
			return nil
		}
		time.Sleep(pollingInterval)
	}
	return fmt.Errorf("timeout waiting for operation to finish")
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"testing"
)

func TestGetReadPostsFromTumblr(t *testing.T) {
	defer os.Unsetenv(readPostsFromTumblrEnv)
	os.Setenv(readPostsFromTumblrEnv, "false")
	if getReadPostsFromTumblr() {
		t.Fail()
	}
	os.Setenv(readPostsFromTumblrEnv, "True")
	if !getReadPostsFromTumblr() {
		t.Fail()
	}
}
<commit_msg>Add test for getNewRelicApp<commit_after>package main

import (
	"os"
	"testing"
)

func TestGetReadPostsFromTumblr(t *testing.T) {
	defer os.Unsetenv(readPostsFromTumblrEnv)
	os.Setenv(readPostsFromTumblrEnv, "false")
	if getReadPostsFromTumblr() {
		t.Fail()
	}
	os.Setenv(readPostsFromTumblrEnv, "True")
	if !getReadPostsFromTumblr() {
		t.Fail()
	}
}

func TestGetNewRelicApp(t *testing.T) {
	app := getNewRelicApp()
	if app == nil {
		t.Fail()
	}
}
<|endoftext|>


<commit_before>package golenv

import (
	"os"
	"strings"
)

func OverrideIfEnv(envVar string, defaultValue string) string {
	if HasEnv(envVar) {
		envVarValue := os.Getenv(envVar)
		if envVarValue != "" {
			return envVarValue
		}
	}
	return defaultValue
}

func HasEnv(envVar string) bool {
	for _, envKeyVal := range os.Environ() {
		if strings.Split(envKeyVal, "=")[0] == envVar {
			return true
		}
	}
	return false
}

func EnvMap() map[string]string {
	var key_val = make(map[string]string)

	for _, envKeyVal := range os.Environ() {
		envKeyValSplit := strings.Split(envKeyVal, "=")
		key_val[envKeyValSplit[0]] = envKeyValSplit[1]
	}
	return key_val
}
<commit_msg>[golenv] added OverrideIfEnvBool, OverrideIfEnvInt, OverrideIfEnvUint64 for typed config from env<commit_after>package golenv

import (
	"os"
	"strings"

	"github.com/abhishekkr/gol/golconv"
)

func OverrideIfEnv(envVar string, defaultValue string) string {
	if HasEnv(envVar) {
		envVarValue := os.Getenv(envVar)
		if envVarValue != "" {
			return envVarValue
		}
	}
	return defaultValue
}

func OverrideIfEnvBool(envVar string, defaultValue bool) bool {
	if HasEnv(envVar) {
		envVarValue := os.Getenv(envVar)
		return golconv.StringToBool(envVarValue, defaultValue)
	}
	return defaultValue
}

func OverrideIfEnvInt(envVar string, defaultValue int) int {
	if HasEnv(envVar) {
		envVarValue := os.Getenv(envVar)
		return golconv.StringToInt(envVarValue, defaultValue)
	}
	return defaultValue
}

func OverrideIfEnvUint64(envVar string, defaultValue uint64) uint64 {
	if HasEnv(envVar) {
		envVarValue := os.Getenv(envVar)
		return golconv.StringToUint64(envVarValue, defaultValue)
	}
	return defaultValue
}

func HasEnv(envVar string) bool {
	for _, envKeyVal := range os.Environ() {
		if strings.Split(envKeyVal, "=")[0] == envVar {
			return true
		}
	}
	return false
}

func EnvMap() map[string]string {
	var key_val = make(map[string]string)

	for _, envKeyVal := range os.Environ() {
		envKeyValSplit := strings.Split(envKeyVal, "=")
		key_val[envKeyValSplit[0]] = envKeyValSplit[1]
	}
	return key_val
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: MIT

package createtable_test

import (
	"testing"

	"github.com/issue9/assert"

	"github.com/issue9/orm/v3/core"
	"github.com/issue9/orm/v3/internal/createtable"
	"github.com/issue9/orm/v3/internal/sqltest"
	"github.com/issue9/orm/v3/internal/test"
)

var mysqlCreateTable = []string{`CREATE TABLE fk_table(
	id BIGINT(20) NOT NULL,
	PRIMARY KEY(id)
	)`,
	`CREATE TABLE users (
	id BIGINT(20) NOT NULL,
	created BIGINT NOT NULL,
	nickname VARCHAR(50) NOT NULL,
	state SMALLINT(6) NOT NULL,
	username VARCHAR(50) NOT NULL,
	mobile VARCHAR(18) NOT NULL,
	email VARCHAR(50) NOT NULL,
	password VARCHAR(128) NOT NULL,
	PRIMARY KEY (id,username),
	UNIQUE KEY u_user_xx1 (mobile,username),
	UNIQUE KEY u_user_email1 (email,username),
	UNIQUE KEY unique_id (id),
	KEY index_user_mobile (mobile),
	CONSTRAINT xxx_fk FOREIGN KEY (id) REFERENCES fk_table (id),
	CONSTRAINT xxx CHECK(created > 0)
	)`,
}

func TestParseMysqlCreateTable(t *testing.T) {
	a := assert.New(t)

	suite := test.NewSuite(a)
	defer suite.Close()

	suite.ForEach(func(t *test.Driver) {
		defer func() {
			_, err := t.DB.Exec("DROP TABLE users")
			a.NotError(err)

			_, err = t.DB.Exec("DROP TABLE fk_table")
			a.NotError(err)
		}()
		for _, query := range mysqlCreateTable {
			_, err := t.DB.Exec(query)
			t.NotError(err)
		}

		table, err := createtable.ParseMysqlCreateTable("users", t.DB)
		t.NotError(err).NotNil(table)

		t.Equal(len(table.Columns), 8)
		sqltest.Equal(a, table.Columns["id"], "bigint NOT NULL")
		sqltest.Equal(a, table.Columns["created"], "bigint NOT NULL")
		sqltest.Equal(a, table.Columns["nickname"], "varchar(50) NOT NULL")
		sqltest.Equal(a, table.Columns["state"], "smallint NOT NULL")
		sqltest.Equal(a, table.Columns["username"], "varchar(50) NOT NULL")
		sqltest.Equal(a, table.Columns["mobile"], "varchar(18) NOT NULL")
		sqltest.Equal(a, table.Columns["email"], "varchar(50) NOT NULL")
		sqltest.Equal(a, table.Columns["password"], "varchar(128) NOT NULL")
		t.Equal(len(table.Constraints), 6).
			Equal(table.Constraints["u_user_xx1"], core.ConstraintUnique).
			Equal(table.Constraints["u_user_email1"], core.ConstraintUnique).
			Equal(table.Constraints["unique_id"], core.ConstraintUnique).
			Equal(table.Constraints["xxx_fk"], core.ConstraintFK).
			Equal(table.Constraints["xxx"], core.ConstraintCheck).
			Equal(table.Constraints["users_pk"], core.ConstraintPK) // 主键约束名为固定值
		t.Equal(len(table.Indexes), 1).
			Equal(table.Indexes["index_user_mobile"], core.IndexDefault)
	}, test.Mysql, test.Mariadb)
}
<commit_msg>fix: 修正 mariadb 错误<commit_after>// SPDX-License-Identifier: MIT

package createtable_test

import (
	"testing"

	"github.com/issue9/assert"

	"github.com/issue9/orm/v3/core"
	"github.com/issue9/orm/v3/internal/createtable"
	"github.com/issue9/orm/v3/internal/sqltest"
	"github.com/issue9/orm/v3/internal/test"
)

var mysqlCreateTable = []string{`CREATE TABLE fk_table(
	id BIGINT(20) NOT NULL,
	PRIMARY KEY(id)
	)`,
	`CREATE TABLE users (
	id BIGINT(20) NOT NULL,
	created BIGINT NOT NULL,
	nickname VARCHAR(50) NOT NULL,
	state SMALLINT(6) NOT NULL,
	username VARCHAR(50) NOT NULL,
	mobile VARCHAR(18) NOT NULL,
	email VARCHAR(50) NOT NULL,
	password VARCHAR(128) NOT NULL,
	PRIMARY KEY (id,username),
	UNIQUE KEY u_user_xx1 (mobile,username),
	UNIQUE KEY u_user_email1 (email,username),
	UNIQUE KEY unique_id (id),
	KEY index_user_mobile (mobile),
	CONSTRAINT xxx_fk FOREIGN KEY (id) REFERENCES fk_table (id),
	CONSTRAINT xxx CHECK(created > 0)
	)`,
}

func TestParseMysqlCreateTable(t *testing.T) {
	a := assert.New(t)

	suite := test.NewSuite(a)
	defer suite.Close()

	suite.ForEach(func(t *test.Driver) {
		defer func() {
			_, err := t.DB.Exec("DROP TABLE users")
			a.NotError(err)

			_, err = t.DB.Exec("DROP TABLE fk_table")
			a.NotError(err)
		}()
		for _, query := range mysqlCreateTable {
			_, err := t.DB.Exec(query)
			t.NotError(err)
		}

		table, err := createtable.ParseMysqlCreateTable("users", t.DB)
		t.NotError(err).NotNil(table)

		t.Equal(len(table.Columns), 8)
		sqltest.Equal(a, table.Columns["id"], "bigint NOT NULL")
		sqltest.Equal(a, table.Columns["created"], "bigint NOT NULL")
		sqltest.Equal(a, table.Columns["nickname"], "varchar(50) NOT NULL")
		sqltest.Equal(a, table.Columns["state"], "smallint NOT NULL")
		sqltest.Equal(a, table.Columns["username"], "varchar(50) NOT NULL")
		sqltest.Equal(a, table.Columns["mobile"], "varchar(18) NOT NULL")
		sqltest.Equal(a, table.Columns["email"], "varchar(50) NOT NULL")
		sqltest.Equal(a, table.Columns["password"], "varchar(128) NOT NULL")
		t.Equal(len(table.Constraints), 6).
			Equal(table.Constraints["u_user_xx1"], core.ConstraintUnique).
			Equal(table.Constraints["u_user_email1"], core.ConstraintUnique).
			Equal(table.Constraints["unique_id"], core.ConstraintUnique).
			Equal(table.Constraints["xxx_fk"], core.ConstraintFK).
			Equal(table.Constraints["xxx"], core.ConstraintCheck).
			Equal(table.Constraints["users_pk"], core.ConstraintPK) // 主键约束名为固定值
		t.Equal(len(table.Indexes), 1).
			Equal(table.Indexes["index_user_mobile"], core.IndexDefault)
	}, test.Mysql)
}

// mariadb 稍微和 mysql 有点不一样
func TestParseMysqlCreateTable_mariadb(t *testing.T) {
	a := assert.New(t)

	suite := test.NewSuite(a)
	defer suite.Close()

	suite.ForEach(func(t *test.Driver) {
		defer func() {
			_, err := t.DB.Exec("DROP TABLE users")
			a.NotError(err)

			_, err = t.DB.Exec("DROP TABLE fk_table")
			a.NotError(err)
		}()
		for _, query := range mysqlCreateTable {
			_, err := t.DB.Exec(query)
			t.NotError(err)
		}

		table, err := createtable.ParseMysqlCreateTable("users", t.DB)
		t.NotError(err).NotNil(table)

		t.Equal(len(table.Columns), 8)
		sqltest.Equal(a, table.Columns["id"], "bigint(20) NOT NULL")
		sqltest.Equal(a, table.Columns["created"], "bigint(20) NOT NULL")
		sqltest.Equal(a, table.Columns["nickname"], "varchar(50) NOT NULL")
		sqltest.Equal(a, table.Columns["state"], "smallint(6) NOT NULL")
		sqltest.Equal(a, table.Columns["username"], "varchar(50) NOT NULL")
		sqltest.Equal(a, table.Columns["mobile"], "varchar(18) NOT NULL")
		sqltest.Equal(a, table.Columns["email"], "varchar(50) NOT NULL")
		sqltest.Equal(a, table.Columns["password"], "varchar(128) NOT NULL")
		t.Equal(len(table.Constraints), 6).
			Equal(table.Constraints["u_user_xx1"], core.ConstraintUnique).
			Equal(table.Constraints["u_user_email1"], core.ConstraintUnique).
			Equal(table.Constraints["unique_id"], core.ConstraintUnique).
			Equal(table.Constraints["xxx_fk"], core.ConstraintFK).
			Equal(table.Constraints["xxx"], core.ConstraintCheck).
			Equal(table.Constraints["users_pk"], core.ConstraintPK) // 主键约束名为固定值
		t.Equal(len(table.Indexes), 1).
			Equal(table.Indexes["index_user_mobile"], core.IndexDefault)
	}, test.Mariadb)
}
<|endoftext|>


<commit_before>package api

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"

	"encoding/base64"

	"github.com/sironfoot/go-twitter-bot/data/db"
	"github.com/sironfoot/go-twitter-bot/data/models"
	"golang.org/x/crypto/bcrypt"
	"golang.org/x/net/context"
)

// AccountLogin = POST/PUT: /account/login
func AccountLogin(ctx context.Context, res http.ResponseWriter, req *http.Request) {
	appContext := ctx.Value("appContext").(*AppContext)
	var login models.Login

	defer req.Body.Close()

	err := json.NewDecoder(io.LimitReader(req.Body, maxRequestLength)).Decode(&login)
	if err != nil {
		res.WriteHeader(http.StatusBadRequest)
		appContext.Response = MessageResponse{
			Message: "JSON request body was in invalid format",
		}
		return
	}

	login.Sanitise()
	validationErrors, err := login.Validate()
	if err != nil {
		panic(err)
	}

	if len(validationErrors) > 0 {
		res.WriteHeader(http.StatusBadRequest)
		appContext.Response = updateResponse{
			Message: "Login model is invalid.",
			Errors:  validationErrors,
		}
		return
	}

	// check user exists for email
	user, err := db.UserFromEmail(login.Email)
	if err == db.ErrEntityNotFound {
		res.WriteHeader(http.StatusForbidden)

		var validationErrors []models.ValidationError
		validationErrors = append(validationErrors, models.ValidationError{
			FieldName: "email",
			Type:      models.ValidationTypeNotFound,
			Message:   fmt.Sprintf("User not found on email %s", login.Email),
		})

		appContext.Response = updateResponse{
			Message: "Login model is invalid.",
			Errors:  validationErrors,
		}
		return
	} else if err != nil {
		panic(err)
	}

	// check password
	errHashCompare := bcrypt.CompareHashAndPassword([]byte(user.HashedPassword), []byte(login.Password))
	if errHashCompare != nil {
		if errHashCompare == bcrypt.ErrMismatchedHashAndPassword || errHashCompare == bcrypt.ErrHashTooShort {
			res.WriteHeader(http.StatusForbidden)

			var validationErrors []models.ValidationError
			validationErrors = append(validationErrors, models.ValidationError{
				FieldName: "password",
				Type:      models.ValidationTypeNotFound,
				Message:   fmt.Sprintf("Password not correct for user %s", login.Email),
			})

			appContext.Response = updateResponse{
				Message: "Login model is invalid.",
				Errors:  validationErrors,
			}
			return
		}

		panic(errHashCompare)
	}

	randomBytes := make([]byte, 32)
	_, err = rand.Read(randomBytes)
	if err != nil {
		panic(err)
	}

	user.AuthToken = sql.NullString{
		String: base64.StdEncoding.EncodeToString(randomBytes),
		Valid:  true,
	}

	err = user.Save()
	if err != nil {
		panic(err)
	}

	block, err := aes.NewCipher([]byte(appContext.Settings.AppSettings.EncryptionKey))
	if err != nil {
		panic(err)
	}

	plaintextToken := []byte(user.ID + "_" + user.AuthToken.String)
	encryptedToken := make([]byte, aes.BlockSize+len(plaintextToken))

	// iv = initialization vector
	iv := encryptedToken[:aes.BlockSize]
	_, err = rand.Read(iv)
	if err != nil {
		panic(err)
	}

	cfb := cipher.NewCFBEncrypter(block, iv)
	cfb.XORKeyStream(encryptedToken[aes.BlockSize:], plaintextToken)
	accessToken := base64.StdEncoding.EncodeToString(encryptedToken)

	appContext.Response = struct {
		Message     string `json:"message"`
		AccessToken string `json:"accessToken"`
	}{
		Message:     ok,
		AccessToken: accessToken,
	}
}

// AccountLogout = POST/PUT: /account/logout
func AccountLogout(ctx context.Context, res http.ResponseWriter, req *http.Request) {
	res.WriteHeader(http.StatusNotImplemented)
}

// AccountSignup = POST/PUT: /account/signup
func AccountSignup(ctx context.Context, res http.ResponseWriter, req *http.Request) {
	res.WriteHeader(http.StatusNotImplemented)
}
<commit_msg>added: PUT: /account/logout API endpoint.<commit_after>package api

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"

	"encoding/base64"

	"github.com/sironfoot/go-twitter-bot/data/db"
	"github.com/sironfoot/go-twitter-bot/data/models"
	"golang.org/x/crypto/bcrypt"
	"golang.org/x/net/context"
)

// AccountLogin = PUT: /account/login
func AccountLogin(ctx context.Context, res http.ResponseWriter, req *http.Request) {
	appContext := ctx.Value("appContext").(*AppContext)
	var login models.Login

	defer req.Body.Close()

	err := json.NewDecoder(io.LimitReader(req.Body, maxRequestLength)).Decode(&login)
	if err != nil {
		res.WriteHeader(http.StatusBadRequest)
		appContext.Response = MessageResponse{
			Message: "JSON request body was in invalid format",
		}
		return
	}

	login.Sanitise()
	validationErrors, err := login.Validate()
	if err != nil {
		panic(err)
	}

	if len(validationErrors) > 0 {
		res.WriteHeader(http.StatusBadRequest)
		appContext.Response = updateResponse{
			Message: "Login model is invalid.",
			Errors:  validationErrors,
		}
		return
	}

	// check user exists for email
	user, err := db.UserFromEmail(login.Email)
	if err == db.ErrEntityNotFound {
		res.WriteHeader(http.StatusForbidden)

		var validationErrors []models.ValidationError
		validationErrors = append(validationErrors, models.ValidationError{
			FieldName: "email",
			Type:      models.ValidationTypeNotFound,
			Message:   fmt.Sprintf("User not found on email %s", login.Email),
		})

		appContext.Response = updateResponse{
			Message: "Login model is invalid.",
			Errors:  validationErrors,
		}
		return
	} else if err != nil {
		panic(err)
	}

	// check password
	errHashCompare := bcrypt.CompareHashAndPassword([]byte(user.HashedPassword), []byte(login.Password))
	if errHashCompare != nil {
		if errHashCompare == bcrypt.ErrMismatchedHashAndPassword || errHashCompare == bcrypt.ErrHashTooShort {
			res.WriteHeader(http.StatusForbidden)

			var validationErrors []models.ValidationError
			validationErrors = append(validationErrors, models.ValidationError{
				FieldName: "password",
				Type:      models.ValidationTypeNotFound,
				Message:   fmt.Sprintf("Password not correct for user %s", login.Email),
			})

			appContext.Response = updateResponse{
				Message: "Login model is invalid.",
				Errors:  validationErrors,
			}
			return
		}

		panic(errHashCompare)
	}

	randomBytes := make([]byte, 32)
	_, err = rand.Read(randomBytes)
	if err != nil {
		panic(err)
	}

	user.AuthToken = sql.NullString{
		String: base64.StdEncoding.EncodeToString(randomBytes),
		Valid:  true,
	}

	err = user.Save()
	if err != nil {
		panic(err)
	}

	block, err := aes.NewCipher([]byte(appContext.Settings.AppSettings.EncryptionKey))
	if err != nil {
		panic(err)
	}

	plaintextToken := []byte(user.ID + "_" + user.AuthToken.String)
	encryptedToken := make([]byte, aes.BlockSize+len(plaintextToken))

	// iv = initialization vector
	iv := encryptedToken[:aes.BlockSize]
	_, err = rand.Read(iv)
	if err != nil {
		panic(err)
	}

	cfb := cipher.NewCFBEncrypter(block, iv)
	cfb.XORKeyStream(encryptedToken[aes.BlockSize:], plaintextToken)
	accessToken := base64.StdEncoding.EncodeToString(encryptedToken)

	appContext.Response = struct {
		Message     string `json:"message"`
		AccessToken string `json:"accessToken"`
	}{
		Message:     ok,
		AccessToken: accessToken,
	}
}

// AccountLogout = PUT: /account/logout
func AccountLogout(ctx context.Context, res http.ResponseWriter, req *http.Request) {
	appContext := ctx.Value("appContext").(*AppContext)

	if appContext.AuthUser != nil {
		authUser := appContext.AuthUser

		authUser.AuthToken = sql.NullString{
			Valid: false,
		}

		err := authUser.Save()
		if err != nil {
			panic(err)
		}
	}

	appContext.Response = MessageResponse{
		Message: ok,
	}
}

// AccountSignup = POST/PUT: /account/signup
func AccountSignup(ctx context.Context, res http.ResponseWriter, req *http.Request) {
	res.WriteHeader(http.StatusNotImplemented)
}
<|endoftext|>


<commit_before>package realtime

import (
	"encoding/json"
	"errors"
	"fmt"
	mongomodels "koding/db/models"
	"koding/db/mongodb/modelhelper"
	"socialapi/models"
	"socialapi/workers/helper"
	notificationmodels "socialapi/workers/notification/models"
	"strconv"

	"github.com/koding/logging"
	"github.com/koding/rabbitmq"
	"github.com/koding/worker"
	"github.com/streadway/amqp"
	"labix.org/v2/mgo"
)

type Action func(*RealtimeWorkerController, []byte) error

type RealtimeWorkerController struct {
	routes  map[string]Action
	log     logging.Logger
	rmqConn *amqp.Connection
}

type NotificationEvent struct {
	RoutingKey string              `json:"routingKey"`
	Content    NotificationContent `json:"contents"`
}

type NotificationContent struct {
	TypeConstant string `json:"type"`
	TargetId     int64  `json:"targetId"`
	ActorId      string `json:"actorId"`
}

func (r *RealtimeWorkerController) DefaultErrHandler(delivery amqp.Delivery, err error) bool {
	r.log.Error("an error occured deleting realtime event", err)
	delivery.Ack(false)
	return false
}

func NewRealtimeWorkerController(rmq *rabbitmq.RabbitMQ, log logging.Logger) (*RealtimeWorkerController, error) {
	rmqConn, err := rmq.Connect("NewRealtimeWorkerController")
	if err != nil {
		return nil, err
	}

	ffc := &RealtimeWorkerController{
		log:     log,
		rmqConn: rmqConn.Conn(),
	}

	routes := map[string]Action{
		"api.channel_message_created": (*RealtimeWorkerController).MessageSaved,
		"api.channel_message_updated": (*RealtimeWorkerController).MessageUpdated,
		"api.channel_message_deleted": (*RealtimeWorkerController).MessageDeleted,

		"api.interaction_created": (*RealtimeWorkerController).InteractionSaved,
		"api.interaction_deleted": (*RealtimeWorkerController).InteractionDeleted,

		"api.message_reply_created": (*RealtimeWorkerController).MessageReplySaved,
		"api.message_reply_deleted": (*RealtimeWorkerController).MessageReplyDeleted,

		"api.channel_message_list_created": (*RealtimeWorkerController).MessageListSaved,
		"api.channel_message_list_updated": (*RealtimeWorkerController).MessageListUpdated,
		"api.channel_message_list_deleted": (*RealtimeWorkerController).MessageListDeleted,

		"api.channel_participant_created": (*RealtimeWorkerController).ChannelParticipantAdded,
		"api.channel_participant_deleted": (*RealtimeWorkerController).ChannelParticipantRemoved,

		"notification.notification_created": (*RealtimeWorkerController).NotifyUser,
		"notification.notification_updated": (*RealtimeWorkerController).NotifyUser,
	}

	ffc.routes = routes

	return ffc, nil
}

func (f *RealtimeWorkerController) HandleEvent(event string, data []byte) error {
	f.log.Debug("New Event Received %s", event)
	handler, ok := f.routes[event]
	if !ok {
		return worker.HandlerNotFoundErr
	}

	return handler(f, data)
}

// no operation for message save for now
func (f *RealtimeWorkerController) MessageSaved(data []byte) error {
	return nil
}

// no operation for message delete for now
// channel_message_delete will handle message deletions from the
func (f *RealtimeWorkerController) MessageDeleted(data []byte) error {
	return nil
}

func (f *RealtimeWorkerController) MessageUpdated(data []byte) error {
	cm, err := helper.MapToChannelMessage(data)
	if err != nil {
		return err
	}

	err = f.sendInstanceEvent(cm.GetId(), cm, "updateInstance")
	if err != nil {
		fmt.Println(err)
		return err
	}

	return nil
}

func (f *RealtimeWorkerController) ChannelParticipantAdded(data []byte) error {
	return f.handleChannelParticipantEvent("AddedToChannel", data)
}

func (f *RealtimeWorkerController) ChannelParticipantRemoved(data []byte) error {
	return f.handleChannelParticipantEvent("RemovedFromChannel", data)
}

func (f *RealtimeWorkerController) handleChannelParticipantEvent(eventName string, data []byte) error {
	cp := models.NewChannelParticipant()
	if err := json.Unmarshal(data, cp); err != nil {
		return err
	}

	c := models.NewChannel()
	if err := c.ById(cp.ChannelId); err != nil {
		return err
	}

	return f.sendNotification(cp.AccountId, eventName, c)
}

func (f *RealtimeWorkerController) InteractionSaved(data []byte) error {
	return f.handleInteractionEvent("InteractionAdded", data)
}

func (f *RealtimeWorkerController) InteractionDeleted(data []byte) error {
	return f.handleInteractionEvent("InteractionRemoved", data)
}

func (f *RealtimeWorkerController) handleInteractionEvent(eventName string, data []byte) error {
	i, err := helper.MapToInteraction(data)
	if err != nil {
		return err
	}

	count, err := i.Count(i.TypeConstant)
	if err != nil {
		return err
	}

	oldId, err := models.AccountOldIdById(i.AccountId)
	if err != nil {
		return err
	}

	res := map[string]interface{}{
		"messageId":    i.MessageId,
		"accountId":    i.AccountId,
		"accountOldId": oldId,
		"typeConstant": i.TypeConstant,
		"count":        count,
	}

	err = f.sendInstanceEvent(i.MessageId, res, eventName)
	if err != nil {
		fmt.Println(err)
		return err
	}

	return nil
}

func (f *RealtimeWorkerController) MessageReplySaved(data []byte) error {
	i, err := helper.MapToMessageReply(data)
	if err != nil {
		return err
	}

	reply := models.NewChannelMessage()
	if err := reply.ById(i.ReplyId); err != nil {
		return err
	}

	cmc, err := reply.BuildEmptyMessageContainer()
	if err != nil {
		return err
	}

	err = f.sendInstanceEvent(i.MessageId, cmc, "ReplyAdded")
	if err != nil {
		fmt.Println(err)
		return err
	}

	return nil
}

func (f *RealtimeWorkerController) MessageReplyDeleted(data []byte) error {
	i, err := helper.MapToMessageReply(data)
	if err != nil {
		return err
	}

	err = f.sendInstanceEvent(i.MessageId, i, "ReplyRemoved")
	if err != nil {
		fmt.Println(err)
		return err
	}

	return nil
}

// send message to the channel
func (f *RealtimeWorkerController) MessageListSaved(data []byte) error {
	cml, err := helper.MapToChannelMessageList(data)
	if err != nil {
		return err
	}

	err = f.sendChannelEvent(cml, "MessageAdded")
	if err != nil {
		return err
	}

	return nil
}

// no operation for channel_message_list_updated event
func (f *RealtimeWorkerController) MessageListUpdated(data []byte) error {
	return nil
}

func (f *RealtimeWorkerController) MessageListDeleted(data []byte) error {
	cml, err := helper.MapToChannelMessageList(data)
	if err != nil {
		return err
	}

	err = f.sendChannelEvent(cml, "MessageRemoved")
	if err != nil {
		return err
	}
	return nil
}

func (f *RealtimeWorkerController) NotifyUser(data []byte) error {
	channel, err := f.rmqConn.Channel()
	if err != nil {
		return errors.New("channel connection error")
	}
	defer channel.Close()

	notification := notificationmodels.NewNotification()
	if err := notification.MapMessage(data); err != nil {
		return err
	}

	activity, nc, err := notification.FetchLastActivity()
	if err != nil {
		return err
	}

	// do not notify actor for her own action
	if activity.ActorId == notification.AccountId {
		return nil
	}

	// do not notify user when notification is not yet activated
	if notification.ActivatedAt.IsZero() {
		return nil
	}

	oldAccount, err := fetchNotifierOldAccount(notification.AccountId)
	if err != nil {
		return fmt.Errorf("an error occurred while fetching old account: %s", err)
	}

	// fetch user profile name from bongo as routing key
	ne := &NotificationEvent{}

	ne.Content = NotificationContent{
		TargetId:     nc.TargetId,
		TypeConstant: nc.TypeConstant,
	}
	ne.Content.ActorId, _ = models.AccountOldIdById(activity.ActorId)

	notificationMessage, err := json.Marshal(ne)
	if err != nil {
		return err
	}

	routingKey := oldAccount.Profile.Nickname

	err = channel.Publish(
		"notification",
		routingKey,
		false,
		false,
		amqp.Publishing{Body: notificationMessage},
	)
	if err != nil {
		return fmt.Errorf("an error occurred while notifying user: %s", err)
	}

	return nil
}

// fetchNotifierOldAccount fetches mongo account of a given new account id.
// this function must be used under another file for further use
func fetchNotifierOldAccount(accountId int64) (*mongomodels.Account, error) {
	newAccount := models.NewAccount()
	if err := newAccount.ById(accountId); err != nil {
		return nil, err
	}

	account, err := modelhelper.GetAccountById(newAccount.OldId)
	if err != nil {
		if err == mgo.ErrNotFound {
			return nil, errors.New("old account not found")
		}

		return nil, err
	}

	return account, nil
}

func (f *RealtimeWorkerController) sendInstanceEvent(instanceId int64, message interface{}, eventName string) error {
	channel, err := f.rmqConn.Channel()
	if err != nil {
		return err
	}
	defer channel.Close()

	routingKey := "oid." + strconv.FormatInt(instanceId, 10) + ".event." + eventName

	updateMessage, err := json.Marshal(message)
	if err != nil {
		return err
	}

	updateArr := make([]string, 1)
	if eventName == "updateInstance" {
		updateArr[0] = fmt.Sprintf("{\"$set\":%s}", string(updateMessage))
	} else {
		updateArr[0] = string(updateMessage)
	}

	msg, err := json.Marshal(updateArr)
	if err != nil {
		return err
	}

	return channel.Publish(
		"updateInstances", // exchange name
		routingKey,        // routing key
		false,             // mandatory
		false,             // immediate
		amqp.Publishing{Body: msg}, // message
	)
}

func (f *RealtimeWorkerController) sendChannelEvent(cml *models.ChannelMessageList, eventName string) error {
	channel, err := f.rmqConn.Channel()
	if err != nil {
		return err
	}
	defer channel.Close()

	secretNames, err := fetchSecretNames(cml.ChannelId)
	if err != nil {
		return err
	}

	// if we dont have any secret names, just return
	if len(secretNames) < 1 {
		f.log.Info("Channel %d doest have any secret name", cml.ChannelId)
		return nil
	}

	cm := models.NewChannelMessage()
	if err := cm.ById(cml.MessageId); err != nil {
		return err
	}

	cmc, err := cm.BuildEmptyMessageContainer()
	if err != nil {
		return err
	}

	byteMessage, err := json.Marshal(cmc)
	if err != nil {
		return err
	}

	for _, secretName := range secretNames {
		routingKey := "socialapi.channelsecret." + secretName + "." + eventName

		if err := channel.Publish(
			"broker",   // exchange name
			routingKey, // routing key
			false,      // mandatory
			false,      // immediate
			amqp.Publishing{Body: byteMessage}, // message
		); err != nil {
			return err
		}
	}
	return nil
}

func fetchSecretNames(channelId int64) ([]string, error) {
	names := make([]string, 0)
	c, err := fetchChannel(channelId)
	if err != nil {
		return names, err
	}

	name := fmt.Sprintf(
		"socialapi-group-%s-type-%s-name-%s",
		c.GroupName,
		c.TypeConstant,
		c.Name,
	)

	names, err = modelhelper.FetchFlattenedSecretName(name)
	return names, nil
}

func fetchChannel(channelId int64) (*models.Channel, error) {
	c := models.NewChannel()
	if err := c.ById(channelId); err != nil {
		return nil, err
	}
	return c, nil
}

func (f *RealtimeWorkerController) sendNotification(accountId int64, eventName string, data interface{}) error {
	channel, err := f.rmqConn.Channel()
	if err != nil {
		return err
	}
	defer channel.Close()

	oldAccount, err := modelhelper.GetAccountBySocialApiId(accountId)
	if err != nil {
		return err
	}

	notification := map[string]interface{}{
		"event":    eventName,
		"contents": data,
	}

	byteNotification, err := json.Marshal(notification)
	if err != nil {
		return err
	}

	return channel.Publish(
		"notification",
		oldAccount.Profile.Nickname, // this is routing key
		false,
		false,
		amqp.Publishing{Body: byteNotification},
	)
}
<commit_msg>Notification: when realtime notifier could not connect related an account related with new Account Id instead of returning error it just logs with Warning level<commit_after>package realtime

import (
	"encoding/json"
	"errors"
	"fmt"
	mongomodels "koding/db/models"
	"koding/db/mongodb/modelhelper"
	"socialapi/models"
	"socialapi/workers/helper"
	notificationmodels "socialapi/workers/notification/models"
	"strconv"

	"github.com/koding/logging"
	"github.com/koding/rabbitmq"
	"github.com/koding/worker"
	"github.com/streadway/amqp"
	"labix.org/v2/mgo"
)

type Action func(*RealtimeWorkerController, []byte) error

type RealtimeWorkerController struct {
	routes  map[string]Action
	log     logging.Logger
	rmqConn *amqp.Connection
}

type NotificationEvent struct {
	RoutingKey string              `json:"routingKey"`
	Content    NotificationContent `json:"contents"`
}

type NotificationContent struct {
	TypeConstant string `json:"type"`
	TargetId     int64  `json:"targetId"`
	ActorId      string `json:"actorId"`
}

func (r *RealtimeWorkerController) DefaultErrHandler(delivery amqp.Delivery, err error) bool {
	r.log.Error("an error occured deleting realtime event", err)
	delivery.Ack(false)
	return false
}

func NewRealtimeWorkerController(rmq *rabbitmq.RabbitMQ, log logging.Logger) (*RealtimeWorkerController, error) {
	rmqConn, err := rmq.Connect("NewRealtimeWorkerController")
	if err != nil {
		return nil, err
	}

	ffc := &RealtimeWorkerController{
		log:     log,
		rmqConn: rmqConn.Conn(),
	}

	routes := map[string]Action{
		"api.channel_message_created": (*RealtimeWorkerController).MessageSaved,
		"api.channel_message_updated": (*RealtimeWorkerController).MessageUpdated,
		"api.channel_message_deleted": (*RealtimeWorkerController).MessageDeleted,

		"api.interaction_created": (*RealtimeWorkerController).InteractionSaved,
		"api.interaction_deleted": (*RealtimeWorkerController).InteractionDeleted,

		"api.message_reply_created": (*RealtimeWorkerController).MessageReplySaved,
		"api.message_reply_deleted": (*RealtimeWorkerController).MessageReplyDeleted,

		"api.channel_message_list_created": (*RealtimeWorkerController).MessageListSaved,
		"api.channel_message_list_updated": (*RealtimeWorkerController).MessageListUpdated,
		"api.channel_message_list_deleted": (*RealtimeWorkerController).MessageListDeleted,

		"api.channel_participant_created": (*RealtimeWorkerController).ChannelParticipantAdded,
		"api.channel_participant_deleted": (*RealtimeWorkerController).ChannelParticipantRemoved,

		"notification.notification_created": (*RealtimeWorkerController).NotifyUser,
		"notification.notification_updated": (*RealtimeWorkerController).NotifyUser,
	}

	ffc.routes = routes

	return ffc, nil
}

func (f *RealtimeWorkerController) HandleEvent(event string, data []byte) error {
	f.log.Debug("New Event Received %s", event)
	handler, ok := f.routes[event]
	if !ok {
		return worker.HandlerNotFoundErr
	}

	return handler(f, data)
}

// no operation for message save for now
func (f *RealtimeWorkerController) MessageSaved(data []byte) error {
	return nil
}

// no operation for message delete for now
// channel_message_delete will handle message deletions from the
func (f *RealtimeWorkerController) MessageDeleted(data []byte) error {
	return nil
}

func (f *RealtimeWorkerController) MessageUpdated(data []byte) error {
	cm, err := helper.MapToChannelMessage(data)
	if err != nil {
		return err
	}

	err = f.sendInstanceEvent(cm.GetId(), cm, "updateInstance")
	if err != nil {
		fmt.Println(err)
		return err
	}

	return nil
}

func (f *RealtimeWorkerController) ChannelParticipantAdded(data []byte) error {
	return f.handleChannelParticipantEvent("AddedToChannel", data)
}

func (f *RealtimeWorkerController) ChannelParticipantRemoved(data []byte) error {
	return f.handleChannelParticipantEvent("RemovedFromChannel", data)
}

func (f *RealtimeWorkerController) handleChannelParticipantEvent(eventName string, data []byte) error {
	cp := models.NewChannelParticipant()
	if err := json.Unmarshal(data, cp); err != nil {
		return err
	}

	c := models.NewChannel()
	if err := c.ById(cp.ChannelId); err != nil {
		return err
	}

	return f.sendNotification(cp.AccountId, eventName, c)
}

func (f *RealtimeWorkerController) InteractionSaved(data []byte) error {
	return f.handleInteractionEvent("InteractionAdded", data)
}

func (f *RealtimeWorkerController) InteractionDeleted(data []byte) error {
	return f.handleInteractionEvent("InteractionRemoved", data)
}

func (f *RealtimeWorkerController) handleInteractionEvent(eventName string, data []byte) error {
	i, err := helper.MapToInteraction(data)
	if err != nil {
		return err
	}

	count, err := i.Count(i.TypeConstant)
	if err != nil {
		return err
	}

	oldId, err := models.AccountOldIdById(i.AccountId)
	if err != nil {
		return err
	}

	res := map[string]interface{}{
		"messageId":    i.MessageId,
		"accountId":    i.AccountId,
		"accountOldId": oldId,
		"typeConstant": i.TypeConstant,
		"count":        count,
	}

	err = f.sendInstanceEvent(i.MessageId, res, eventName)
	if err != nil {
		fmt.Println(err)
		return err
	}

	return nil
}

func (f *RealtimeWorkerController) MessageReplySaved(data []byte) error {
	i, err := helper.MapToMessageReply(data)
	if err != nil {
		return err
	}

	reply := models.NewChannelMessage()
	if err := reply.ById(i.ReplyId); err != nil {
		return err
	}

	cmc, err := reply.BuildEmptyMessageContainer()
	if err != nil {
		return err
	}

	err = f.sendInstanceEvent(i.MessageId, cmc, "ReplyAdded")
	if err != nil {
		fmt.Println(err)
		return err
	}

	return nil
}

func (f *RealtimeWorkerController) MessageReplyDeleted(data []byte) error {
	i, err := helper.MapToMessageReply(data)
	if err != nil {
		return err
	}

	err = f.sendInstanceEvent(i.MessageId, i, "ReplyRemoved")
	if err != nil {
		fmt.Println(err)
		return err
	}

	return nil
}

// send message to the channel
func (f *RealtimeWorkerController) MessageListSaved(data []byte) error {
	cml, err := helper.MapToChannelMessageList(data)
	if err != nil {
		return err
	}

	err = f.sendChannelEvent(cml, "MessageAdded")
	if err != nil {
		return err
	}

	return nil
}

// no operation for channel_message_list_updated event
func (f *RealtimeWorkerController) MessageListUpdated(data []byte) error {
	return nil
}

func (f *RealtimeWorkerController) MessageListDeleted(data []byte) error {
	cml, err := helper.MapToChannelMessageList(data)
	if err != nil {
		return err
	}

	err = f.sendChannelEvent(cml, "MessageRemoved")
	if err != nil {
		return err
	}
	return nil
}

func (f *RealtimeWorkerController) NotifyUser(data []byte) error {
	channel, err := f.rmqConn.Channel()
	if err != nil {
		return errors.New("channel connection error")
	}
	defer channel.Close()

	notification := notificationmodels.NewNotification()
	if err := notification.MapMessage(data); err != nil {
		return err
	}

	activity, nc, err := notification.FetchLastActivity()
	if err != nil {
		return err
	}

	// do not notify actor for her own action
	if activity.ActorId == notification.AccountId {
		return nil
	}

	// do not notify user when notification is not yet activated
	if notification.ActivatedAt.IsZero() {
		return nil
	}

	oldAccount, err := fetchNotifierOldAccount(notification.AccountId)
	if err != nil {
		f.log.Warning("an error occurred while fetching old account: %s", err)
		return nil
	}

	// fetch user profile name from bongo as routing key
	ne := &NotificationEvent{}

	ne.Content = NotificationContent{
		TargetId:     nc.TargetId,
		TypeConstant: nc.TypeConstant,
	}
	ne.Content.ActorId, _ = models.AccountOldIdById(activity.ActorId)

	notificationMessage, err := json.Marshal(ne)
	if err != nil {
		return err
	}

	routingKey := oldAccount.Profile.Nickname

	err = channel.Publish(
		"notification",
		routingKey,
		false,
		false,
		amqp.Publishing{Body: notificationMessage},
	)
	if err != nil {
		return fmt.Errorf("an error occurred while notifying user: %s", err)
	}

	return nil
}

// fetchNotifierOldAccount fetches mongo account of a given new account id.
// this function must be used under another file for further use
func fetchNotifierOldAccount(accountId int64) (*mongomodels.Account, error) {
	newAccount := models.NewAccount()
	if err := newAccount.ById(accountId); err != nil {
		return nil, err
	}

	account, err := modelhelper.GetAccountById(newAccount.OldId)
	if err != nil {
		if err == mgo.ErrNotFound {
			return nil, errors.New("old account not found")
		}

		return nil, err
	}

	return account, nil
}

func (f *RealtimeWorkerController) sendInstanceEvent(instanceId int64, message interface{}, eventName string) error {
	channel, err := f.rmqConn.Channel()
	if err != nil {
		return err
	}
	defer channel.Close()

	routingKey := "oid." + strconv.FormatInt(instanceId, 10) + ".event." + eventName

	updateMessage, err := json.Marshal(message)
	if err != nil {
		return err
	}

	updateArr := make([]string, 1)
	if eventName == "updateInstance" {
		updateArr[0] = fmt.Sprintf("{\"$set\":%s}", string(updateMessage))
	} else {
		updateArr[0] = string(updateMessage)
	}

	msg, err := json.Marshal(updateArr)
	if err != nil {
		return err
	}

	return channel.Publish(
		"updateInstances", // exchange name
		routingKey,        // routing key
		false,             // mandatory
		false,             // immediate
		amqp.Publishing{Body: msg}, // message
	)
}

func (f *RealtimeWorkerController) sendChannelEvent(cml *models.ChannelMessageList, eventName string) error {
	channel, err := f.rmqConn.Channel()
	if err != nil {
		return err
	}
	defer channel.Close()

	secretNames, err := fetchSecretNames(cml.ChannelId)
	if err != nil {
		return err
	}

	// if we dont have any secret names, just return
	if len(secretNames) < 1 {
		f.log.Info("Channel %d doest have any secret name", cml.ChannelId)
		return nil
	}

	cm := models.NewChannelMessage()
	if err := cm.ById(cml.MessageId); err != nil {
		return err
	}

	cmc, err := cm.BuildEmptyMessageContainer()
	if err != nil {
		return err
	}

	byteMessage, err := json.Marshal(cmc)
	if err != nil {
		return err
	}

	for _, secretName := range secretNames {
		routingKey := "socialapi.channelsecret." + secretName + "." + eventName

		if err := channel.Publish(
			"broker",   // exchange name
			routingKey, // routing key
			false,      // mandatory
			false,      // immediate
			amqp.Publishing{Body: byteMessage}, // message
		); err != nil {
			return err
		}
	}
	return nil
}

func fetchSecretNames(channelId int64) ([]string, error) {
	names := make([]string, 0)
	c, err := fetchChannel(channelId)
	if err != nil {
		return names, err
	}

	name := fmt.Sprintf(
		"socialapi-group-%s-type-%s-name-%s",
		c.GroupName,
		c.TypeConstant,
		c.Name,
	)

	names, err = modelhelper.FetchFlattenedSecretName(name)
	return names, nil
}

func fetchChannel(channelId int64) (*models.Channel, error) {
	c := models.NewChannel()
	if err := c.ById(channelId); err != nil {
		return nil, err
	}
	return c, nil
}

func (f *RealtimeWorkerController) sendNotification(accountId int64, eventName string, data interface{}) error {
	channel, err := f.rmqConn.Channel()
	if err != nil {
		return err
	}
	defer channel.Close()

	oldAccount, err := modelhelper.GetAccountBySocialApiId(accountId)
	if err != nil {
		return err
	}

	notification := map[string]interface{}{
		"event":    eventName,
		"contents": data,
	}

	byteNotification, err := json.Marshal(notification)
	if err != nil {
		return err
	}

	return channel.Publish(
		"notification",
		oldAccount.Profile.Nickname, // this is routing key
		false,
		false,
		amqp.Publishing{Body: byteNotification},
	)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"
)

type Response map[string]interface{}

func TestPostLogin(t *testing.T) {
	var expected, response Response

	expected = Response{"error": "Must send content as text/json"}

	r, _ := http.NewRequest("POST", "/login", nil)
	w := httptest.NewRecorder()

	LoginHandler(w, r)

	if w.Code != http.StatusBadRequest {
		t.Errorf("Expected status %d got %v", http.StatusBadRequest, w.Code)
	}

	b, err := ioutil.ReadAll(w.Body)
	if err != nil {
		t.Error("Can't read the response body")
	}

	err = json.Unmarshal(b, &response)
	if err != nil {
		t.Error("Can't read the response json")
	}

	for key, v := range expected {
		if _, ok := response[key]; !ok {
			t.Errorf("Expected %s:%v but key is missing in %v", key, v, len(response))
		} else if response[key] != v {
			t.Errorf("Expected %s:%v got %v", key, v, response[key])
		}
	}

}
<commit_msg>added body to the POST request<commit_after>package main

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

type Response map[string]interface{}

func TestInvalidHeaderLogin(t *testing.T) {
	var expected, response Response

	expected = Response{"error": "Must send content as text/json"}

	r, _ := http.NewRequest("POST", "/login", strings.NewReader("TeST"))
	w := httptest.NewRecorder()

	LoginHandler(w, r)

	if w.Code != http.StatusBadRequest {
		t.Errorf("Expected status %d got %v", http.StatusBadRequest, w.Code)
	}

	b, err := ioutil.ReadAll(w.Body)
	if err != nil {
		t.Error("Can't read the response body")
	}

	err = json.Unmarshal(b, &response)
	if err != nil {
		t.Error("Can't read the response json")
	}

	for key, v := range expected {
		if _, ok := response[key]; !ok {
			t.Errorf("Expected %s:%v but key is missing in %v", key, v, len(response))
		} else if response[key] != v {
			t.Errorf("Expected %s:%v got %v", key, v, response[key])
		}
	}

}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"io"
	"math/rand"
	"net"
	"os"
	"strconv"
	"testing"
	"time"

	"github.com/xindong/frontd/aes256cbc"
	"github.com/xindong/frontd/reuse"
)

var (
	_echoServerAddr      = []byte("127.0.0.1:62863")
	_expectAESCiphertext = []byte("U2FsdGVkX19KIJ9OQJKT/yHGMrS+5SsBAAjetomptQ0=")
	_secret              = []byte("p0S8rX680*48")
	_defaultFrontdAddr   = "127.0.0.1:" + strconv.Itoa(_DefaultPort)
)

var (
	// use -reuse with go test enable SO_REUSEPORT
	// go test -parallel 6553 -benchtime 60s -bench BenchmarkEchoParallel -reuse
	// but it seems will not working with single backend addr because of
	// http://stackoverflow.com/questions/14388706/socket-options-so-reuseaddr-and-so-reuseport-how-do-they-differ-do-they-mean-t
	reuseTest = flag.Bool("reuse", false, "test reuseport dialer")
)

func TestMain(m *testing.M) {
	flag.Parse()
	if *reuseTest {
		fmt.Println("testing SO_REUSEPORT")
	}

	// start echo server
	go servEcho()

	// start listen
	os.Setenv("SECRET", string(_secret))

	go main()

	rand.Seed(time.Now().UnixNano())

	// TODO: better way to wait for server to start
	time.Sleep(time.Second)
	os.Exit(m.Run())
}

func servEcho() {
	l, err := net.Listen("tcp", string(_echoServerAddr))
	if err != nil {
		fmt.Println("Error listening:", err.Error())
		os.Exit(1)
	}
	// Close the listener when the application closes.
	defer l.Close()
	fmt.Println("Listening on " + string(_echoServerAddr))
	for {
		// Listen for an incoming connection.
		c, err := l.Accept()
		if err != nil {
			fmt.Println("Error accepting: ", err.Error())
			os.Exit(1)
		}
		// Handle connections in a new goroutine.
		go func(c net.Conn) {
			defer c.Close()

			_, err := io.Copy(c, c)
			switch err {
			case io.EOF:
				err = nil
				return
			case nil:
				return
			}
			panic(err)
		}(c)
	}
}

func TestTextDecryptAES(t *testing.T) {
	o := aes256cbc.New()

	dec, err := o.Decrypt(_secret, _expectAESCiphertext)
	if err != nil {
		panic(err)
	}
	if !bytes.Equal(dec, _echoServerAddr) {
		panic(errors.New("not match"))
	}
}

func encryptText(plaintext, passphrase []byte) ([]byte, error) {
	o := aes256cbc.New()

	return o.Encrypt(passphrase, plaintext)
}

func randomBytes(n int) []byte {

	b := make([]byte, n)
	// A src.Int63() generates 63 random bits, enough for letterIdxMax characters!
	for i := 0; i < n; i++ {
		b[i] = byte(rand.Int())
	}

	return b
}

func testEchoRound(conn net.Conn) {
	conn.SetDeadline(time.Now().Add(time.Second * 10))

	n := rand.Int()%2048 + 10
	out := randomBytes(n)
	n0, err := conn.Write(out)
	if err != nil {
		panic(err)
	}

	rcv := make([]byte, n)
	n1, err := io.ReadFull(conn, rcv)
	if err != nil && err != io.EOF {
		panic(err)
	}
	if !bytes.Equal(out[:n0], rcv[:n1]) {
		fmt.Println("out: ", n0, "in:", n1)

		fmt.Println("out: ", hex.EncodeToString(out), "in:", hex.EncodeToString(rcv))
		panic(errors.New("echo server reply is not match"))
	}
}

func TestEchoServer(t *testing.T) {
	var conn net.Conn
	var err error
	if *reuseTest {
		conn, err = reuseport.Dial("tcp", "127.0.0.1:0", string(_echoServerAddr))
	} else {
		conn, err = net.Dial("tcp", string(_echoServerAddr))
	}
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	n := rand.Int() % 10
	for i := 0; i < n; i++ {
		testEchoRound(conn)
	}
}

func testProtocol(cipherAddr []byte) {
	// * test decryption
	var conn net.Conn
	var err error
	if *reuseTest {
		conn, err = reuseport.Dial("tcp", "127.0.0.1:0", _defaultFrontdAddr)
	} else {
		conn, err = net.Dial("tcp", _defaultFrontdAddr)
	}

	if err != nil {
		panic(err)
	}
	defer conn.Close()

	_, err = conn.Write(cipherAddr)
	if err != nil {
		panic(err)
	}

	_, err = conn.Write([]byte("\n"))
	if err != nil {
		panic(err)
	}

	for i := 0; i < 5; i++ {
		testEchoRound(conn)
	}
}

func TestProtocolDecrypt(*testing.T) {
	b, err := encryptText(_echoServerAddr, _secret)
	if err != nil {
		panic(err)
	}
	testProtocol(b)
}

// TODO: test decryption with extra bytes in packet and check data

// TODO: test decryption with seperated packet simulate loss connection and check data

// TODO: benchmark 100, 1000 connect with 1k 10k 100k 1m data

func BenchmarkEncryptText(b *testing.B) {
	s1 := randomBytes(255)
	s2 := randomBytes(32)
	for i := 0; i < b.N; i++ {
		_, err := encryptText(s1, s2)
		if err != nil {
			panic(err)
		}
	}
}

func BenchmarkDecryptText(b *testing.B) {
	for i := 0; i < b.N; i++ {
		o := aes256cbc.New()
		_, err := o.Decrypt(_secret, _expectAESCiphertext)
		if err != nil {
			panic(err)
		}
	}
}

func BenchmarkEcho(b *testing.B) {
	for i := 0; i < b.N; i++ {
		TestEchoServer(&testing.T{})
	}
}

func BenchmarkLatency(b *testing.B) {
	cipherAddr, err := encryptText(_echoServerAddr, _secret)
	if err != nil {
		panic(err)
	}

	for i := 0; i < b.N; i++ {
		testProtocol(cipherAddr)
	}
}

func BenchmarkNoHitLatency(b *testing.B) {
	for i := 0; i < b.N; i++ {
		TestProtocolDecrypt(&testing.T{})
	}
}

func BenchmarkEchoParallel(b *testing.B) {
	if !reuseport.Available() {
		return
	}
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			TestEchoServer(&testing.T{})
		}
	})
}

// with echo server with random hanging
// * benchmark latency
// * benchmark throughput
// * benchmark copy-on-write performance BackendAddrCache
// * benchmark memory footprint
<commit_msg>add parallel test cases<commit_after>package main

import (
	"bytes"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"io"
	"math/rand"
	"net"
	"os"
	"strconv"
	"testing"
	"time"

	"github.com/xindong/frontd/aes256cbc"
	"github.com/xindong/frontd/reuse"
)

var (
	_echoServerAddr      = []byte("127.0.0.1:62863")
	_expectAESCiphertext = []byte("U2FsdGVkX19KIJ9OQJKT/yHGMrS+5SsBAAjetomptQ0=")
	_secret              = []byte("p0S8rX680*48")
	_defaultFrontdAddr   = "127.0.0.1:" + strconv.Itoa(_DefaultPort)
)

var (
	// use -reuse with go test enable SO_REUSEPORT
	// go test -parallel 6553 -benchtime 60s -bench BenchmarkEchoParallel -reuse
	// but it seems will not working with single backend addr because of
	// http://stackoverflow.com/questions/14388706/socket-options-so-reuseaddr-and-so-reuseport-how-do-they-differ-do-they-mean-t
	reuseTest = flag.Bool("reuse", false, "test reuseport dialer")
)

func TestMain(m *testing.M) {
	flag.Parse()
	if *reuseTest {
		fmt.Println("testing SO_REUSEPORT")
	}

	// start echo server
	go servEcho()

	// start listen
	os.Setenv("SECRET", string(_secret))

	go main()

	rand.Seed(time.Now().UnixNano())

	// TODO: better way to wait for server to start
	time.Sleep(time.Second)
	os.Exit(m.Run())
}

func servEcho() {
	l, err := net.Listen("tcp", string(_echoServerAddr))
	if err != nil {
		fmt.Println("Error listening:", err.Error())
		os.Exit(1)
	}
	// Close the listener when the application closes.
	defer l.Close()
	fmt.Println("Listening on " + string(_echoServerAddr))
	for {
		// Listen for an incoming connection.
		c, err := l.Accept()
		if err != nil {
			fmt.Println("Error accepting: ", err.Error())
			os.Exit(1)
		}
		// Handle connections in a new goroutine.
		go func(c net.Conn) {
			defer c.Close()

			_, err := io.Copy(c, c)
			switch err {
			case io.EOF:
				err = nil
				return
			case nil:
				return
			}
			panic(err)
		}(c)
	}
}

func TestTextDecryptAES(t *testing.T) {
	o := aes256cbc.New()

	dec, err := o.Decrypt(_secret, _expectAESCiphertext)
	if err != nil {
		panic(err)
	}
	if !bytes.Equal(dec, _echoServerAddr) {
		panic(errors.New("not match"))
	}
}

func encryptText(plaintext, passphrase []byte) ([]byte, error) {
	o := aes256cbc.New()

	return o.Encrypt(passphrase, plaintext)
}

func randomBytes(n int) []byte {

	b := make([]byte, n)
	// A src.Int63() generates 63 random bits, enough for letterIdxMax characters!
	for i := 0; i < n; i++ {
		b[i] = byte(rand.Int())
	}

	return b
}

func testEchoRound(conn net.Conn) {
	conn.SetDeadline(time.Now().Add(time.Second * 10))

	n := rand.Int()%2048 + 10
	out := randomBytes(n)
	n0, err := conn.Write(out)
	if err != nil {
		panic(err)
	}

	rcv := make([]byte, n)
	n1, err := io.ReadFull(conn, rcv)
	if err != nil && err != io.EOF {
		panic(err)
	}
	if !bytes.Equal(out[:n0], rcv[:n1]) {
		fmt.Println("out: ", n0, "in:", n1)

		fmt.Println("out: ", hex.EncodeToString(out), "in:", hex.EncodeToString(rcv))
		panic(errors.New("echo server reply is not match"))
	}
}

func TestEchoServer(t *testing.T) {
	var conn net.Conn
	var err error
	if *reuseTest {
		conn, err = reuseport.Dial("tcp", "127.0.0.1:0", string(_echoServerAddr))
	} else {
		conn, err = net.Dial("tcp", string(_echoServerAddr))
	}
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	n := rand.Int() % 10
	for i := 0; i < n; i++ {
		testEchoRound(conn)
	}
}

func testProtocol(cipherAddr []byte) {
	// * test decryption
	var conn net.Conn
	var err error
	if *reuseTest {
		conn, err = reuseport.Dial("tcp", "127.0.0.1:0", _defaultFrontdAddr)
	} else {
		conn, err = net.Dial("tcp", _defaultFrontdAddr)
	}

	if err != nil {
		panic(err)
	}
	defer conn.Close()

	_, err = conn.Write(cipherAddr)
	if err != nil {
		panic(err)
	}

	_, err = conn.Write([]byte("\n"))
	if err != nil {
		panic(err)
	}

	for i := 0; i < 5; i++ {
		testEchoRound(conn)
	}
}

func TestProtocolDecrypt(*testing.T) {
	b, err := encryptText(_echoServerAddr, _secret)
	if err != nil {
		panic(err)
	}
	testProtocol(b)
}

// TODO: test decryption with extra bytes in packet and check data

// TODO: test decryption with seperated packet simulate loss connection and check data

// TODO: benchmark 100, 1000 connect with 1k 10k 100k 1m data

func BenchmarkEncryptText(b *testing.B) {
	s1 := randomBytes(255)
	s2 := randomBytes(32)
	for i := 0; i < b.N; i++ {
		_, err := encryptText(s1, s2)
		if err != nil {
			panic(err)
		}
	}
}

func BenchmarkDecryptText(b *testing.B) {
	for i := 0; i < b.N; i++ {
		o := aes256cbc.New()
		_, err := o.Decrypt(_secret, _expectAESCiphertext)
		if err != nil {
			panic(err)
		}
	}
}

func BenchmarkEcho(b *testing.B) {
	for i := 0; i < b.N; i++ {
		TestEchoServer(&testing.T{})
	}
}

func BenchmarkLatency(b *testing.B) {
	cipherAddr, err := encryptText(_echoServerAddr, _secret)
	if err != nil {
		panic(err)
	}

	for i := 0; i < b.N; i++ {
		testProtocol(cipherAddr)
	}
}

func BenchmarkNoHitLatency(b *testing.B) {
	for i := 0; i < b.N; i++ {
		TestProtocolDecrypt(&testing.T{})
	}
}

func BenchmarkEchoParallel(b *testing.B) {

	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			TestEchoServer(&testing.T{})
		}
	})
}

func BenchmarkLatencyParallel(b *testing.B) {
	cipherAddr, err := encryptText(_echoServerAddr, _secret)
	if err != nil {
		panic(err)
	}
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			testProtocol(cipherAddr)
		}
	})
}

func BenchmarkNoHitLatencyParallel(b *testing.B) {
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			TestProtocolDecrypt(&testing.T{})
		}
	})
}

// with echo server with random hanging
// * benchmark latency
// * benchmark throughput
// * benchmark copy-on-write performance BackendAddrCache
// * benchmark memory footprint
<|endoftext|>