<commit_before>package utils

import (
	"bytes"
	"crypto/md5"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"math/rand"
	"mime/multipart"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/eirka/eirka-libs/amazon"
	"github.com/eirka/eirka-libs/config"
	"github.com/eirka/eirka-libs/db"
	_ "github.com/eirka/eirka-libs/errors"

	local "github.com/eirka/eirka-post/config"
)

// valid file extensions
var validExt = map[string]bool{
	".jpg":  true,
	".jpeg": true,
	".png":  true,
	".gif":  true,
	".webm": true,
}

type ImageType struct {
	File        multipart.File
	Header      *multipart.FileHeader
	Ib          uint
	Filename    string
	Thumbnail   string
	Filepath    string
	Thumbpath   string
	Ext         string
	MD5         string
	OrigWidth   int
	OrigHeight  int
	ThumbWidth  int
	ThumbHeight int
	image       *bytes.Buffer
	mime        string
	duration    int
	video       bool
	avatar      bool
}

func (i *ImageType) IsValid() bool {

	if i.Filename == "" {
		return false
	}

	if i.Filepath == "" {
		return false
	}

	if i.Thumbnail == "" {
		return false
	}

	if i.Thumbpath == "" {
		return false
	}

	if i.Ib == 0 {
		return false
	}

	if i.Ext == "" {
		return false
	}

	if i.MD5 == "" {
		return false
	}

	if i.mime == "" {
		return false
	}

	return true
}

func (i *ImageType) IsValidPost() bool {
	if i.OrigWidth == 0 {
		return false
	}

	if i.OrigHeight == 0 {
		return false
	}

	if i.ThumbWidth == 0 {
		return false
	}

	if i.ThumbHeight == 0 {
		return false
	}

	return true
}

// save an image file
func (i *ImageType) SaveImage() (err error) {

	// check given file ext
	err = i.checkReqExt()
	if err != nil {
		return
	}

	// get file md5
	err = i.getMD5()
	if err != nil {
		return
	}

	// check to see if the file already exists
	err = i.checkDuplicate()
	if err != nil {
		return
	}

	// check file magic sig
	err = i.checkMagic()
	if err != nil {
		return
	}

	// check image stats
	err = i.getStats()
	if err != nil {
		return
	}

	// save the file to disk
	err = i.saveFile()
	if err != nil {
		return
	}

	// process a webm
	if i.video {
		// check the webm info
		err = i.checkWebM()
		if err != nil {
			return
		}

		// create thumbnail from webm
		err = i.createWebMThumbnail()
		if err != nil {
			return
		}

	}

	// create a thumbnail
	err = i.createThumbnail(config.Settings.Limits.ThumbnailMaxWidth, config.Settings.Limits.ThumbnailMaxHeight)
	if err != nil {
		return
	}

	// copy the file to s3
	err = i.copyToS3()
	if err != nil {
		return
	}

	// check final state
	if !i.IsValidPost() {
		return errors.New("ImageType is not valid")
	}

	return

}

// Get file extension from request header
func (i *ImageType) checkReqExt() (err error) {
	// Get ext from request header
	name := i.Header.Filename
	ext := filepath.Ext(name)

	if ext == "" {
		return errors.New("No file extension")
	}

	// Check to see if extension is allowed
	if !isAllowedExt(ext) {
		return errors.New("Format not supported")
	}

	return

}

// Check if file ext allowed
func isAllowedExt(ext string) bool {
	return validExt[strings.ToLower(ext)]
}

// Get image MD5 and write file into buffer
func (i *ImageType) getMD5() (err error) {

	defer i.File.Close()

	hasher := md5.New()

	i.image = new(bytes.Buffer)

	// Save file and also read into hasher for md5
	_, err = io.Copy(i.image, io.TeeReader(i.File, hasher))
	if err != nil {
		return errors.New("Problem copying file")
	}

	// Set md5sum from hasher
	i.MD5 = hex.EncodeToString(hasher.Sum(nil))

	return

}

// check if the md5 is already in the database
func (i *ImageType) checkDuplicate() (err error) {

	// Get Database handle
	dbase, err := db.GetDb()
	if err != nil {
		return
	}

	if i.Ib == 0 {
		return errors.New("No imageboard set on duplicate check")
	}

	var check bool
	var thread, post sql.NullInt64

	err = dbase.QueryRow(`select count(1),posts.post_num,threads.thread_id from threads
	LEFT JOIN posts on threads.thread_id = posts.thread_id
	LEFT JOIN images on posts.post_id = images.post_id
	WHERE image_hash = ? AND ib_id = ?`, i.MD5, i.Ib).Scan(&check, &post, &thread)
	if err != nil {
		return
	}

	// return error if it exists
	if check {
		return fmt.Errorf("Image has already been posted. Thread: %d Post: %d", thread.Int64, post.Int64)
	}

	return
}

func (i *ImageType) checkMagic() (err error) {

	// detect the mime type
	i.mime = http.DetectContentType(i.image.Bytes())

	switch i.mime {
	case "image/png":
		i.Ext = ".png"
	case "image/jpeg":
		i.Ext = ".jpg"
	case "image/gif":
		i.Ext = ".gif"
	case "video/webm":
		i.Ext = ".webm"
		i.video = true
	default:
		return errors.New("Unknown file type")
	}

	// Check to see if extension is allowed
	if !isAllowedExt(i.Ext) {
		return errors.New("Format not supported")
	}

	return

}

func (i *ImageType) getStats() (err error) {

	// skip if its a video since we cant decode it
	if i.video {
		return
	}

	// decode image config
	img, _, err := image.DecodeConfig(bytes.NewReader(i.image.Bytes()))
	if err != nil {
		return errors.New("Problem decoding image")
	}

	// set original width
	i.OrigWidth = img.Width
	// set original height
	i.OrigHeight = img.Height

	// Check against maximum sizes
	switch {
	case i.OrigWidth > config.Settings.Limits.ImageMaxWidth:
		return errors.New("Image width too large")
	case img.Width < config.Settings.Limits.ImageMinWidth:
		return errors.New("Image width too small")
	case i.OrigHeight > config.Settings.Limits.ImageMaxHeight:
		return errors.New("Image height too large")
	case img.Height < config.Settings.Limits.ImageMinHeight:
		return errors.New("Image height too small")
	case i.image.Len() > config.Settings.Limits.ImageMaxSize:
		return errors.New("Image size too large")
	}

	return

}

func (i *ImageType) saveFile() (err error) {

	defer i.image.Reset()

	i.makeFilenames()

	// avatar filename is the users id
	if i.avatar {
		i.Thumbnail = fmt.Sprintf("%d.png", i.Ib)
		i.Thumbpath = filepath.Join(local.Settings.Directories.ThumbnailDir, i.Thumbnail)
	}

	if !i.IsValid() {
		return errors.New("ImageType is not valid")
	}

	imagefile := filepath.Join(local.Settings.Directories.ImageDir, i.Filename)

	image, err := os.Create(imagefile)
	if err != nil {
		return errors.New("Problem saving file")
	}
	defer image.Close()

	_, err = io.Copy(image, bytes.NewReader(i.image.Bytes()))
	if err != nil {
		return errors.New("Problem saving file")
	}

	return

}

// Make a random unix time filename
func (i *ImageType) makeFilenames() {

	// Create seed for random
	rand.Seed(time.Now().UnixNano())

	// Get random 3 digit int to append to unix time
	rand_t := rand.Intn(899) + 100

	// Get current unix time
	time_t := time.Now().Unix()

	// Append random int to unix time
	file_t := fmt.Sprintf("%d%d", time_t, rand_t)

	// Append ext to filename
	i.Filename = fmt.Sprintf("%s%s", file_t, i.Ext)

	// Append jpg to thumbnail name because it is always a jpg
	i.Thumbnail = fmt.Sprintf("%ss.jpg", file_t)

	// set the full file path
	i.Filepath = filepath.Join(local.Settings.Directories.ImageDir, i.Filename)

	// set the full thumbnail path
	i.Thumbpath = filepath.Join(local.Settings.Directories.ThumbnailDir, i.Thumbnail)

	return

}

func (i *ImageType) createThumbnail(maxwidth, maxheight int) (err error) {

	var imagef string

	if i.video {
		imagef = fmt.Sprintf("%s[0]", i.Thumbpath)
	} else {
		imagef = fmt.Sprintf("%s[0]", i.Filepath)
	}

	orig_dimensions := fmt.Sprintf("%dx%d", i.OrigWidth, i.OrigHeight)

	var args []string

	// different options for avatars
	if i.avatar {
		args = []string{
			"-size",
			orig_dimensions,
			imagef,
			"-background",
			"none",
			"-thumbnail",
			fmt.Sprintf("%dx%d^", maxwidth, maxheight),
			"-gravity",
			"center",
			"-extent",
			fmt.Sprintf("%dx%d", maxwidth, maxheight),
			i.Thumbpath,
		}
	} else {
		args = []string{
			"-background",
			"white",
			"-flatten",
			"-size",
			orig_dimensions,
			"-resize",
			fmt.Sprintf("%dx%d>", maxwidth, maxheight),
			"-quality",
			"90",
			imagef,
			i.Thumbpath,
		}
	}

	_, err = exec.Command("convert", args...).Output()
	if err != nil {
		return errors.New("Problem making thumbnail")
	}

	thumb, err := os.Open(i.Thumbpath)
	if err != nil {
		return errors.New("Problem making thumbnail")
	}
	defer thumb.Close()

	img, _, err := image.DecodeConfig(thumb)
	if err != nil {
		return errors.New("Problem decoding thumbnail")
	}

	i.ThumbWidth = img.Width
	i.ThumbHeight = img.Height

	return

}

func (i *ImageType) copyToS3() (err error) {

	s3 := amazon.New()

	err = s3.Save(i.Filepath, fmt.Sprintf("src/%s", i.Filename), i.mime, false)
	if err != nil {
		return
	}

	err = s3.Save(i.Thumbpath, fmt.Sprintf("thumb/%s", i.Thumbnail), "image/jpeg", false)
	if err != nil {
		return
	}

	return
}
<commit_msg>add random avatar when registering<commit_after>package utils

import (
	"bytes"
	"crypto/md5"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"math/rand"
	"mime/multipart"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/eirka/eirka-libs/amazon"
	"github.com/eirka/eirka-libs/config"
	"github.com/eirka/eirka-libs/db"
	_ "github.com/eirka/eirka-libs/errors"

	local "github.com/eirka/eirka-post/config"
)

// valid file extensions
var validExt = map[string]bool{
	".jpg":  true,
	".jpeg": true,
	".png":  true,
	".gif":  true,
	".webm": true,
}

type ImageType struct {
	File        multipart.File
	Header      *multipart.FileHeader
	Ib          uint
	Filename    string
	Thumbnail   string
	Filepath    string
	Thumbpath   string
	Ext         string
	MD5         string
	OrigWidth   int
	OrigHeight  int
	ThumbWidth  int
	ThumbHeight int
	image       *bytes.Buffer
	mime        string
	duration    int
	video       bool
	avatar      bool
}

func (i *ImageType) IsValid() bool {

	if i.Filename == "" {
		return false
	}

	if i.Filepath == "" {
		return false
	}

	if i.Thumbnail == "" {
		return false
	}

	if i.Thumbpath == "" {
		return false
	}

	if i.Ib == 0 {
		return false
	}

	if i.Ext == "" {
		return false
	}

	if i.MD5 == "" {
		return false
	}

	if i.mime == "" {
		return false
	}

	return true
}

func (i *ImageType) IsValidPost() bool {
	if i.OrigWidth == 0 {
		return false
	}

	if i.OrigHeight == 0 {
		return false
	}

	if i.ThumbWidth == 0 {
		return false
	}

	if i.ThumbHeight == 0 {
		return false
	}

	return true
}

// save an image file
func (i *ImageType) SaveImage() (err error) {

	// check given file ext
	err = i.checkReqExt()
	if err != nil {
		return
	}

	// get file md5
	err = i.getMD5()
	if err != nil {
		return
	}

	// check to see if the file already exists
	err = i.checkDuplicate()
	if err != nil {
		return
	}

	// check file magic sig
	err = i.checkMagic()
	if err != nil {
		return
	}

	// check image stats
	err = i.getStats()
	if err != nil {
		return
	}

	// save the file to disk
	err = i.saveFile()
	if err != nil {
		return
	}

	// process a webm
	if i.video {
		// check the webm info
		err = i.checkWebM()
		if err != nil {
			return
		}

		// create thumbnail from webm
		err = i.createWebMThumbnail()
		if err != nil {
			return
		}

	}

	// create a thumbnail
	err = i.createThumbnail(config.Settings.Limits.ThumbnailMaxWidth, config.Settings.Limits.ThumbnailMaxHeight)
	if err != nil {
		return
	}

	// copy the file to s3
	err = i.copyToS3()
	if err != nil {
		return
	}

	// check final state
	if !i.IsValidPost() {
		return errors.New("ImageType is not valid")
	}

	return

}

// Get file extension from request header
func (i *ImageType) checkReqExt() (err error) {
	// Get ext from request header
	name := i.Header.Filename
	ext := filepath.Ext(name)

	if ext == "" {
		return errors.New("No file extension")
	}

	// Check to see if extension is allowed
	if !isAllowedExt(ext) {
		return errors.New("Format not supported")
	}

	return

}

// Check if file ext allowed
func isAllowedExt(ext string) bool {
	return validExt[strings.ToLower(ext)]
}

// Get image MD5 and write file into buffer
func (i *ImageType) getMD5() (err error) {

	defer i.File.Close()

	hasher := md5.New()

	i.image = new(bytes.Buffer)

	// Save file and also read into hasher for md5
	_, err = io.Copy(i.image, io.TeeReader(i.File, hasher))
	if err != nil {
		return errors.New("Problem copying file")
	}

	// Set md5sum from hasher
	i.MD5 = hex.EncodeToString(hasher.Sum(nil))

	return

}

// check if the md5 is already in the database
func (i *ImageType) checkDuplicate() (err error) {

	// Get Database handle
	dbase, err := db.GetDb()
	if err != nil {
		return
	}

	if i.Ib == 0 {
		return errors.New("No imageboard set on duplicate check")
	}

	var check bool
	var thread, post sql.NullInt64

	err = dbase.QueryRow(`select count(1),posts.post_num,threads.thread_id from threads
	LEFT JOIN posts on threads.thread_id = posts.thread_id
	LEFT JOIN images on posts.post_id = images.post_id
	WHERE image_hash = ? AND ib_id = ?`, i.MD5, i.Ib).Scan(&check, &post, &thread)
	if err != nil {
		return
	}

	// return error if it exists
	if check {
		return fmt.Errorf("Image has already been posted. Thread: %d Post: %d", thread.Int64, post.Int64)
	}

	return
}

func (i *ImageType) checkMagic() (err error) {

	// detect the mime type
	i.mime = http.DetectContentType(i.image.Bytes())

	switch i.mime {
	case "image/png":
		i.Ext = ".png"
	case "image/jpeg":
		i.Ext = ".jpg"
	case "image/gif":
		i.Ext = ".gif"
	case "video/webm":
		i.Ext = ".webm"
		i.video = true
	default:
		return errors.New("Unknown file type")
	}

	// Check to see if extension is allowed
	if !isAllowedExt(i.Ext) {
		return errors.New("Format not supported")
	}

	return

}

func (i *ImageType) getStats() (err error) {

	// skip if its a video since we cant decode it
	if i.video {
		return
	}

	// decode image config
	img, _, err := image.DecodeConfig(bytes.NewReader(i.image.Bytes()))
	if err != nil {
		return errors.New("Problem decoding image")
	}

	// set original width
	i.OrigWidth = img.Width
	// set original height
	i.OrigHeight = img.Height

	// Check against maximum sizes
	switch {
	case i.OrigWidth > config.Settings.Limits.ImageMaxWidth:
		return errors.New("Image width too large")
	case img.Width < config.Settings.Limits.ImageMinWidth:
		return errors.New("Image width too small")
	case i.OrigHeight > config.Settings.Limits.ImageMaxHeight:
		return errors.New("Image height too large")
	case img.Height < config.Settings.Limits.ImageMinHeight:
		return errors.New("Image height too small")
	case i.image.Len() > config.Settings.Limits.ImageMaxSize:
		return errors.New("Image size too large")
	}

	return

}

func (i *ImageType) saveFile() (err error) {

	defer i.image.Reset()

	i.makeFilenames()

	// avatar filename is the users id
	if i.avatar {
		i.Thumbnail = fmt.Sprintf("%d.png", i.Ib)
		i.Thumbpath = filepath.Join(local.Settings.Directories.ThumbnailDir, i.Thumbnail)
	}

	if !i.IsValid() {
		return errors.New("ImageType is not valid")
	}

	imagefile := filepath.Join(local.Settings.Directories.ImageDir, i.Filename)

	image, err := os.Create(imagefile)
	if err != nil {
		return errors.New("Problem saving file")
	}
	defer image.Close()

	_, err = io.Copy(image, bytes.NewReader(i.image.Bytes()))
	if err != nil {
		return errors.New("Problem saving file")
	}

	return

}

// Make a random unix time filename
func (i *ImageType) makeFilenames() {

	// Create seed for random
	rand.Seed(time.Now().UnixNano())

	// Get random 3 digit int to append to unix time
	rand_t := rand.Intn(899) + 100

	// Get current unix time
	time_t := time.Now().Unix()

	// Append random int to unix time
	file_t := fmt.Sprintf("%d%d", time_t, rand_t)

	// Append ext to filename
	i.Filename = fmt.Sprintf("%s%s", file_t, i.Ext)

	// Append jpg to thumbnail name because it is always a jpg
	i.Thumbnail = fmt.Sprintf("%ss.jpg", file_t)

	// set the full file path
	i.Filepath = filepath.Join(local.Settings.Directories.ImageDir, i.Filename)

	// set the full thumbnail path
	i.Thumbpath = filepath.Join(local.Settings.Directories.ThumbnailDir, i.Thumbnail)

	return

}

func (i *ImageType) createThumbnail(maxwidth, maxheight int) (err error) {

	var imagef string

	if i.video {
		imagef = fmt.Sprintf("%s[0]", i.Thumbpath)
	} else {
		imagef = fmt.Sprintf("%s[0]", i.Filepath)
	}

	orig_dimensions := fmt.Sprintf("%dx%d", i.OrigWidth, i.OrigHeight)

	var args []string

	// different options for avatars
	if i.avatar {
		args = []string{
			"-size",
			orig_dimensions,
			imagef,
			"-background",
			"none",
			"-thumbnail",
			fmt.Sprintf("%dx%d^", maxwidth, maxheight),
			"-gravity",
			"center",
			"-extent",
			fmt.Sprintf("%dx%d", maxwidth, maxheight),
			i.Thumbpath,
		}
	} else {
		args = []string{
			"-background",
			"white",
			"-flatten",
			"-size",
			orig_dimensions,
			"-resize",
			fmt.Sprintf("%dx%d>", maxwidth, maxheight),
			"-quality",
			"90",
			imagef,
			i.Thumbpath,
		}
	}

	_, err = exec.Command("convert", args...).Output()
	if err != nil {
		return errors.New("Problem making thumbnail")
	}

	thumb, err := os.Open(i.Thumbpath)
	if err != nil {
		return errors.New("Problem creating thumbnail file")
	}
	defer thumb.Close()

	img, _, err := image.DecodeConfig(thumb)
	if err != nil {
		return errors.New("Problem decoding thumbnail")
	}

	i.ThumbWidth = img.Width
	i.ThumbHeight = img.Height

	return

}

func (i *ImageType) copyToS3() (err error) {

	s3 := amazon.New()

	err = s3.Save(i.Filepath, fmt.Sprintf("src/%s", i.Filename), i.mime, false)
	if err != nil {
		return
	}

	err = s3.Save(i.Thumbpath, fmt.Sprintf("thumb/%s", i.Thumbnail), "image/jpeg", false)
	if err != nil {
		return
	}

	return
}
<|endoftext|>


<commit_before>package video

import "errors"
import "image"

var ErrInvalidReservedBits = errors.New("invalid reserved bits")

type Slice struct {
	slice_start_code                  StartCode // 32 bslbf
	slice_vertical_position_extension uint32    // 3 uimsbf
	priority_breakpoint               uint32    // 7 uimsbf
	quantiser_scale_code              uint32    // 5 uimsbf
	intra_slice_flag                  bool      // 1 bslbf
	intra_slice                       bool      // 1 uimsbf
	extra_information                 []uint8
}

func (br *VideoSequence) slice(frame *image.YCbCr) error {

	code, err := br.Read32(32)
	if err != nil {
		return err
	}

	if StartCode(code).isSlice() == false {
		return ErrUnexpectedStartCode
	}

	br.resetDCPredictors()

	// Reset motion vector predictors: Start if each slice (7.6.4.3)
	br.pMV.reset()

	s := Slice{}
	s.slice_start_code = StartCode(code)

	mb_row := int((code & 0xFF) - 1)

	frameSlice := &image.YCbCr{
		Y:       frame.Y[16*mb_row*frame.YStride:],
		Cb:      frame.Cb[8*mb_row*frame.CStride:],
		Cr:      frame.Cr[8*mb_row*frame.CStride:],
		YStride: frame.YStride,
		CStride: frame.CStride,
		Rect:    image.Rect(frame.Rect.Min.X, mb_row*16, frame.Rect.Max.X, mb_row*16+16),
	}

	if br.SequenceHeader.vertical_size_value > 2800 {
		s.slice_vertical_position_extension, err = br.Read32(3)
		if err != nil {
			return err
		}
	}

	// TODO
	if false /* (<sequence_scalable_extension() is present in the bitstream>) */ {
		if false /* (scalable_mode == “data partitioning” ) */ {
			s.priority_breakpoint, err = br.Read32(7)
			if err != nil {
				return err
			}
		}
	}

	if qsc, err := br.Read32(5); err != nil {
		return err
	} else {
		s.quantiser_scale_code = qsc
		br.currentQSC = qsc
	}

	if nextbits, err := br.Peek32(1); err != nil {
		return err
	} else if nextbits == 1 {
		s.intra_slice_flag, err = br.ReadBit()
		if err != nil {
			return err
		}

		s.intra_slice, err = br.ReadBit()
		if err != nil {
			return err
		}

		if reserved_bits, err := br.Read32(7); err != nil {
			return err
		} else if reserved_bits != 0 {
			return ErrInvalidReservedBits
		}

		for {
			if nextbits, err := br.Peek32(1); err != nil {
				return err
			} else if nextbits != 1 {
				break
			}

			if err = br.Trash(1); err != nil {
				return err
			}

			if data, err := br.Read32(8); err != nil {
				return err
			} else {
				s.extra_information = append(s.extra_information, uint8(data))
			}
		}
	}

	if err := br.Trash(1); err != nil {
		return err
	}

	var mb_address int = -1
	for {
		mb_address, err = br.macroblock(mb_address, mb_row, frameSlice)
		if err != nil {
			return err
		}

		if nextbits, err := br.Peek32(23); err != nil {
			return err
		} else if nextbits == 0 {
			break
		}
	}

	return next_start_code(br)
}
<commit_msg>updating todo<commit_after>package video

import "errors"
import "image"

var ErrInvalidReservedBits = errors.New("invalid reserved bits")

type Slice struct {
	slice_start_code                  StartCode // 32 bslbf
	slice_vertical_position_extension uint32    // 3 uimsbf
	priority_breakpoint               uint32    // 7 uimsbf
	quantiser_scale_code              uint32    // 5 uimsbf
	intra_slice_flag                  bool      // 1 bslbf
	intra_slice                       bool      // 1 uimsbf
	extra_information                 []uint8
}

func (br *VideoSequence) slice(frame *image.YCbCr) error {

	code, err := br.Read32(32)
	if err != nil {
		return err
	}

	if StartCode(code).isSlice() == false {
		return ErrUnexpectedStartCode
	}

	br.resetDCPredictors()

	// Reset motion vector predictors: Start if each slice (7.6.4.3)
	br.pMV.reset()

	s := Slice{}
	s.slice_start_code = StartCode(code)

	mb_row := int((code & 0xFF) - 1)

	frameSlice := &image.YCbCr{
		Y:       frame.Y[16*mb_row*frame.YStride:],
		Cb:      frame.Cb[8*mb_row*frame.CStride:],
		Cr:      frame.Cr[8*mb_row*frame.CStride:],
		YStride: frame.YStride,
		CStride: frame.CStride,
		Rect:    image.Rect(frame.Rect.Min.X, mb_row*16, frame.Rect.Max.X, mb_row*16+16),
	}

	if br.SequenceHeader.vertical_size_value > 2800 {
		s.slice_vertical_position_extension, err = br.Read32(3)
		if err != nil {
			return err
		}
	}

	// TODO(jh): sequence_scalable support
	if false /* (<sequence_scalable_extension() is present in the bitstream>) */ {
		if false /* (scalable_mode == “data partitioning” ) */ {
			s.priority_breakpoint, err = br.Read32(7)
			if err != nil {
				return err
			}
		}
	}

	if qsc, err := br.Read32(5); err != nil {
		return err
	} else {
		s.quantiser_scale_code = qsc
		br.currentQSC = qsc
	}

	if nextbits, err := br.Peek32(1); err != nil {
		return err
	} else if nextbits == 1 {
		s.intra_slice_flag, err = br.ReadBit()
		if err != nil {
			return err
		}

		s.intra_slice, err = br.ReadBit()
		if err != nil {
			return err
		}

		if reserved_bits, err := br.Read32(7); err != nil {
			return err
		} else if reserved_bits != 0 {
			return ErrInvalidReservedBits
		}

		for {
			if nextbits, err := br.Peek32(1); err != nil {
				return err
			} else if nextbits != 1 {
				break
			}

			if err = br.Trash(1); err != nil {
				return err
			}

			if data, err := br.Read32(8); err != nil {
				return err
			} else {
				s.extra_information = append(s.extra_information, uint8(data))
			}
		}
	}

	if err := br.Trash(1); err != nil {
		return err
	}

	var mb_address int = -1
	for {
		mb_address, err = br.macroblock(mb_address, mb_row, frameSlice)
		if err != nil {
			return err
		}

		if nextbits, err := br.Peek32(23); err != nil {
			return err
		} else if nextbits == 0 {
			break
		}
	}

	return next_start_code(br)
}
<|endoftext|>


<commit_before>package main

import (
	"code.google.com/p/gopacket/layers"
	"crypto/sha256"
	"flag"
	"fmt"
	"github.com/davecheney/profile"
	"github.com/zettio/weave"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"
)

func ensureInterface(ifaceName string, wait int) (iface *net.Interface, err error) {
	iface, err = findInterface(ifaceName)
	if err == nil || wait == 0 {
		return
	}
	log.Println("Waiting for interface", ifaceName, "to come up")
	for ; err != nil && wait > 0; wait -= 1 {
		time.Sleep(1 * time.Second)
		iface, err = findInterface(ifaceName)
	}
	if err == nil {
		log.Println("Interface", ifaceName, "is up")
	}
	return
}

func findInterface(ifaceName string) (iface *net.Interface, err error) {
	iface, err = net.InterfaceByName(ifaceName)
	if err != nil {
		return iface, fmt.Errorf("Unable to find interface %s", ifaceName)
	}
	if 0 == (net.FlagUp & iface.Flags) {
		return iface, fmt.Errorf("Interface %s is not up", ifaceName)
	}
	return
}

func main() {

	log.SetPrefix(weave.Protocol + " ")
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)
	log.Println(os.Args)

	procs := runtime.NumCPU()
	// packet sniffing can block an OS thread, so we need one thread
	// for that plus at least one more.
	if procs < 2 {
		procs = 2
	}
	runtime.GOMAXPROCS(procs)

	var (
		ifaceName  string
		routerName string
		password   string
		wait       int
		debug      bool
		prof       string
		peers      []string
		connLimit  int
		bufSz      int
	)

	flag.StringVar(&ifaceName, "iface", "", "name of interface to read from")
	flag.StringVar(&routerName, "name", "", "name of router (defaults to MAC)")
	flag.StringVar(&password, "password", "", "network password")
	flag.IntVar(&wait, "wait", 0, "number of seconds to wait for interface to be created and come up (defaults to 0, i.e. don't wait)")
	flag.BoolVar(&debug, "debug", false, "enable debug logging")
	flag.StringVar(&prof, "profile", "", "enable profiling and write profiles to given path")
	flag.IntVar(&connLimit, "connlimit", 10, "connection limit (defaults to 10, set to 0 for unlimited)")
	flag.IntVar(&bufSz, "bufsz", 8, "capture buffer size in MB (defaults to 8MB)")
	flag.Parse()
	peers = flag.Args()

	if ifaceName == "" {
		fmt.Println("Missing required parameter 'iface'")
		os.Exit(1)
	}
	iface, err := ensureInterface(ifaceName, wait)
	if err != nil {
		log.Fatal(err)
	}

	if connLimit < 0 {
		connLimit = 0
	}

	if routerName == "" {
		routerName = iface.HardwareAddr.String()
	}

	ourName, err := weave.PeerNameFromUserInput(routerName)
	if err != nil {
		log.Fatal(err)
	}

	var logFrame func(string, []byte, *layers.Ethernet)
	if debug {
		logFrame = func(prefix string, frame []byte, eth *layers.Ethernet) {
			h := fmt.Sprintf("%x", sha256.Sum256(frame))
			if eth == nil {
				log.Println(prefix, len(frame), "bytes (", h, ")")
			} else {
				log.Println(prefix, len(frame), "bytes (", h, "):", eth.SrcMAC, "->", eth.DstMAC)
			}
		}
	} else {
		logFrame = func(prefix string, frame []byte, eth *layers.Ethernet) {}
	}

	if prof != "" {
		p := *profile.CPUProfile
		p.ProfilePath = prof
		defer profile.Start(&p).Stop()
	}

	router := weave.NewRouter(iface, ourName, []byte(password), connLimit, bufSz*1024*1024, logFrame)
	router.Start()
	for _, peer := range peers {
		router.ConnectionMaker.InitiateConnection(peer)
	}
	go handleHttp(router)
	handleSignals(router)
}

func handleHttp(router *weave.Router) {
	http.HandleFunc("/status", func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, router.Status())
	})
	address := fmt.Sprintf(":%d", weave.StatusPort)
	err := http.ListenAndServe(address, nil)
	if err != nil {
		log.Fatal("Unable to create http listener: ", err)
	}
}

func handleSignals(router *weave.Router) {
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGQUIT, syscall.SIGUSR1)
	buf := make([]byte, 1<<20)
	for {
		sig := <-sigs
		switch sig {
		case syscall.SIGQUIT:
			runtime.Stack(buf, true)
			log.Printf("=== received SIGQUIT ===\n*** goroutine dump...\n%s\n*** end\n", buf)
		case syscall.SIGUSR1:
			log.Printf("=== received SIGUSR1 ===\n*** status...\n%s\n*** end\n", router.Status())
		}
	}
}
<commit_msg>Resolve peer names to IP address at startup<commit_after>package main

import (
	"code.google.com/p/gopacket/layers"
	"crypto/sha256"
	"flag"
	"fmt"
	"github.com/davecheney/profile"
	"github.com/zettio/weave"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"
)

func ensureInterface(ifaceName string, wait int) (iface *net.Interface, err error) {
	iface, err = findInterface(ifaceName)
	if err == nil || wait == 0 {
		return
	}
	log.Println("Waiting for interface", ifaceName, "to come up")
	for ; err != nil && wait > 0; wait -= 1 {
		time.Sleep(1 * time.Second)
		iface, err = findInterface(ifaceName)
	}
	if err == nil {
		log.Println("Interface", ifaceName, "is up")
	}
	return
}

func findInterface(ifaceName string) (iface *net.Interface, err error) {
	iface, err = net.InterfaceByName(ifaceName)
	if err != nil {
		return iface, fmt.Errorf("Unable to find interface %s", ifaceName)
	}
	if 0 == (net.FlagUp & iface.Flags) {
		return iface, fmt.Errorf("Interface %s is not up", ifaceName)
	}
	return
}

func main() {

	log.SetPrefix(weave.Protocol + " ")
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)
	log.Println(os.Args)

	procs := runtime.NumCPU()
	// packet sniffing can block an OS thread, so we need one thread
	// for that plus at least one more.
	if procs < 2 {
		procs = 2
	}
	runtime.GOMAXPROCS(procs)

	var (
		ifaceName  string
		routerName string
		password   string
		wait       int
		debug      bool
		prof       string
		peers      []string
		connLimit  int
		bufSz      int
	)

	flag.StringVar(&ifaceName, "iface", "", "name of interface to read from")
	flag.StringVar(&routerName, "name", "", "name of router (defaults to MAC)")
	flag.StringVar(&password, "password", "", "network password")
	flag.IntVar(&wait, "wait", 0, "number of seconds to wait for interface to be created and come up (defaults to 0, i.e. don't wait)")
	flag.BoolVar(&debug, "debug", false, "enable debug logging")
	flag.StringVar(&prof, "profile", "", "enable profiling and write profiles to given path")
	flag.IntVar(&connLimit, "connlimit", 10, "connection limit (defaults to 10, set to 0 for unlimited)")
	flag.IntVar(&bufSz, "bufsz", 8, "capture buffer size in MB (defaults to 8MB)")
	flag.Parse()
	peers = flag.Args()

	if ifaceName == "" {
		fmt.Println("Missing required parameter 'iface'")
		os.Exit(1)
	}
	iface, err := ensureInterface(ifaceName, wait)
	if err != nil {
		log.Fatal(err)
	}

	if connLimit < 0 {
		connLimit = 0
	}

	if routerName == "" {
		routerName = iface.HardwareAddr.String()
	}

	ourName, err := weave.PeerNameFromUserInput(routerName)
	if err != nil {
		log.Fatal(err)
	}

	var logFrame func(string, []byte, *layers.Ethernet)
	if debug {
		logFrame = func(prefix string, frame []byte, eth *layers.Ethernet) {
			h := fmt.Sprintf("%x", sha256.Sum256(frame))
			if eth == nil {
				log.Println(prefix, len(frame), "bytes (", h, ")")
			} else {
				log.Println(prefix, len(frame), "bytes (", h, "):", eth.SrcMAC, "->", eth.DstMAC)
			}
		}
	} else {
		logFrame = func(prefix string, frame []byte, eth *layers.Ethernet) {}
	}

	if prof != "" {
		p := *profile.CPUProfile
		p.ProfilePath = prof
		defer profile.Start(&p).Stop()
	}

	router := weave.NewRouter(iface, ourName, []byte(password), connLimit, bufSz*1024*1024, logFrame)
	router.Start()
	for _, peer := range peers {
		normPeer := weave.NormalisePeerAddr(peer)
		addr, err := net.ResolveTCPAddr("tcp4", normPeer)
		if err == nil {
			router.ConnectionMaker.InitiateConnection(addr.String())
		} else {
			log.Fatal(err)
		}
	}
	go handleHttp(router)
	handleSignals(router)
}

func handleHttp(router *weave.Router) {
	http.HandleFunc("/status", func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, router.Status())
	})
	address := fmt.Sprintf(":%d", weave.StatusPort)
	err := http.ListenAndServe(address, nil)
	if err != nil {
		log.Fatal("Unable to create http listener: ", err)
	}
}

func handleSignals(router *weave.Router) {
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGQUIT, syscall.SIGUSR1)
	buf := make([]byte, 1<<20)
	for {
		sig := <-sigs
		switch sig {
		case syscall.SIGQUIT:
			runtime.Stack(buf, true)
			log.Printf("=== received SIGQUIT ===\n*** goroutine dump...\n%s\n*** end\n", buf)
		case syscall.SIGUSR1:
			log.Printf("=== received SIGUSR1 ===\n*** status...\n%s\n*** end\n", router.Status())
		}
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2014 - 2015 Workiva, LLC

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

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

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

package rest

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
)

// Supported HTTP Methods
const (
	httpDelete = "DELETE"
	httpGet    = "GET"
	httpPost   = "POST"
	httpPut    = "PUT"
)

// Client is the type that encapsulates and uses the Authorizer to sign any REST
// requests that are performed.
type Client struct {
	*http.Client
}

// Response is unmarshaled struct returned from an HTTP request.
type Response struct {
	Status   int            // HTTP status code.
	Reason   string         // Reason message for the status code.
	Messages []string       // Any server messages attached to the Response.
	Next     string         // A cursor to the next result set.
	Result   interface{}    // The decoded result of the REST request.
	Raw      *http.Response // The raw HTTP response.
}

// Wraps response decoding error in a helpful way
type ResponseDecodeError struct {
	StatusCode  int    // Response status code
	Status      string // Response status message
	Response    []byte // Payload of the response that could not be decoded
	DecodeError error  // Error that occurred while decoding the response
}

func (rde *ResponseDecodeError) Error() string {
	return fmt.Sprintf("(Error, Status Code, Status Message, Payload) = (%s, %d, %s, %s)",
		rde.DecodeError.Error(), rde.StatusCode, rde.Status, string(rde.Response))
}

// Get will perform an HTTP GET on the specified URL and return the response.
func (c *Client) Get(url string, header http.Header) (*Response, error) {
	return do(c.Client, httpGet, url, nil, header)
}

// Post will perform an HTTP POST on the specified URL and return the response.
func (c *Client) Post(url string, body interface{}, header http.Header) (*Response, error) {
	return do(c.Client, httpPost, url, body, header)
}

// Put will perform an HTTP PUT on the specified URL and return the response.
func (c *Client) Put(url string, body interface{}, header http.Header) (*Response, error) {
	return do(c.Client, httpPut, url, body, header)
}

// Delete will perform an HTTP DELETE on the specified URL and return the response.
func (c *Client) Delete(url string, header http.Header) (*Response, error) {
	return do(c.Client, httpDelete, url, nil, header)
}

var do = func(c *http.Client, method, url string, body interface{}, header http.Header) (*Response, error) {
	if header == nil {
		header = http.Header{}
	}

	var reqBody io.Reader
	switch method {
	case httpPost, httpPut:
		body, err := json.Marshal(body)
		if err != nil {
			return nil, err
		}
		reqBody = bytes.NewReader(body)
		header.Set("Content-Type", "application/json")
	}

	req, err := http.NewRequest(method, url, reqBody)
	if err != nil {
		return nil, err
	}

	req.Header = header

	resp, err := c.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	// Don't try to decode the response on 404.
	if resp.StatusCode == http.StatusNotFound {
		return &Response{
			Status:   resp.StatusCode,
			Reason:   http.StatusText(resp.StatusCode),
			Messages: []string{},
			Next:     "",
			Raw:      resp,
		}, nil
	}

	rawResp, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return decodeResponse(rawResp, resp)
}

func decodeResponse(response []byte, r *http.Response) (*Response, error) {
	var payload map[string]interface{}
	if err := json.Unmarshal(response, &payload); err != nil {
		return nil, &ResponseDecodeError{
			StatusCode:  r.StatusCode,
			Status:      r.Status,
			DecodeError: err,
			Response:    response,
		}
	}

	messages := []string{}
	for _, message := range payload["messages"].([]interface{}) {
		messages = append(messages, message.(string))
	}

	next := ""
	if n, ok := payload["next"]; ok {
		next = n.(string)
	}

	result, ok := payload["result"]
	if !ok {
		result = payload["results"]
	}

	resp := &Response{
		Status:   int(payload["status"].(float64)),
		Reason:   payload["reason"].(string),
		Messages: messages,
		Next:     next,
		Result:   result,
		Raw:      r,
	}

	return resp, nil
}
<commit_msg>Add initial rest client middleware<commit_after>/*
Copyright 2014 - 2015 Workiva, LLC

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

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

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

package rest

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
)

// Supported HTTP Methods
const (
	httpDelete = "DELETE"
	httpGet    = "GET"
	httpPost   = "POST"
	httpPut    = "PUT"
)

type InovcationHandler func(c *http.Client, method, url string, body interface{}, header http.Header) (*Response, error)

type ClientMiddleware func(InovcationHandler) InovcationHandler

type HttpClient interface {
	Do(req *http.Request) (resp *http.Response, err error)
	Get(url string) (resp *http.Response, err error)
	Post(url string, bodyType string, body io.Reader) (resp *http.Response, err error)
	PostForm(url string, data url.Values) (resp *http.Response, err error)
	Head(url string) (resp *http.Response, err error)
}

type RestClient interface {
	// Get will perform an HTTP GET on the specified URL and return the response.
	Get(url string, header http.Header) (*Response, error)

	// Post will perform an HTTP POST on the specified URL and return the response.
	Post(url string, body interface{}, header http.Header) (*Response, error)

	// Put will perform an HTTP PUT on the specified URL and return the response.
	Put(url string, body interface{}, header http.Header) (*Response, error)

	// Delete will perform an HTTP DELETE on the specified URL and return the response.
	Delete(url string, header http.Header) (*Response, error)
}

func NewRestClient(c HttpClient, middleware ...ClientMiddleware) RestClient {
	return &client{c, middleware}
}

// Client is the type that encapsulates and uses the Authorizer to sign any REST
// requests that are performed.
type Client struct {
	HttpClient
}

type client struct {
	HttpClient
	middleware []ClientMiddleware
}

// Response is unmarshaled struct returned from an HTTP request.
type Response struct {
	Status   int            // HTTP status code.
	Reason   string         // Reason message for the status code.
	Messages []string       // Any server messages attached to the Response.
	Next     string         // A cursor to the next result set.
	Result   interface{}    // The decoded result of the REST request.
	Raw      *http.Response // The raw HTTP response.
}

// Wraps response decoding error in a helpful way
type ResponseDecodeError struct {
	StatusCode  int    // Response status code
	Status      string // Response status message
	Response    []byte // Payload of the response that could not be decoded
	DecodeError error  // Error that occurred while decoding the response
}

func (rde *ResponseDecodeError) Error() string {
	return fmt.Sprintf("(Error, Status Code, Status Message, Payload) = (%s, %d, %s, %s)",
		rde.DecodeError.Error(), rde.StatusCode, rde.Status, string(rde.Response))
}

func (c *client) applyMiddleware(method InovcationHandler) InovcationHandler {
	var m InovcationHandler
	for _, middleware := range c.middleware {
		m = middleware(method)
	}
	return m
}

func (c *client) process(method, url string, body interface{}, header http.Header) (*Response, error) {
	m := c.applyMiddleware(do)
	return m(c.HttpClient.(*http.Client), method, url, body, header)
}

// Get will perform an HTTP GET on the specified URL and return the response.
func (c *client) Get(url string, header http.Header) (*Response, error) {
	return c.process(httpGet, url, nil, header)
}

// Post will perform an HTTP POST on the specified URL and return the response.
func (c *client) Post(url string, body interface{}, header http.Header) (*Response, error) {
	return c.process(httpPost, url, body, header)
}

// Put will perform an HTTP PUT on the specified URL and return the response.
func (c *client) Put(url string, body interface{}, header http.Header) (*Response, error) {
	return c.process(httpPut, url, body, header)
}

// Delete will perform an HTTP DELETE on the specified URL and return the response.
func (c *client) Delete(url string, header http.Header) (*Response, error) {
	return c.process(httpDelete, url, nil, header)
}

/* Internal client calls. Would like to move up to the top level with a major
version change. Keeping this around for now to maintain backwards compat and allow
consumers to still construct the Client struct directly vs using the new constructor
method that returns the internal implementation that maps to the interface.
*/

// Get will perform an HTTP GET on the specified URL and return the response.
func (c *Client) Get(url string, header http.Header) (*Response, error) {
	return do(c.HttpClient.(*http.Client), httpGet, url, nil, header)
}

// Post will perform an HTTP POST on the specified URL and return the response.
func (c *Client) Post(url string, body interface{}, header http.Header) (*Response, error) {
	return do(c.HttpClient.(*http.Client), httpPost, url, body, header)
}

// Put will perform an HTTP PUT on the specified URL and return the response.
func (c *Client) Put(url string, body interface{}, header http.Header) (*Response, error) {
	return do(c.HttpClient.(*http.Client), httpPut, url, body, header)
}

// Delete will perform an HTTP DELETE on the specified URL and return the response.
func (c *Client) Delete(url string, header http.Header) (*Response, error) {
	return do(c.HttpClient.(*http.Client), httpDelete, url, nil, header)
}

var do = func(c *http.Client, method, url string, body interface{}, header http.Header) (*Response, error) {
	if header == nil {
		header = http.Header{}
	}

	var reqBody io.Reader
	switch method {
	case httpPost, httpPut:
		body, err := json.Marshal(body)
		if err != nil {
			return nil, err
		}
		reqBody = bytes.NewReader(body)
		header.Set("Content-Type", "application/json")
	}

	req, err := http.NewRequest(method, url, reqBody)
	if err != nil {
		return nil, err
	}

	req.Header = header

	resp, err := c.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	// Don't try to decode the response on 404.
	if resp.StatusCode == http.StatusNotFound {
		return &Response{
			Status:   resp.StatusCode,
			Reason:   http.StatusText(resp.StatusCode),
			Messages: []string{},
			Next:     "",
			Raw:      resp,
		}, nil
	}

	rawResp, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return decodeResponse(rawResp, resp)
}

func decodeResponse(response []byte, r *http.Response) (*Response, error) {
	var payload map[string]interface{}
	if err := json.Unmarshal(response, &payload); err != nil {
		return nil, &ResponseDecodeError{
			StatusCode:  r.StatusCode,
			Status:      r.Status,
			DecodeError: err,
			Response:    response,
		}
	}

	messages := []string{}
	for _, message := range payload["messages"].([]interface{}) {
		messages = append(messages, message.(string))
	}

	next := ""
	if n, ok := payload["next"]; ok {
		next = n.(string)
	}

	result, ok := payload["result"]
	if !ok {
		result = payload["results"]
	}

	resp := &Response{
		Status:   int(payload["status"].(float64)),
		Reason:   payload["reason"].(string),
		Messages: messages,
		Next:     next,
		Result:   result,
		Raw:      r,
	}

	return resp, nil
}
<|endoftext|>


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

package cbft

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"runtime/pprof"
	"sort"
	"strings"
	"time"

	"github.com/gorilla/mux"
)

// DiagHandler represents a part of the DiagGetHandler response.
// Different modules can provide their parts of a DiagGetHandler
// response via their own DiagHandler's.
type DiagHandler struct {
	Name        string
	Handler     http.Handler
	HandlerFunc http.HandlerFunc
}

// DiagGetHandler is a REST handler that retrieves diagnostic
// information for a node.
type DiagGetHandler struct {
	versionMain string
	mgr         *Manager
	mr          *MsgRing
}

func NewDiagGetHandler(versionMain string,
	mgr *Manager, mr *MsgRing) *DiagGetHandler {
	return &DiagGetHandler{versionMain: versionMain, mgr: mgr, mr: mr}
}

func (h *DiagGetHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	handlers := []DiagHandler{
		{"/api/cfg", NewCfgGetHandler(h.mgr), nil},
		{"/api/index", NewListIndexHandler(h.mgr), nil},
		{"/api/log", NewLogGetHandler(h.mgr, h.mr), nil},
		{"/api/managerMeta", NewManagerMetaHandler(h.mgr, nil), nil},
		{"/api/pindex", NewListPIndexHandler(h.mgr), nil},
		{"/api/runtime", NewRuntimeGetHandler(h.versionMain, h.mgr), nil},
		{"/api/runtime/args", nil, restGetRuntimeArgs},
		{"/api/runtime/stats", nil, restGetRuntimeStats},
		{"/api/runtime/statsMem", nil, restGetRuntimeStatsMem},
		{"/api/stats", NewStatsHandler(h.mgr), nil},
		{"/debug/pprof/goroutine?debug=2", nil,
			func(w http.ResponseWriter, r *http.Request) {
				var b bytes.Buffer
				debug := 2
				pprof.Lookup("goroutine").WriteTo(&b, debug)
				mustEncode(w, b.String())
			}},
	}

	for _, t := range PIndexImplTypes {
		for _, h := range t.DiagHandlers {
			handlers = append(handlers, h)
		}
	}

	w.Write(jsonOpenBrace)
	for i, handler := range handlers {
		if i > 0 {
			w.Write(jsonComma)
		}
		w.Write([]byte(fmt.Sprintf(`"%s":`, handler.Name)))
		if handler.Handler != nil {
			handler.Handler.ServeHTTP(w, req)
		}
		if handler.HandlerFunc != nil {
			handler.HandlerFunc.ServeHTTP(w, req)
		}
	}

	var first = true
	var visit func(path string, f os.FileInfo, err error) error
	visit = func(path string, f os.FileInfo, err error) error {
		m := map[string]interface{}{
			"Path":    path,
			"Name":    f.Name(),
			"Size":    f.Size(),
			"Mode":    f.Mode(),
			"ModTime": f.ModTime().Format(time.RFC3339Nano),
			"IsDir":   f.IsDir(),
		}
		if strings.HasPrefix(f.Name(), "PINDEX_") || // Matches PINDEX_xxx_META.
			strings.HasSuffix(f.Name(), "_META") || // Matches PINDEX_META.
			strings.HasSuffix(f.Name(), ".json") { // Matches index_meta.json.
			b, err := ioutil.ReadFile(path)
			if err == nil {
				m["Contents"] = string(b)
			}
		}
		buf, err := json.Marshal(m)
		if err == nil {
			if !first {
				w.Write(jsonComma)
			}
			w.Write(buf)
			first = false
		}
		return nil
	}

	w.Write([]byte(`,"dataDir":[`))
	filepath.Walk(h.mgr.dataDir, visit)
	w.Write([]byte(`]`))

	entries, err := AssetDir("static/dist")
	if err == nil {
		for _, name := range entries {
			// Ex: "static/dist/manifest.txt".
			a, err := Asset("static/dist/" + name)
			if err == nil {
				j, err := json.Marshal(strings.TrimSpace(string(a)))
				if err == nil {
					w.Write([]byte(`,"`))
					w.Write([]byte("/static/dist/" + name))
					w.Write([]byte(`":`))
					w.Write(j)
				}
			}
		}
	}

	w.Write(jsonCloseBrace)
}

// ---------------------------------------------------

// StatsHandler is a REST handler that provides stats/metrics for a
// node.
type StatsHandler struct {
	mgr *Manager
}

func NewStatsHandler(mgr *Manager) *StatsHandler {
	return &StatsHandler{mgr: mgr}
}

var statsFeedsPrefix = []byte("\"feeds\":{")
var statsPIndexesPrefix = []byte("\"pindexes\":{")
var statsManagerPrefix = []byte(",\"manager\":")
var statsNamePrefix = []byte("\"")
var statsNameSuffix = []byte("\":")

func (h *StatsHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	indexName := mux.Vars(req)["indexName"]

	feeds, pindexes := h.mgr.CurrentMaps()
	feedNames := make([]string, 0, len(feeds))
	for feedName := range feeds {
		feedNames = append(feedNames, feedName)
	}
	sort.Strings(feedNames)

	pindexNames := make([]string, 0, len(pindexes))
	for pindexName := range pindexes {
		pindexNames = append(pindexNames, pindexName)
	}
	sort.Strings(pindexNames)

	w.Write(jsonOpenBrace)

	first := true
	w.Write(statsFeedsPrefix)
	for _, feedName := range feedNames {
		if indexName == "" || indexName == feeds[feedName].IndexName() {
			if !first {
				w.Write(jsonComma)
			}
			first = false
			w.Write(statsNamePrefix)
			w.Write([]byte(feedName))
			w.Write(statsNameSuffix)
			feeds[feedName].Stats(w)
		}
	}
	w.Write(jsonCloseBraceComma)

	first = true
	w.Write(statsPIndexesPrefix)
	for _, pindexName := range pindexNames {
		if indexName == "" || indexName == pindexes[pindexName].IndexName {
			if !first {
				w.Write(jsonComma)
			}
			first = false
			w.Write(statsNamePrefix)
			w.Write([]byte(pindexName))
			w.Write(statsNameSuffix)
			pindexes[pindexName].Dest.Stats(w)
		}
	}
	w.Write(jsonCloseBrace)

	if indexName == "" {
		w.Write(statsManagerPrefix)
		var mgrStats ManagerStats
		h.mgr.stats.AtomicCopyTo(&mgrStats)
		mgrStatsJSON, err := json.Marshal(&mgrStats)
		if err == nil && len(mgrStatsJSON) > 0 {
			w.Write(mgrStatsJSON)
		} else {
			w.Write(jsonNULL)
		}
	}

	w.Write(jsonCloseBrace)
}

// ---------------------------------------------------

// ManagerKickHandler is a REST handler that processes a request to
// kick a manager.
type ManagerKickHandler struct {
	mgr *Manager
}

func NewManagerKickHandler(mgr *Manager) *ManagerKickHandler {
	return &ManagerKickHandler{mgr: mgr}
}

func (h *ManagerKickHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	h.mgr.Kick(req.FormValue("msg"))
	mustEncode(w, struct {
		Status string `json:"status"`
	}{Status: "ok"})
}

// ---------------------------------------------------

// CfgGetHandler is a REST handler that retrieves the contents of the
// Cfg system.
type CfgGetHandler struct {
	mgr *Manager
}

func NewCfgGetHandler(mgr *Manager) *CfgGetHandler {
	return &CfgGetHandler{mgr: mgr}
}

func (h *CfgGetHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	// TODO: Might need to scrub auth passwords from this output.
	cfg := h.mgr.Cfg()
	indexDefs, indexDefsCAS, indexDefsErr :=
		CfgGetIndexDefs(cfg)
	nodeDefsWanted, nodeDefsWantedCAS, nodeDefsWantedErr :=
		CfgGetNodeDefs(cfg, NODE_DEFS_WANTED)
	nodeDefsKnown, nodeDefsKnownCAS, nodeDefsKnownErr :=
		CfgGetNodeDefs(cfg, NODE_DEFS_KNOWN)
	planPIndexes, planPIndexesCAS, planPIndexesErr :=
		CfgGetPlanPIndexes(cfg)
	mustEncode(w, struct {
		Status            string        `json:"status"`
		IndexDefs         *IndexDefs    `json:"indexDefs"`
		IndexDefsCAS      uint64        `json:"indexDefsCAS"`
		IndexDefsErr      error         `json:"indexDefsErr"`
		NodeDefsWanted    *NodeDefs     `json:"nodeDefsWanted"`
		NodeDefsWantedCAS uint64        `json:"nodeDefsWantedCAS"`
		NodeDefsWantedErr error         `json:"nodeDefsWantedErr"`
		NodeDefsKnown     *NodeDefs     `json:"nodeDefsKnown"`
		NodeDefsKnownCAS  uint64        `json:"nodeDefsKnownCAS"`
		NodeDefsKnownErr  error         `json:"nodeDefsKnownErr"`
		PlanPIndexes      *PlanPIndexes `json:"planPIndexes"`
		PlanPIndexesCAS   uint64        `json:"planPIndexesCAS"`
		PlanPIndexesErr   error         `json:"planPIndexesErr"`
	}{
		Status:            "ok",
		IndexDefs:         indexDefs,
		IndexDefsCAS:      indexDefsCAS,
		IndexDefsErr:      indexDefsErr,
		NodeDefsWanted:    nodeDefsWanted,
		NodeDefsWantedCAS: nodeDefsWantedCAS,
		NodeDefsWantedErr: nodeDefsWantedErr,
		NodeDefsKnown:     nodeDefsKnown,
		NodeDefsKnownCAS:  nodeDefsKnownCAS,
		NodeDefsKnownErr:  nodeDefsKnownErr,
		PlanPIndexes:      planPIndexes,
		PlanPIndexesCAS:   planPIndexesCAS,
		PlanPIndexesErr:   planPIndexesErr,
	})
}

// ---------------------------------------------------

// CfgRefreshHandler is a REST handler that processes a request for
// the manager/node to refresh its cached snapshot of the Cfg system
// contents.
type CfgRefreshHandler struct {
	mgr *Manager
}

func NewCfgRefreshHandler(mgr *Manager) *CfgRefreshHandler {
	return &CfgRefreshHandler{mgr: mgr}
}

func (h *CfgRefreshHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	h.mgr.Cfg().Refresh()
	h.mgr.GetIndexDefs(true)
	h.mgr.GetPlanPIndexes(true)
	mustEncode(w, struct {
		Status string `json:"status"`
	}{Status: "ok"})
}
<commit_msg>REST API diag has more /debug/pprof capture<commit_after>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
//  except in compliance with the License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing, software distributed under the
//  License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
//  either express or implied. See the License for the specific language governing permissions
//  and limitations under the License.

package cbft

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"runtime/pprof"
	"sort"
	"strings"
	"time"

	"github.com/gorilla/mux"
)

// DiagHandler represents a part of the DiagGetHandler response.
// Different modules can provide their parts of a DiagGetHandler
// response via their own DiagHandler's.
type DiagHandler struct {
	Name        string
	Handler     http.Handler
	HandlerFunc http.HandlerFunc
}

// DiagGetHandler is a REST handler that retrieves diagnostic
// information for a node.
type DiagGetHandler struct {
	versionMain string
	mgr         *Manager
	mr          *MsgRing
}

func NewDiagGetHandler(versionMain string,
	mgr *Manager, mr *MsgRing) *DiagGetHandler {
	return &DiagGetHandler{versionMain: versionMain, mgr: mgr, mr: mr}
}

func (h *DiagGetHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	handlers := []DiagHandler{
		{"/api/cfg", NewCfgGetHandler(h.mgr), nil},
		{"/api/index", NewListIndexHandler(h.mgr), nil},
		{"/api/log", NewLogGetHandler(h.mgr, h.mr), nil},
		{"/api/managerMeta", NewManagerMetaHandler(h.mgr, nil), nil},
		{"/api/pindex", NewListPIndexHandler(h.mgr), nil},
		{"/api/runtime", NewRuntimeGetHandler(h.versionMain, h.mgr), nil},
		{"/api/runtime/args", nil, restGetRuntimeArgs},
		{"/api/runtime/stats", nil, restGetRuntimeStats},
		{"/api/runtime/statsMem", nil, restGetRuntimeStatsMem},
		{"/api/stats", NewStatsHandler(h.mgr), nil},
		{"/debug/pprof/block?debug=1", nil,
			func(w http.ResponseWriter, r *http.Request) {
				DiagGetPProf(w, "block", 2)
			}},
		{"/debug/pprof/goroutine?debug=2", nil,
			func(w http.ResponseWriter, r *http.Request) {
				DiagGetPProf(w, "goroutine", 2)
			}},
		{"/debug/pprof/heap?debug=1", nil,
			func(w http.ResponseWriter, r *http.Request) {
				DiagGetPProf(w, "heap", 1)
			}},
		{"/debug/pprof/threadcreate?debug=1", nil,
			func(w http.ResponseWriter, r *http.Request) {
				DiagGetPProf(w, "threadcreate", 1)
			}},
	}

	for _, t := range PIndexImplTypes {
		for _, h := range t.DiagHandlers {
			handlers = append(handlers, h)
		}
	}

	w.Write(jsonOpenBrace)
	for i, handler := range handlers {
		if i > 0 {
			w.Write(jsonComma)
		}
		w.Write([]byte(fmt.Sprintf(`"%s":`, handler.Name)))
		if handler.Handler != nil {
			handler.Handler.ServeHTTP(w, req)
		}
		if handler.HandlerFunc != nil {
			handler.HandlerFunc.ServeHTTP(w, req)
		}
	}

	var first = true
	var visit func(path string, f os.FileInfo, err error) error
	visit = func(path string, f os.FileInfo, err error) error {
		m := map[string]interface{}{
			"Path":    path,
			"Name":    f.Name(),
			"Size":    f.Size(),
			"Mode":    f.Mode(),
			"ModTime": f.ModTime().Format(time.RFC3339Nano),
			"IsDir":   f.IsDir(),
		}
		if strings.HasPrefix(f.Name(), "PINDEX_") || // Matches PINDEX_xxx_META.
			strings.HasSuffix(f.Name(), "_META") || // Matches PINDEX_META.
			strings.HasSuffix(f.Name(), ".json") { // Matches index_meta.json.
			b, err := ioutil.ReadFile(path)
			if err == nil {
				m["Contents"] = string(b)
			}
		}
		buf, err := json.Marshal(m)
		if err == nil {
			if !first {
				w.Write(jsonComma)
			}
			w.Write(buf)
			first = false
		}
		return nil
	}

	w.Write([]byte(`,"dataDir":[`))
	filepath.Walk(h.mgr.dataDir, visit)
	w.Write([]byte(`]`))

	entries, err := AssetDir("static/dist")
	if err == nil {
		for _, name := range entries {
			// Ex: "static/dist/manifest.txt".
			a, err := Asset("static/dist/" + name)
			if err == nil {
				j, err := json.Marshal(strings.TrimSpace(string(a)))
				if err == nil {
					w.Write([]byte(`,"`))
					w.Write([]byte("/static/dist/" + name))
					w.Write([]byte(`":`))
					w.Write(j)
				}
			}
		}
	}

	w.Write(jsonCloseBrace)
}

func DiagGetPProf(w http.ResponseWriter, profile string, debug int) {
	var b bytes.Buffer
	pprof.Lookup(profile).WriteTo(&b, debug)
	mustEncode(w, b.String())
}

// ---------------------------------------------------

// StatsHandler is a REST handler that provides stats/metrics for a
// node.
type StatsHandler struct {
	mgr *Manager
}

func NewStatsHandler(mgr *Manager) *StatsHandler {
	return &StatsHandler{mgr: mgr}
}

var statsFeedsPrefix = []byte("\"feeds\":{")
var statsPIndexesPrefix = []byte("\"pindexes\":{")
var statsManagerPrefix = []byte(",\"manager\":")
var statsNamePrefix = []byte("\"")
var statsNameSuffix = []byte("\":")

func (h *StatsHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	indexName := mux.Vars(req)["indexName"]

	feeds, pindexes := h.mgr.CurrentMaps()
	feedNames := make([]string, 0, len(feeds))
	for feedName := range feeds {
		feedNames = append(feedNames, feedName)
	}
	sort.Strings(feedNames)

	pindexNames := make([]string, 0, len(pindexes))
	for pindexName := range pindexes {
		pindexNames = append(pindexNames, pindexName)
	}
	sort.Strings(pindexNames)

	w.Write(jsonOpenBrace)

	first := true
	w.Write(statsFeedsPrefix)
	for _, feedName := range feedNames {
		if indexName == "" || indexName == feeds[feedName].IndexName() {
			if !first {
				w.Write(jsonComma)
			}
			first = false
			w.Write(statsNamePrefix)
			w.Write([]byte(feedName))
			w.Write(statsNameSuffix)
			feeds[feedName].Stats(w)
		}
	}
	w.Write(jsonCloseBraceComma)

	first = true
	w.Write(statsPIndexesPrefix)
	for _, pindexName := range pindexNames {
		if indexName == "" || indexName == pindexes[pindexName].IndexName {
			if !first {
				w.Write(jsonComma)
			}
			first = false
			w.Write(statsNamePrefix)
			w.Write([]byte(pindexName))
			w.Write(statsNameSuffix)
			pindexes[pindexName].Dest.Stats(w)
		}
	}
	w.Write(jsonCloseBrace)

	if indexName == "" {
		w.Write(statsManagerPrefix)
		var mgrStats ManagerStats
		h.mgr.stats.AtomicCopyTo(&mgrStats)
		mgrStatsJSON, err := json.Marshal(&mgrStats)
		if err == nil && len(mgrStatsJSON) > 0 {
			w.Write(mgrStatsJSON)
		} else {
			w.Write(jsonNULL)
		}
	}

	w.Write(jsonCloseBrace)
}

// ---------------------------------------------------

// ManagerKickHandler is a REST handler that processes a request to
// kick a manager.
type ManagerKickHandler struct {
	mgr *Manager
}

func NewManagerKickHandler(mgr *Manager) *ManagerKickHandler {
	return &ManagerKickHandler{mgr: mgr}
}

func (h *ManagerKickHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	h.mgr.Kick(req.FormValue("msg"))
	mustEncode(w, struct {
		Status string `json:"status"`
	}{Status: "ok"})
}

// ---------------------------------------------------

// CfgGetHandler is a REST handler that retrieves the contents of the
// Cfg system.
type CfgGetHandler struct {
	mgr *Manager
}

func NewCfgGetHandler(mgr *Manager) *CfgGetHandler {
	return &CfgGetHandler{mgr: mgr}
}

func (h *CfgGetHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	// TODO: Might need to scrub auth passwords from this output.
	cfg := h.mgr.Cfg()
	indexDefs, indexDefsCAS, indexDefsErr :=
		CfgGetIndexDefs(cfg)
	nodeDefsWanted, nodeDefsWantedCAS, nodeDefsWantedErr :=
		CfgGetNodeDefs(cfg, NODE_DEFS_WANTED)
	nodeDefsKnown, nodeDefsKnownCAS, nodeDefsKnownErr :=
		CfgGetNodeDefs(cfg, NODE_DEFS_KNOWN)
	planPIndexes, planPIndexesCAS, planPIndexesErr :=
		CfgGetPlanPIndexes(cfg)
	mustEncode(w, struct {
		Status            string        `json:"status"`
		IndexDefs         *IndexDefs    `json:"indexDefs"`
		IndexDefsCAS      uint64        `json:"indexDefsCAS"`
		IndexDefsErr      error         `json:"indexDefsErr"`
		NodeDefsWanted    *NodeDefs     `json:"nodeDefsWanted"`
		NodeDefsWantedCAS uint64        `json:"nodeDefsWantedCAS"`
		NodeDefsWantedErr error         `json:"nodeDefsWantedErr"`
		NodeDefsKnown     *NodeDefs     `json:"nodeDefsKnown"`
		NodeDefsKnownCAS  uint64        `json:"nodeDefsKnownCAS"`
		NodeDefsKnownErr  error         `json:"nodeDefsKnownErr"`
		PlanPIndexes      *PlanPIndexes `json:"planPIndexes"`
		PlanPIndexesCAS   uint64        `json:"planPIndexesCAS"`
		PlanPIndexesErr   error         `json:"planPIndexesErr"`
	}{
		Status:            "ok",
		IndexDefs:         indexDefs,
		IndexDefsCAS:      indexDefsCAS,
		IndexDefsErr:      indexDefsErr,
		NodeDefsWanted:    nodeDefsWanted,
		NodeDefsWantedCAS: nodeDefsWantedCAS,
		NodeDefsWantedErr: nodeDefsWantedErr,
		NodeDefsKnown:     nodeDefsKnown,
		NodeDefsKnownCAS:  nodeDefsKnownCAS,
		NodeDefsKnownErr:  nodeDefsKnownErr,
		PlanPIndexes:      planPIndexes,
		PlanPIndexesCAS:   planPIndexesCAS,
		PlanPIndexesErr:   planPIndexesErr,
	})
}

// ---------------------------------------------------

// CfgRefreshHandler is a REST handler that processes a request for
// the manager/node to refresh its cached snapshot of the Cfg system
// contents.
type CfgRefreshHandler struct {
	mgr *Manager
}

func NewCfgRefreshHandler(mgr *Manager) *CfgRefreshHandler {
	return &CfgRefreshHandler{mgr: mgr}
}

func (h *CfgRefreshHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	h.mgr.Cfg().Refresh()
	h.mgr.GetIndexDefs(true)
	h.mgr.GetPlanPIndexes(true)
	mustEncode(w, struct {
		Status string `json:"status"`
	}{Status: "ok"})
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"net/http/httputil"
	"runtime"
	"time"
)

// Console flags
var (
	listen                = flag.String("l", ":8888", "port to accept requests")
	targetProduction      = flag.String("a", "localhost:8080", "where production traffic goes. http://localhost:8080/production")
	altTarget             = flag.String("b", "localhost:8081", "where testing traffic goes. response are skipped. http://localhost:8081/test")
	debug                 = flag.Bool("debug", false, "more logging, showing ignored output")
	productionTimeout     = flag.Int("a.timeout", 3, "timeout in seconds for production traffic")
	alternateTimeout      = flag.Int("b.timeout", 1, "timeout in seconds for alternate site traffic")
	productionHostRewrite = flag.Bool("a.rewrite", false, "rewrite the host header when proxying production traffic")
	alternateHostRewrite  = flag.Bool("b.rewrite", false, "rewrite the host header when proxying alternate site traffic")
	percent               = flag.Float64("p", 100.0, "float64 percentage of traffic to send to testing")
)

// handler contains the address of the main Target and the one for the Alternative target
type handler struct {
	Target      string
	Alternative string
	Randomizer  rand.Rand
}

// ServeHTTP duplicates the incoming request (req) and does the request to the Target and the Alternate target discading the Alternate response
func (h handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var req1, req2 *http.Request
	if *percent == 100.0 || h.Randomizer.Float64()*100 < *percent {
		req1, req2 = DuplicateRequest(req)
		go func() {
			defer func() {
				if r := recover(); r != nil && *debug {
					fmt.Println("Recovered in f", r)
				}
			}()
			// Open new TCP connection to the server
			clientTcpConn, err := net.DialTimeout("tcp", h.Alternative, time.Duration(time.Duration(*alternateTimeout)*time.Second))
			if err != nil {
				if *debug {
					fmt.Printf("Failed to connect to %s\n", h.Alternative)
				}
				return
			}
			clientHttpConn := httputil.NewClientConn(clientTcpConn, nil) // Start a new HTTP connection on it
			defer clientHttpConn.Close()                                 // Close the connection to the server
			if *alternateHostRewrite {
				req1.Host = h.Alternative
			}
			err = clientHttpConn.Write(req1) // Pass on the request
			if err != nil {
				if *debug {
					fmt.Printf("Failed to send to %s: %v\n", h.Alternative, err)
				}
				return
			}
			_, err = clientHttpConn.Read(req1) // Read back the reply
			if err != nil {
				if *debug {
					fmt.Printf("Failed to receive from %s: %v\n", h.Alternative, err)
				}
				return
			}
		}()
	} else {
		req2 = req
	}
	defer func() {
		if r := recover(); r != nil && *debug {
			fmt.Println("Recovered in f", r)
		}
	}()

	// Open new TCP connection to the server
	clientTcpConn, err := net.DialTimeout("tcp", h.Target, time.Duration(time.Duration(*productionTimeout)*time.Second))
	if err != nil {
		fmt.Printf("Failed to connect to %s\n", h.Target)
		return
	}
	clientHttpConn := httputil.NewClientConn(clientTcpConn, nil) // Start a new HTTP connection on it
	defer clientHttpConn.Close()                                 // Close the connection to the server
	if *productionHostRewrite {
		req2.Host = h.Target
	}
	err = clientHttpConn.Write(req2) // Pass on the request
	if err != nil {
		fmt.Printf("Failed to send to %s: %v\n", h.Target, err)
		return
	}
	resp, err := clientHttpConn.Read(req2) // Read back the reply
	if err != nil {
		fmt.Printf("Failed to receive from %s: %v\n", h.Target, err)
		return
	}
	for k, v := range resp.Header {
		w.Header()[k] = v
	}
	w.WriteHeader(resp.StatusCode)
	body, _ := ioutil.ReadAll(resp.Body)
	w.Write(body)
}

func main() {
	flag.Parse()

	runtime.GOMAXPROCS(runtime.NumCPU())

	local, err := net.Listen("tcp", *listen)
	if err != nil {
		fmt.Printf("Failed to listen to %s\n", *listen)
		return
	}
	h := handler{
		Target:      *targetProduction,
		Alternative: *altTarget,
		Randomizer:  *rand.New(rand.NewSource(time.Now().UnixNano())),
	}
	http.Serve(local, h)
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() error { return nil }

func DuplicateRequest(request *http.Request) (request1 *http.Request, request2 *http.Request) {
	b1 := new(bytes.Buffer)
	b2 := new(bytes.Buffer)
	w := io.MultiWriter(b1, b2)
	io.Copy(w, request.Body)
	defer request.Body.Close()
	request1 = &http.Request{
		Method:        request.Method,
		URL:           request.URL,
		Proto:         request.Proto,
		ProtoMajor:    request.ProtoMajor,
		ProtoMinor:    request.ProtoMinor,
		Header:        request.Header,
		Body:          nopCloser{b1},
		Host:          request.Host,
		ContentLength: request.ContentLength,
	}
	request2 = &http.Request{
		Method:        request.Method,
		URL:           request.URL,
		Proto:         request.Proto,
		ProtoMajor:    request.ProtoMajor,
		ProtoMinor:    request.ProtoMinor,
		Header:        request.Header,
		Body:          nopCloser{b2},
		Host:          request.Host,
		ContentLength: request.ContentLength,
	}
	return
}
<commit_msg>support https<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"net/http/httputil"
	"runtime"
	"time"
)

// Console flags
var (
	listen                = flag.String("l", ":8888", "port to accept requests")
	targetProduction      = flag.String("a", "localhost:8080", "where production traffic goes. http://localhost:8080/production")
	altTarget             = flag.String("b", "localhost:8081", "where testing traffic goes. response are skipped. http://localhost:8081/test")
	debug                 = flag.Bool("debug", false, "more logging, showing ignored output")
	productionTimeout     = flag.Int("a.timeout", 3, "timeout in seconds for production traffic")
	alternateTimeout      = flag.Int("b.timeout", 1, "timeout in seconds for alternate site traffic")
	productionHostRewrite = flag.Bool("a.rewrite", false, "rewrite the host header when proxying production traffic")
	alternateHostRewrite  = flag.Bool("b.rewrite", false, "rewrite the host header when proxying alternate site traffic")
	percent               = flag.Float64("p", 100.0, "float64 percentage of traffic to send to testing")
	tlsPrivateKey         = flag.String("key.file", "", "path to the TLS private key file")
	tlsCertificate        = flag.String("cert.file", "", "path to the TLS certificate file")
)

// handler contains the address of the main Target and the one for the Alternative target
type handler struct {
	Target      string
	Alternative string
	Randomizer  rand.Rand
}

// ServeHTTP duplicates the incoming request (req) and does the request to the Target and the Alternate target discading the Alternate response
func (h handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var productionRequest, alternativeRequest *http.Request
	if *percent == 100.0 || h.Randomizer.Float64()*100 < *percent {
		alternativeRequest, productionRequest = DuplicateRequest(req)
		go func() {
			defer func() {
				if r := recover(); r != nil && *debug {
					fmt.Println("Recovered in f", r)
				}
			}()
			// Open new TCP connection to the server
			clientTcpConn, err := net.DialTimeout("tcp", h.Alternative, time.Duration(time.Duration(*alternateTimeout)*time.Second))
			if err != nil {
				if *debug {
					fmt.Printf("Failed to connect to %s\n", h.Alternative)
				}
				return
			}
			clientHttpConn := httputil.NewClientConn(clientTcpConn, nil) // Start a new HTTP connection on it
			defer clientHttpConn.Close()                                 // Close the connection to the server
			if *alternateHostRewrite {
				alternativeRequest.Host = h.Alternative
			}
			err = clientHttpConn.Write(alternativeRequest) // Pass on the request
			if err != nil {
				if *debug {
					fmt.Printf("Failed to send to %s: %v\n", h.Alternative, err)
				}
				return
			}
			_, err = clientHttpConn.Read(alternativeRequest) // Read back the reply
			if err != nil {
				if *debug {
					fmt.Printf("Failed to receive from %s: %v\n", h.Alternative, err)
				}
				return
			}
		}()
	} else {
		productionRequest = req
	}
	defer func() {
		if r := recover(); r != nil && *debug {
			fmt.Println("Recovered in f", r)
		}
	}()

	// Open new TCP connection to the server
	clientTcpConn, err := net.DialTimeout("tcp", h.Target, time.Duration(time.Duration(*productionTimeout)*time.Second))
	if err != nil {
		fmt.Printf("Failed to connect to %s\n", h.Target)
		return
	}
	clientHttpConn := httputil.NewClientConn(clientTcpConn, nil) // Start a new HTTP connection on it
	defer clientHttpConn.Close()                                 // Close the connection to the server
	if *productionHostRewrite {
		productionRequest.Host = h.Target
	}
	err = clientHttpConn.Write(productionRequest) // Pass on the request
	if err != nil {
		fmt.Printf("Failed to send to %s: %v\n", h.Target, err)
		return
	}
	resp, err := clientHttpConn.Read(productionRequest) // Read back the reply
	if err != nil {
		fmt.Printf("Failed to receive from %s: %v\n", h.Target, err)
		return
	}
	for k, v := range resp.Header {
		w.Header()[k] = v
	}
	w.WriteHeader(resp.StatusCode)
	body, _ := ioutil.ReadAll(resp.Body)
	w.Write(body)
}

func main() {
	flag.Parse()

	runtime.GOMAXPROCS(runtime.NumCPU())

	var err error

	var listener net.Listener

	if len(*tlsPrivateKey) > 0 {
		cer, err := tls.LoadX509KeyPair(*tlsCertificate, *tlsPrivateKey)
		if err != nil {
			fmt.Printf("Failed to load certficate: %s and private key: %s", *tlsCertificate, *tlsPrivateKey)
			return
		}

		config := &tls.Config{Certificates: []tls.Certificate{cer}}
		listener, err = tls.Listen("tcp", *listen, config)
		if err != nil {
			fmt.Printf("Failed to listen to %s: %s\n", *listen, err)
			return
		}
	} else {
		listener, err = net.Listen("tcp", *listen)
		if err != nil {
			fmt.Printf("Failed to listen to %s: %s\n", *listen, err)
			return
		}
	}

	h := handler{
		Target:      *targetProduction,
		Alternative: *altTarget,
		Randomizer:  *rand.New(rand.NewSource(time.Now().UnixNano())),
	}
	http.Serve(listener, h)
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() error { return nil }

func DuplicateRequest(request *http.Request) (request1 *http.Request, request2 *http.Request) {
	b1 := new(bytes.Buffer)
	b2 := new(bytes.Buffer)
	w := io.MultiWriter(b1, b2)
	io.Copy(w, request.Body)
	defer request.Body.Close()
	request1 = &http.Request{
		Method:        request.Method,
		URL:           request.URL,
		Proto:         request.Proto,
		ProtoMajor:    request.ProtoMajor,
		ProtoMinor:    request.ProtoMinor,
		Header:        request.Header,
		Body:          nopCloser{b1},
		Host:          request.Host,
		ContentLength: request.ContentLength,
	}
	request2 = &http.Request{
		Method:        request.Method,
		URL:           request.URL,
		Proto:         request.Proto,
		ProtoMajor:    request.ProtoMajor,
		ProtoMinor:    request.ProtoMinor,
		Header:        request.Header,
		Body:          nopCloser{b2},
		Host:          request.Host,
		ContentLength: request.ContentLength,
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api"
	bt "github.com/ikool-cn/gobeanstalk-connection-pool"
)

var (
	telegramBot *TelegramBot
)

// TelegramBot ...
type TelegramBot struct {
	Name           string
	SelfChatID     int64
	ChannelChatID  int64
	ComicPath      string
	PixivPath      string
	TwitterImgPath string
	DeleteDelay    time.Duration
	Client         *tgbotapi.BotAPI
	Queue          *bt.Pool
	Tube           string
}

// NewTelegramBot ...
func NewTelegramBot(cfg *Config) (t *TelegramBot) {
	bot, err := tgbotapi.NewBotAPI(cfg.Telegram.Token)
	if err != nil {
		logger.Panicf("tg bot init failed: %+v", err)
	}
	delay, err := time.ParseDuration(cfg.Telegram.DeleteDelay)
	if err != nil {
		logger.Panicf("delete delay error: %+v", err)
	}

	t = &TelegramBot{
		Name:           bot.Self.UserName,
		SelfChatID:     cfg.Telegram.SelfChatID,
		ChannelChatID:  cfg.Telegram.ChannelChatID,
		ComicPath:      cfg.Telegram.ComicPath,
		PixivPath:      cfg.Pixiv.ImgPath,
		TwitterImgPath: cfg.Twitter.ImgPath,
		DeleteDelay:    delay,
		Client:         bot,
		Tube:           "tg",
	}
	t.Queue = &bt.Pool{
		Dial: func() (*bt.Conn, error) {
			return bt.Dial(cfg.BeanstalkAddr)
		},
		MaxIdle:     10,
		MaxActive:   100,
		IdleTimeout: 60 * time.Second,
		MaxLifetime: 180 * time.Second,
		Wait:        true,
	}
	return
}

func (t *TelegramBot) putQueue(msg []byte) {
	conn, err := t.Queue.Get()
	if err != nil {
		logger.Errorf("%+v: %s", err, string(msg))
		return
	}
	conn.Use(t.Tube)
	_, err = conn.Put(msg, 1, t.DeleteDelay, time.Minute)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
}

func (t *TelegramBot) send(chat int64, msg string) (tgbotapi.Message, error) {
	logger.Debugf("[%d]%s", chat, msg)
	return t.Client.Send(tgbotapi.NewMessage(chat, msg))
}

func (t *TelegramBot) delMessage() {
	for {
		conn, err := t.Queue.Get()
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		conn.Watch(t.Tube)
		job, err := conn.Reserve()
		if err != nil {
			logger.Warningf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		msg := &tgbotapi.Message{}
		err = json.Unmarshal(job.Body, msg)
		if err != nil {
			logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		delMsg := tgbotapi.DeleteMessageConfig{
			ChatID:    msg.Chat.ID,
			MessageID: msg.MessageID,
		}
		logger.Infof(":[%s]{%s}", getMsgTitle(msg), strconv.Quote(msg.Text))

		_, err = t.Client.DeleteMessage(delMsg)
		if err != nil {
			logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		err = conn.Delete(job.ID)
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
		}
		t.Queue.Release(conn, false)
	}
}

func (t *TelegramBot) tgBot() {
	u := tgbotapi.NewUpdate(0)
	u.Timeout = 60
	for {
		updates, err := t.Client.GetUpdatesChan(u)
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		var message *tgbotapi.Message
		for update := range updates {
			if update.Message != nil {
				message = update.Message
			} else if update.EditedMessage != nil {
				message = update.EditedMessage
			} else if update.CallbackQuery != nil {
				logger.Infof(
					"recv:(%s)[%s]reaction:{%s}",
					update.CallbackQuery.ChatInstance,
					update.CallbackQuery.From.String(),
					update.CallbackQuery.Data,
				)
				data := strings.SplitN(update.CallbackQuery.Data, ":", 2)
				switch data[0] {
				case "comic", "pic", "pixiv":
					go onReaction(t, update.CallbackQuery)
				}
				continue
			} else {
				continue
			}
			if message.Chat.IsGroup() {
				logger.Infof(
					"recv:(%s)[%s]{%s}",
					message.Chat.Title,
					message.From.String(),
					strconv.Quote(message.Text))
			} else {
				logger.Infof(
					"recv:[%s]{%s}",
					message.From.String(),
					strconv.Quote(message.Text),
				)
			}

			if message.IsCommand() {
				switch message.Command() {
				case "start":
					go onStart(t, message)
				case "comic":
					go onComic(t, message)
				case "pic":
					go onPic(t, message)
				case "pixiv":
					go onPixiv(t, message)
				default:
					logger.Infof("ignore unknown cmd: %+v", message.Command())
					continue
				}
			} else {
				if message.Text == "" {
					continue
				}
				checkRepeat(t, message)
			}
		}
		logger.Warning("tg bot restarted.")
		time.Sleep(3 * time.Second)
	}
}

func checkRepeat(t *TelegramBot, message *tgbotapi.Message) {
	key := "tg_last_" + strconv.FormatInt(message.Chat.ID, 10)
	flattendMsg := strings.TrimSpace(message.Text)
	defer redisClient.LTrim(key, 0, 10)
	defer redisClient.LPush(key, flattendMsg)

	lastMsgs, err := redisClient.LRange(key, 0, 6).Result()
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	i := 0
	for _, s := range lastMsgs {
		if s == flattendMsg {
			i++
		}
	}
	if i > 1 {
		redisClient.Del(key)
		logger.Infof("repeat: %s", strconv.Quote(message.Text))
		msg := tgbotapi.NewMessage(message.Chat.ID, message.Text)
		t.Client.Send(msg)
	}
}

func onStart(t *TelegramBot, message *tgbotapi.Message) {
	msg := tgbotapi.NewMessage(message.Chat.ID, "呀呀呀")
	msg.ReplyToMessageID = message.MessageID
	t.Client.Send(msg)
}

func onComic(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(filepath.Join(t.ComicPath, "*.epub"))
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	rand.Seed(time.Now().UnixNano())
	file := files[rand.Intn(len(files))]
	number := strings.Split(strings.Split(file, "@")[1], ".")[0]
	msg := tgbotapi.NewMessage(message.Chat.ID, "🔞 https://nhentai.net/g/"+number)

	msg.ReplyMarkup = buildInlineKeyboardMarkup("comic", number)

	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))
	msgSent, err := t.Client.Send(msg)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	data, err := json.Marshal(msgSent)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	t.putQueue(data)
}

func onPic(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(filepath.Join(t.TwitterImgPath, "*"))
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	if files == nil {
		logger.Error("find no pics")
	}
	rand.Seed(time.Now().UnixNano())
	file := files[rand.Intn(len(files))]

	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))

	msg := tgbotapi.NewDocumentUpload(message.Chat.ID, file)
	msg.ReplyMarkup = buildInlineKeyboardMarkup("pic", filepath.Base(file))

	_, err = t.Client.Send(msg)
	if err != nil {
		logger.Errorf("%+v", err)
	}
}

func onPixiv(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(filepath.Join(t.PixivPath, "*"))
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	if files == nil {
		logger.Error("find no pics")
	}
	rand.Seed(time.Now().UnixNano())
	file := files[rand.Intn(len(files))]
	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))
	msg := tgbotapi.NewDocumentUpload(message.Chat.ID, file)
	msg.ReplyMarkup = buildInlineKeyboardMarkup("pixiv", filepath.Base(file))

	_, err = t.Client.Send(msg)
	if err != nil {
		logger.Errorf("%+v", err)
	}
}

func onReaction(t *TelegramBot, callbackQuery *tgbotapi.CallbackQuery) {
	var callbackText string

	_type, _id, reaction, err := saveReaction(callbackQuery.Data, callbackQuery.From.ID)
	if err == nil {
		diss := redisClient.SCard(buildReactionKey(_type, _id, "diss")).Val()
		if diss <= 1 {
			msg := tgbotapi.NewEditMessageReplyMarkup(
				callbackQuery.Message.Chat.ID,
				callbackQuery.Message.MessageID,
				buildInlineKeyboardMarkup(_type, _id),
			)
			_, err = t.Client.Send(msg)
		} else {
			delMsg := tgbotapi.DeleteMessageConfig{
				ChatID:    callbackQuery.Message.Chat.ID,
				MessageID: callbackQuery.Message.MessageID,
			}
			_, err = t.Client.DeleteMessage(delMsg)
			if err == nil {
				err = probate(_type, _id)
			}
		}
	}

	if err != nil {
		logger.Debugf("%+v", err)
		callbackText = err.Error()
	} else {
		callbackText = reaction + " " + _id + "!"
	}

	callbackMsg := tgbotapi.NewCallback(callbackQuery.ID, callbackText)
	_, err = t.Client.AnswerCallbackQuery(callbackMsg)
	if err != nil {
		logger.Errorf("%+v", err)
	}
}

func getMsgTitle(m *tgbotapi.Message) string {
	if m.Chat.IsGroup() {
		return m.Chat.Title
	}
	return m.From.String()
}

func buildReactionData(_type, _id, reaction string) string {
	return _type + ":" + _id + ":" + reaction
}
func buildReactionKey(_type, _id, reaction string) string {
	return "reaction_" + buildReactionData(_type, _id, reaction)
}

func buildInlineKeyboardMarkup(_type, _id string) tgbotapi.InlineKeyboardMarkup {

	likeCount, _ := redisClient.SCard(buildReactionKey(_type, _id, "like")).Result()
	dissCount, _ := redisClient.SCard(buildReactionKey(_type, _id, "diss")).Result()

	likeText := "❤"
	if likeCount > 0 {
		likeText = likeText + " " + strconv.FormatInt(likeCount, 10)
	}
	dissText := "💔"
	if dissCount > 0 {
		dissText = dissText + " " + strconv.FormatInt(dissCount, 10)
	}

	row := tgbotapi.NewInlineKeyboardRow(
		tgbotapi.NewInlineKeyboardButtonData(likeText, buildReactionData(_type, _id, "like")),
		tgbotapi.NewInlineKeyboardButtonData(dissText, buildReactionData(_type, _id, "diss")),
	)
	return tgbotapi.NewInlineKeyboardMarkup(row)
}

func saveReaction(key string, user int) (_type, _id, reaction string, err error) {
	token := strings.Split(key, ":")
	if len(token) != 3 {
		err = fmt.Errorf("react data error: %s", key)
		return
	}
	_type = token[0]
	_id = token[1]
	reaction = token[2]

	pipe := redisClient.Pipeline()
	switch reaction {
	case "like":
		likeCount := pipe.SAdd(buildReactionKey(_type, _id, "like"), strconv.Itoa(user))
		dissCount := pipe.SRem(buildReactionKey(_type, _id, "diss"), strconv.Itoa(user))
		_, err = pipe.Exec()
		if err == nil {
			if likeCount.Val()+dissCount.Val() == 0 {
				err = fmt.Errorf("not modified")
			}
		}
	case "diss":
		dissCount := pipe.SAdd(buildReactionKey(_type, _id, "diss"), strconv.Itoa(user))
		likeCount := pipe.SRem(buildReactionKey(_type, _id, "like"), strconv.Itoa(user))
		_, err = pipe.Exec()
		if err == nil {
			if likeCount.Val()+dissCount.Val() == 0 {
				err = fmt.Errorf("not modified")
			}
		}
	default:
		err = fmt.Errorf("react type error: %s", key)
	}
	return
}
<commit_msg>Return if nothing found<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api"
	bt "github.com/ikool-cn/gobeanstalk-connection-pool"
)

var (
	telegramBot *TelegramBot
)

// TelegramBot ...
type TelegramBot struct {
	Name           string
	SelfChatID     int64
	ChannelChatID  int64
	ComicPath      string
	PixivPath      string
	TwitterImgPath string
	DeleteDelay    time.Duration
	Client         *tgbotapi.BotAPI
	Queue          *bt.Pool
	Tube           string
}

// NewTelegramBot ...
func NewTelegramBot(cfg *Config) (t *TelegramBot) {
	bot, err := tgbotapi.NewBotAPI(cfg.Telegram.Token)
	if err != nil {
		logger.Panicf("tg bot init failed: %+v", err)
	}
	delay, err := time.ParseDuration(cfg.Telegram.DeleteDelay)
	if err != nil {
		logger.Panicf("delete delay error: %+v", err)
	}

	t = &TelegramBot{
		Name:           bot.Self.UserName,
		SelfChatID:     cfg.Telegram.SelfChatID,
		ChannelChatID:  cfg.Telegram.ChannelChatID,
		ComicPath:      cfg.Telegram.ComicPath,
		PixivPath:      cfg.Pixiv.ImgPath,
		TwitterImgPath: cfg.Twitter.ImgPath,
		DeleteDelay:    delay,
		Client:         bot,
		Tube:           "tg",
	}
	t.Queue = &bt.Pool{
		Dial: func() (*bt.Conn, error) {
			return bt.Dial(cfg.BeanstalkAddr)
		},
		MaxIdle:     10,
		MaxActive:   100,
		IdleTimeout: 60 * time.Second,
		MaxLifetime: 180 * time.Second,
		Wait:        true,
	}
	return
}

func (t *TelegramBot) putQueue(msg []byte) {
	conn, err := t.Queue.Get()
	if err != nil {
		logger.Errorf("%+v: %s", err, string(msg))
		return
	}
	conn.Use(t.Tube)
	_, err = conn.Put(msg, 1, t.DeleteDelay, time.Minute)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
}

func (t *TelegramBot) send(chat int64, msg string) (tgbotapi.Message, error) {
	logger.Debugf("[%d]%s", chat, msg)
	return t.Client.Send(tgbotapi.NewMessage(chat, msg))
}

func (t *TelegramBot) delMessage() {
	for {
		conn, err := t.Queue.Get()
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		conn.Watch(t.Tube)
		job, err := conn.Reserve()
		if err != nil {
			logger.Warningf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		msg := &tgbotapi.Message{}
		err = json.Unmarshal(job.Body, msg)
		if err != nil {
			logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		delMsg := tgbotapi.DeleteMessageConfig{
			ChatID:    msg.Chat.ID,
			MessageID: msg.MessageID,
		}
		logger.Infof(":[%s]{%s}", getMsgTitle(msg), strconv.Quote(msg.Text))

		_, err = t.Client.DeleteMessage(delMsg)
		if err != nil {
			logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		err = conn.Delete(job.ID)
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
		}
		t.Queue.Release(conn, false)
	}
}

func (t *TelegramBot) tgBot() {
	u := tgbotapi.NewUpdate(0)
	u.Timeout = 60
	for {
		updates, err := t.Client.GetUpdatesChan(u)
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		var message *tgbotapi.Message
		for update := range updates {
			if update.Message != nil {
				message = update.Message
			} else if update.EditedMessage != nil {
				message = update.EditedMessage
			} else if update.CallbackQuery != nil {
				logger.Infof(
					"recv:(%s)[%s]reaction:{%s}",
					update.CallbackQuery.ChatInstance,
					update.CallbackQuery.From.String(),
					update.CallbackQuery.Data,
				)
				data := strings.SplitN(update.CallbackQuery.Data, ":", 2)
				switch data[0] {
				case "comic", "pic", "pixiv":
					go onReaction(t, update.CallbackQuery)
				}
				continue
			} else {
				continue
			}
			if message.Chat.IsGroup() {
				logger.Infof(
					"recv:(%s)[%s]{%s}",
					message.Chat.Title,
					message.From.String(),
					strconv.Quote(message.Text))
			} else {
				logger.Infof(
					"recv:[%s]{%s}",
					message.From.String(),
					strconv.Quote(message.Text),
				)
			}

			if message.IsCommand() {
				switch message.Command() {
				case "start":
					go onStart(t, message)
				case "comic":
					go onComic(t, message)
				case "pic":
					go onPic(t, message)
				case "pixiv":
					go onPixiv(t, message)
				default:
					logger.Infof("ignore unknown cmd: %+v", message.Command())
					continue
				}
			} else {
				if message.Text == "" {
					continue
				}
				checkRepeat(t, message)
			}
		}
		logger.Warning("tg bot restarted.")
		time.Sleep(3 * time.Second)
	}
}

func checkRepeat(t *TelegramBot, message *tgbotapi.Message) {
	key := "tg_last_" + strconv.FormatInt(message.Chat.ID, 10)
	flattendMsg := strings.TrimSpace(message.Text)
	defer redisClient.LTrim(key, 0, 10)
	defer redisClient.LPush(key, flattendMsg)

	lastMsgs, err := redisClient.LRange(key, 0, 6).Result()
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	i := 0
	for _, s := range lastMsgs {
		if s == flattendMsg {
			i++
		}
	}
	if i > 1 {
		redisClient.Del(key)
		logger.Infof("repeat: %s", strconv.Quote(message.Text))
		msg := tgbotapi.NewMessage(message.Chat.ID, message.Text)
		t.Client.Send(msg)
	}
}

func onStart(t *TelegramBot, message *tgbotapi.Message) {
	msg := tgbotapi.NewMessage(message.Chat.ID, "呀呀呀")
	msg.ReplyToMessageID = message.MessageID
	t.Client.Send(msg)
}

func onComic(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(filepath.Join(t.ComicPath, "*.epub"))
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	if files == nil {
		logger.Error("find no comic")
		return
	}
	rand.Seed(time.Now().UnixNano())
	file := files[rand.Intn(len(files))]
	number := strings.Split(strings.Split(file, "@")[1], ".")[0]
	msg := tgbotapi.NewMessage(message.Chat.ID, "🔞 https://nhentai.net/g/"+number)

	msg.ReplyMarkup = buildInlineKeyboardMarkup("comic", number)

	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))
	msgSent, err := t.Client.Send(msg)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	data, err := json.Marshal(msgSent)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	t.putQueue(data)
}

func onPic(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(filepath.Join(t.TwitterImgPath, "*"))
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	if files == nil {
		logger.Error("find no pic")
		return
	}
	rand.Seed(time.Now().UnixNano())
	file := files[rand.Intn(len(files))]

	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))

	msg := tgbotapi.NewDocumentUpload(message.Chat.ID, file)
	msg.ReplyMarkup = buildInlineKeyboardMarkup("pic", filepath.Base(file))

	_, err = t.Client.Send(msg)
	if err != nil {
		logger.Errorf("%+v", err)
	}
}

func onPixiv(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(filepath.Join(t.PixivPath, "*"))
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	if files == nil {
		logger.Error("find no pic")
		return
	}
	rand.Seed(time.Now().UnixNano())
	file := files[rand.Intn(len(files))]
	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))
	msg := tgbotapi.NewDocumentUpload(message.Chat.ID, file)
	msg.ReplyMarkup = buildInlineKeyboardMarkup("pixiv", filepath.Base(file))

	_, err = t.Client.Send(msg)
	if err != nil {
		logger.Errorf("%+v", err)
	}
}

func onReaction(t *TelegramBot, callbackQuery *tgbotapi.CallbackQuery) {
	var callbackText string

	_type, _id, reaction, err := saveReaction(callbackQuery.Data, callbackQuery.From.ID)
	if err == nil {
		diss := redisClient.SCard(buildReactionKey(_type, _id, "diss")).Val()
		if diss <= 1 {
			msg := tgbotapi.NewEditMessageReplyMarkup(
				callbackQuery.Message.Chat.ID,
				callbackQuery.Message.MessageID,
				buildInlineKeyboardMarkup(_type, _id),
			)
			_, err = t.Client.Send(msg)
		} else {
			delMsg := tgbotapi.DeleteMessageConfig{
				ChatID:    callbackQuery.Message.Chat.ID,
				MessageID: callbackQuery.Message.MessageID,
			}
			_, err = t.Client.DeleteMessage(delMsg)
			if err == nil {
				err = probate(_type, _id)
			}
		}
	}

	if err != nil {
		logger.Debugf("%+v", err)
		callbackText = err.Error()
	} else {
		callbackText = reaction + " " + _id + "!"
	}

	callbackMsg := tgbotapi.NewCallback(callbackQuery.ID, callbackText)
	_, err = t.Client.AnswerCallbackQuery(callbackMsg)
	if err != nil {
		logger.Errorf("%+v", err)
	}
}

func getMsgTitle(m *tgbotapi.Message) string {
	if m.Chat.IsGroup() {
		return m.Chat.Title
	}
	return m.From.String()
}

func buildReactionData(_type, _id, reaction string) string {
	return _type + ":" + _id + ":" + reaction
}
func buildReactionKey(_type, _id, reaction string) string {
	return "reaction_" + buildReactionData(_type, _id, reaction)
}

func buildInlineKeyboardMarkup(_type, _id string) tgbotapi.InlineKeyboardMarkup {

	likeCount, _ := redisClient.SCard(buildReactionKey(_type, _id, "like")).Result()
	dissCount, _ := redisClient.SCard(buildReactionKey(_type, _id, "diss")).Result()

	likeText := "❤"
	if likeCount > 0 {
		likeText = likeText + " " + strconv.FormatInt(likeCount, 10)
	}
	dissText := "💔"
	if dissCount > 0 {
		dissText = dissText + " " + strconv.FormatInt(dissCount, 10)
	}

	row := tgbotapi.NewInlineKeyboardRow(
		tgbotapi.NewInlineKeyboardButtonData(likeText, buildReactionData(_type, _id, "like")),
		tgbotapi.NewInlineKeyboardButtonData(dissText, buildReactionData(_type, _id, "diss")),
	)
	return tgbotapi.NewInlineKeyboardMarkup(row)
}

func saveReaction(key string, user int) (_type, _id, reaction string, err error) {
	token := strings.Split(key, ":")
	if len(token) != 3 {
		err = fmt.Errorf("react data error: %s", key)
		return
	}
	_type = token[0]
	_id = token[1]
	reaction = token[2]

	pipe := redisClient.Pipeline()
	switch reaction {
	case "like":
		likeCount := pipe.SAdd(buildReactionKey(_type, _id, "like"), strconv.Itoa(user))
		dissCount := pipe.SRem(buildReactionKey(_type, _id, "diss"), strconv.Itoa(user))
		_, err = pipe.Exec()
		if err == nil {
			if likeCount.Val()+dissCount.Val() == 0 {
				err = fmt.Errorf("not modified")
			}
		}
	case "diss":
		dissCount := pipe.SAdd(buildReactionKey(_type, _id, "diss"), strconv.Itoa(user))
		likeCount := pipe.SRem(buildReactionKey(_type, _id, "like"), strconv.Itoa(user))
		_, err = pipe.Exec()
		if err == nil {
			if likeCount.Val()+dissCount.Val() == 0 {
				err = fmt.Errorf("not modified")
			}
		}
	default:
		err = fmt.Errorf("react type error: %s", key)
	}
	return
}
<|endoftext|>


<commit_before>package main

const (
	GCVIS_TMPL = `

<title>gcvis - {{ .Title }}</title>

{{ .Title }}

	json

The smaller plot is linked to the main plot, so it acts as an overview. Try dragging a selection on either plot, and watch the behavior of the other.

Legend

gc.heapinuse: heap in use after gc
scvg.inuse: virtual memory considered in use by the scavenger
scvg.idle: virtual memory considered unused by the scavenger
scvg.sys: virtual memory requested from the operating system (should aproximate VSS)
scvg.released: virtual memory returned to the operating system by the scavenger
scvg.consumed: virtual memory in use (should roughly match process RSS)

	`
)
<commit_msg>Add units on graph axis<commit_after>package main

const (
	GCVIS_TMPL = `

<title>gcvis - {{ .Title }}</title>

{{ .Title }}

	json

The smaller plot is linked to the main plot, so it acts as an overview. Try dragging a selection on either plot, and watch the behavior of the other.

Legend

gc.heapinuse: heap in use after gc
scvg.inuse: virtual memory considered in use by the scavenger
scvg.idle: virtual memory considered unused by the scavenger
scvg.sys: virtual memory requested from the operating system (should aproximate VSS)
scvg.released: virtual memory returned to the operating system by the scavenger
scvg.consumed: virtual memory in use (should roughly match process RSS)

	`
)
<|endoftext|>


<commit_before>package command

import (
	"testing"

	"gomock.googlecode.com/hg/gomock"

	"testmatcher"
)

const (
	cmdText = "give"
	desc    = "Gives x y to player z"
	usage   = "Usage text"
	message = "/give 1 64 admin"
)

func TestCommandFramework(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()

	mockPlayer := NewMockICommandHandler(mockCtrl)

	cf := NewCommandFramework("/")

	mockPlayer.EXPECT().BroadcastMessage("§dthis is a broadcast", true)
	cf.Process("/say this is a broadcast", mockPlayer)

	mockPlayer.EXPECT().GiveItem(1, 64, 0)
	cf.Process("/give 1 64", mockPlayer)

	mockPlayer.EXPECT().SendMessageToPlayer(&testmatcher.StringPrefix{"Commands:"})
	cf.Process("/help", mockPlayer)

	gomock.InOrder(
		mockPlayer.EXPECT().SendMessageToPlayer("Command: /help"),
		mockPlayer.EXPECT().SendMessageToPlayer("Usage: help|?"),
		mockPlayer.EXPECT().SendMessageToPlayer("Description: Shows a list of all commands."),
	)
	cf.Process("/help help", mockPlayer)
}
<commit_msg>Removed some unused test constants.<commit_after>package command

import (
	"testing"

	"gomock.googlecode.com/hg/gomock"

	"testmatcher"
)

func TestCommandFramework(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()

	mockPlayer := NewMockICommandHandler(mockCtrl)

	cf := NewCommandFramework("/")

	mockPlayer.EXPECT().BroadcastMessage("§dthis is a broadcast", true)
	cf.Process("/say this is a broadcast", mockPlayer)

	mockPlayer.EXPECT().GiveItem(1, 64, 0)
	cf.Process("/give 1 64", mockPlayer)

	mockPlayer.EXPECT().SendMessageToPlayer(&testmatcher.StringPrefix{"Commands:"})
	cf.Process("/help", mockPlayer)

	gomock.InOrder(
		mockPlayer.EXPECT().SendMessageToPlayer("Command: /help"),
		mockPlayer.EXPECT().SendMessageToPlayer("Usage: help|?"),
		mockPlayer.EXPECT().SendMessageToPlayer("Description: Shows a list of all commands."),
	)
	cf.Process("/help help", mockPlayer)
}
<|endoftext|>