<commit_before>package main

import (
	"fmt"
	"io"
	"net"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"sync"
	"time"

	"github.com/howbazaar/loggo"
	"github.com/hpcloud/tail"
	"github.com/papertrail/remote_syslog2/syslog"
	"github.com/papertrail/remote_syslog2/utils"
)

var (
	log = loggo.GetLogger("")
	_   = fmt.Printf
)

type Server struct {
	config   *Config
	logger   *syslog.Logger
	registry *WorkerRegistry
	stopChan chan struct{}
	stopped  bool
	mu       sync.RWMutex
}

func NewServer(config *Config) *Server {
	return &Server{
		config:   config,
		registry: NewWorkerRegistry(),
		stopChan: make(chan struct{}),
	}
}

func (s *Server) Start() error {
	if err := s.config.Validate(); err != nil {
		return err
	}

	if !s.config.NoDetach {
		utils.Daemonize(s.config.DebugLogFile, s.config.PidFile)
	}

	loggo.ConfigureLoggers(s.config.LogLevels)

	raddr := net.JoinHostPort(s.config.Destination.Host, strconv.Itoa(s.config.Destination.Port))
	log.Infof("Connecting to %s over %s", raddr, s.config.Destination.Protocol)

	var err error
	s.logger, err = syslog.Dial(
		s.config.Hostname,
		s.config.Destination.Protocol,
		raddr, s.config.RootCAs,
		s.config.ConnectTimeout,
		s.config.WriteTimeout,
		s.config.TcpMaxLineLength,
	)
	if err != nil {
		log.Errorf("Cannot connect to server: %v", err)
		// If the first dial fails, the connection will stay uninitialized (nil)
		// and all subsequent attempts to reconnect will also fail.
		// Instead of continuing, bail out at this error
		return err
	}

	go s.tailFiles()

	for err = range s.logger.Errors {
		log.Errorf("Syslog error: %v", err)
	}

	return nil
}

func (s *Server) Close() {
	s.mu.Lock()
	defer s.mu.Unlock()

	if !s.stopped {
		s.stopped = true
		s.stopChan <- struct{}{}

		log.Infof("Shutting down...")
		s.logger.Close()
	}
}

func (s *Server) closing() bool {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.stopped
}

// Tails a single file
func (s *Server) tailOne(file, tag string, whence int) {
	defer s.registry.Remove(file)
	s.registry.Add(file)

	t, err := tail.TailFile(file, tail.Config{
		ReOpen:    true,
		Follow:    true,
		MustExist: true,
		Poll:      s.config.Poll,
		Location:  &tail.SeekInfo{0, whence},
	})

	if err != nil {
		log.Errorf("%s", err)
		return
	}

	if tag == "" {
		tag = path.Base(file)
	}

	for {
		select {
		case line := <-t.Lines:
			if s.closing() {
				return
			}

			if !matchExps(line.Text, s.config.ExcludePatterns) {
				log.Tracef("foo: %s", line.Text)
				err := s.logger.Write(syslog.Packet{
					Severity: s.config.Severity,
					Facility: s.config.Facility,
					Time:     time.Now(),
					Hostname: s.logger.ClientHostname,
					Tag:      tag,
					Message:  line.Text,
				})

				if err != nil {
					log.Errorf("%s", err)
					return
				}

				log.Tracef("Forwarding line: %s", line.Text)

			} else {
				log.Tracef("Not Forwarding line: %s", line.Text)
			}

		case <-s.stopChan:
			return
		}
	}

	return
}

// Tails files speficied in the globs and re-evaluates the globs
// at the specified interval
func (s *Server) tailFiles() {
	log.Debugf("Evaluating globs every %s", s.config.NewFileCheckInterval)
	firstPass := true

	for {
		if s.closing() {
			return
		}

		s.globFiles(firstPass)
		time.Sleep(s.config.NewFileCheckInterval)
		firstPass = false
	}
}

//
func (s *Server) globFiles(firstPass bool) {
	log.Debugf("Evaluating file globs")
	for _, glob := range s.config.Files {

		tag := glob.Tag
		files, err := filepath.Glob(utils.ResolvePath(glob.Path))

		if err != nil {
			log.Errorf("Failed to glob %s: %s", glob.Path, err)
		} else if files == nil && firstPass {
			log.Errorf("Cannot forward %s, it may not exist", glob.Path)
		}

		for _, file := range files {
			switch {
			case s.registry.Exists(file):
				log.Debugf("Skipping %s because it is already running", file)
			case matchExps(file, s.config.ExcludeFiles):
				log.Debugf("Skipping %s because it is excluded by regular expression", file)
			default:
				log.Infof("Forwarding file: %s", file)

				whence := io.SeekStart

				// don't read the entire file on startup
				if firstPass {
					whence = io.SeekEnd
				}

				go s.tailOne(file, tag, whence)
			}
		}
	}
}

// Evaluates each regex against the string. If any one is a match
// the function returns true, otherwise it returns false
func matchExps(value string, expressions []*regexp.Regexp) bool {
	for _, exp := range expressions {
		if exp.MatchString(value) {
			return true
		}
	}
	return false
}

func main() {
	c, err := NewConfigFromEnv()
	if err != nil {
		log.Criticalf("Failed to configure the application: %s", err)
		os.Exit(1)
	}

	utils.AddSignalHandlers()

	s := NewServer(c)
	if err = s.Start(); err != nil {
		log.Criticalf("Failed to start server: %v", err)
		os.Exit(255)
	}
}
<commit_msg>Remove errant log line<commit_after>package main

import (
	"fmt"
	"io"
	"net"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"sync"
	"time"

	"github.com/howbazaar/loggo"
	"github.com/hpcloud/tail"
	"github.com/papertrail/remote_syslog2/syslog"
	"github.com/papertrail/remote_syslog2/utils"
)

var (
	log = loggo.GetLogger("")
	_   = fmt.Printf
)

type Server struct {
	config   *Config
	logger   *syslog.Logger
	registry *WorkerRegistry
	stopChan chan struct{}
	stopped  bool
	mu       sync.RWMutex
}

func NewServer(config *Config) *Server {
	return &Server{
		config:   config,
		registry: NewWorkerRegistry(),
		stopChan: make(chan struct{}),
	}
}

func (s *Server) Start() error {
	if err := s.config.Validate(); err != nil {
		return err
	}

	if !s.config.NoDetach {
		utils.Daemonize(s.config.DebugLogFile, s.config.PidFile)
	}

	loggo.ConfigureLoggers(s.config.LogLevels)

	raddr := net.JoinHostPort(s.config.Destination.Host, strconv.Itoa(s.config.Destination.Port))
	log.Infof("Connecting to %s over %s", raddr, s.config.Destination.Protocol)

	var err error
	s.logger, err = syslog.Dial(
		s.config.Hostname,
		s.config.Destination.Protocol,
		raddr, s.config.RootCAs,
		s.config.ConnectTimeout,
		s.config.WriteTimeout,
		s.config.TcpMaxLineLength,
	)
	if err != nil {
		log.Errorf("Cannot connect to server: %v", err)
		// If the first dial fails, the connection will stay uninitialized (nil)
		// and all subsequent attempts to reconnect will also fail.
		// Instead of continuing, bail out at this error
		return err
	}

	go s.tailFiles()

	for err = range s.logger.Errors {
		log.Errorf("Syslog error: %v", err)
	}

	return nil
}

func (s *Server) Close() {
	s.mu.Lock()
	defer s.mu.Unlock()

	if !s.stopped {
		s.stopped = true
		s.stopChan <- struct{}{}

		log.Infof("Shutting down...")
		s.logger.Close()
	}
}

func (s *Server) closing() bool {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.stopped
}

// Tails a single file
func (s *Server) tailOne(file, tag string, whence int) {
	defer s.registry.Remove(file)
	s.registry.Add(file)

	t, err := tail.TailFile(file, tail.Config{
		ReOpen:    true,
		Follow:    true,
		MustExist: true,
		Poll:      s.config.Poll,
		Location:  &tail.SeekInfo{0, whence},
	})

	if err != nil {
		log.Errorf("%s", err)
		return
	}

	if tag == "" {
		tag = path.Base(file)
	}

	for {
		select {
		case line := <-t.Lines:
			if s.closing() {
				return
			}

			if !matchExps(line.Text, s.config.ExcludePatterns) {
				err := s.logger.Write(syslog.Packet{
					Severity: s.config.Severity,
					Facility: s.config.Facility,
					Time:     time.Now(),
					Hostname: s.logger.ClientHostname,
					Tag:      tag,
					Message:  line.Text,
				})

				if err != nil {
					log.Errorf("%s", err)
					return
				}

				log.Tracef("Forwarding line: %s", line.Text)

			} else {
				log.Tracef("Not Forwarding line: %s", line.Text)
			}

		case <-s.stopChan:
			return
		}
	}

	return
}

// Tails files speficied in the globs and re-evaluates the globs
// at the specified interval
func (s *Server) tailFiles() {
	log.Debugf("Evaluating globs every %s", s.config.NewFileCheckInterval)
	firstPass := true

	for {
		if s.closing() {
			return
		}

		s.globFiles(firstPass)
		time.Sleep(s.config.NewFileCheckInterval)
		firstPass = false
	}
}

//
func (s *Server) globFiles(firstPass bool) {
	log.Debugf("Evaluating file globs")
	for _, glob := range s.config.Files {

		tag := glob.Tag
		files, err := filepath.Glob(utils.ResolvePath(glob.Path))

		if err != nil {
			log.Errorf("Failed to glob %s: %s", glob.Path, err)
		} else if files == nil && firstPass {
			log.Errorf("Cannot forward %s, it may not exist", glob.Path)
		}

		for _, file := range files {
			switch {
			case s.registry.Exists(file):
				log.Debugf("Skipping %s because it is already running", file)
			case matchExps(file, s.config.ExcludeFiles):
				log.Debugf("Skipping %s because it is excluded by regular expression", file)
			default:
				log.Infof("Forwarding file: %s", file)

				whence := io.SeekStart

				// don't read the entire file on startup
				if firstPass {
					whence = io.SeekEnd
				}

				go s.tailOne(file, tag, whence)
			}
		}
	}
}

// Evaluates each regex against the string. If any one is a match
// the function returns true, otherwise it returns false
func matchExps(value string, expressions []*regexp.Regexp) bool {
	for _, exp := range expressions {
		if exp.MatchString(value) {
			return true
		}
	}
	return false
}

func main() {
	c, err := NewConfigFromEnv()
	if err != nil {
		log.Criticalf("Failed to configure the application: %s", err)
		os.Exit(1)
	}

	utils.AddSignalHandlers()

	s := NewServer(c)
	if err = s.Start(); err != nil {
		log.Criticalf("Failed to start server: %v", err)
		os.Exit(255)
	}
}
<|endoftext|>


<commit_before>package render

import (
	"github.com/disintegration/gift"
	"image"
	"image/color"
	"math"
)

// Modifications enum
const (
	F_FlipX = iota
	F_FlipY
	F_ApplyColor
	F_FillMask
	F_ApplyMask
	F_Rotate
	F_Scale
)

// Modifiables are Renderables that have functions to change their
// underlying image.
// This may be replaced with the gift library down the line
type Modifiable interface {
	Renderable
	FlipX()
	FlipY()
	ApplyColor(c color.Color)
	Copy() Modifiable
	FillMask(img image.RGBA)
	ApplyMask(img image.RGBA)
	Rotate(degrees int)
	Scale(xRatio float64, yRatio float64)
}

// FlipX returns a new rgba which is flipped
// over the horizontal axis.
func FlipX(rgba *image.RGBA) *image.RGBA {
	bounds := rgba.Bounds()
	w := bounds.Max.X
	h := bounds.Max.Y
	newRgba := image.NewRGBA(image.Rect(0, 0, w, h))
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			newRgba.Set(x, y, rgba.At(w-x, y))
		}
	}
	return newRgba
}

// FlipY returns a new rgba which is flipped
// over the vertical axis.
func FlipY(rgba *image.RGBA) *image.RGBA {
	bounds := rgba.Bounds()
	w := bounds.Max.X
	h := bounds.Max.Y
	newRgba := image.NewRGBA(image.Rect(0, 0, w, h))
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			newRgba.Set(x, y, rgba.At(x, h-y))
		}
	}
	return newRgba
}

// Apply color mixes a color into the rgba values of an image
// and returns that new rgba.
func ApplyColor(rgba *image.RGBA, c color.Color) *image.RGBA {
	r1, g1, b1, a1 := c.RGBA()
	bounds := rgba.Bounds()
	w := bounds.Max.X
	h := bounds.Max.Y
	newRgba := image.NewRGBA(image.Rect(0, 0, w, h))
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			r2, g2, b2, a2 := rgba.At(x, y).RGBA()
			a3 := a1 + a2
			if a3 == 0 {
				newRgba.Set(x, y, color.RGBA{0, 0, 0, 0})
				continue
			}
			tmp := color.RGBA64{
				uint16(((a1 * r1) + (a2 * r2)) / a3),
				uint16(((a1 * g1) + (a2 * g2)) / a3),
				uint16(((a1 * b1) + (a2 * b2)) / a3),
				uint16(a2)}
			newRgba.Set(x, y, tmp)
		}
	}
	return newRgba
}

// FillMask replaces alpha 0 pixels in an RGBA with corresponding
// pixels in a second RGBA.
func FillMask(rgba *image.RGBA, img image.RGBA) *image.RGBA {
	// Instead of static color it just two buffers melding
	bounds := rgba.Bounds()
	w := bounds.Max.X
	h := bounds.Max.Y
	newRgba := image.NewRGBA(image.Rect(0, 0, w, h))
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			newRgba.Set(x, y, rgba.At(x, y))
		}
	}
	bounds = img.Bounds()
	w = bounds.Max.X
	h = bounds.Max.Y
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			r1, g1, b1, a1 := rgba.At(x, y).RGBA()
			r2, g2, b2, a2 := img.At(x, y).RGBA()

			var tmp color.RGBA64

			if a1 == 0 {
				tmp = color.RGBA64{
					uint16(r2),
					uint16(g2),
					uint16(b2),
					uint16(a2),
				}
			} else {
				tmp = color.RGBA64{
					uint16(r1),
					uint16(g1),
					uint16(b1),
					uint16(a1),
				}
			}

			newRgba.Set(x, y, tmp)
		}
	}
	return newRgba
}

// ApplyMask mixes the rgba values of two images, according to
// their alpha levels, and returns that as a new rgba.
func ApplyMask(rgba *image.RGBA, img image.RGBA) *image.RGBA {
	// Instead of static color it just two buffers melding
	bounds := rgba.Bounds()
	w := bounds.Max.X
	h := bounds.Max.Y
	newRgba := image.NewRGBA(image.Rect(0, 0, w, h))
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			newRgba.Set(x, y, rgba.At(x, y))
		}
	}
	bounds = img.Bounds()
	w = bounds.Max.X
	h = bounds.Max.Y
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			r1, g1, b1, a1 := img.At(x, y).RGBA()
			r2, g2, b2, a2 := rgba.At(x, y).RGBA()

			var tmp color.RGBA64

			a3 := a1 + a2
			if a3 == 0 {
				tmp = color.RGBA64{
					0, 0, 0, 0,
				}
				newRgba.Set(x, y, tmp)
				continue
			}

			tmp = color.RGBA64{
				uint16(((a1 * r1) + (a2 * r2)) / a3),
				uint16(((a1 * g1) + (a2 * g2)) / a3),
				uint16(((a1 * b1) + (a2 * b2)) / a3),
				uint16(math.Max(float64(a1), float64(a2)))}

			newRgba.Set(x, y, tmp)
		}
	}
	return newRgba
}

// Rotate returns a rotated rgba.
func Rotate(rgba *image.RGBA, degrees int) *image.RGBA {
	filter := gift.New(
		gift.Rotate(float32(degrees), color.Black, gift.CubicInterpolation))
	dst := image.NewRGBA(filter.Bounds(rgba.Bounds()))
	filter.Draw(dst, rgba)
	return dst

}

// Scale returns a scaled rgba.
func Scale(rgba *image.RGBA, xRatio float64, yRatio float64) *image.RGBA {
	bounds := rgba.Bounds()
	w := int(math.Floor(float64(bounds.Max.X) * xRatio))
	h := int(math.Floor(float64(bounds.Max.Y) * yRatio))
	newRgba := image.NewRGBA(image.Rect(0, 0, w, h))
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			newRgba.Set(x, y, rgba.At(int(math.Floor(float64(x)/xRatio)), int(math.Floor(float64(y)/yRatio))))
		}
	}
	return newRgba
}

func round(f float64) int {
	if f < -0.5 {
		return int(f - 0.5)
	}
	if f > 0.5 {
		return int(f + 0.5)
	}
	return 0
}
<commit_msg>Update apply color alpha<commit_after>package render

import (
	"github.com/disintegration/gift"
	"image"
	"image/color"
	"math"
)

// Modifications enum
const (
	F_FlipX = iota
	F_FlipY
	F_ApplyColor
	F_FillMask
	F_ApplyMask
	F_Rotate
	F_Scale
)

// Modifiables are Renderables that have functions to change their
// underlying image.
// This may be replaced with the gift library down the line
type Modifiable interface {
	Renderable
	FlipX()
	FlipY()
	ApplyColor(c color.Color)
	Copy() Modifiable
	FillMask(img image.RGBA)
	ApplyMask(img image.RGBA)
	Rotate(degrees int)
	Scale(xRatio float64, yRatio float64)
}

// FlipX returns a new rgba which is flipped
// over the horizontal axis.
func FlipX(rgba *image.RGBA) *image.RGBA {
	bounds := rgba.Bounds()
	w := bounds.Max.X
	h := bounds.Max.Y
	newRgba := image.NewRGBA(image.Rect(0, 0, w, h))
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			newRgba.Set(x, y, rgba.At(w-x, y))
		}
	}
	return newRgba
}

// FlipY returns a new rgba which is flipped
// over the vertical axis.
func FlipY(rgba *image.RGBA) *image.RGBA {
	bounds := rgba.Bounds()
	w := bounds.Max.X
	h := bounds.Max.Y
	newRgba := image.NewRGBA(image.Rect(0, 0, w, h))
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			newRgba.Set(x, y, rgba.At(x, h-y))
		}
	}
	return newRgba
}

// Apply color mixes a color into the rgba values of an image
// and returns that new rgba.
func ApplyColor(rgba *image.RGBA, c color.Color) *image.RGBA {
	r1, g1, b1, a1 := c.RGBA()
	bounds := rgba.Bounds()
	w := bounds.Max.X
	h := bounds.Max.Y
	newRgba := image.NewRGBA(image.Rect(0, 0, w, h))
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			r2, g2, b2, a2 := rgba.At(x, y).RGBA()
			a3 := a1 + a2
			if a2 == 0 {
				newRgba.Set(x, y, color.RGBA{0, 0, 0, 0})
				continue
			}
			tmp := color.RGBA64{
				uint16(((a1 * r1) + (a2 * r2)) / a3),
				uint16(((a1 * g1) + (a2 * g2)) / a3),
				uint16(((a1 * b1) + (a2 * b2)) / a3),
				uint16(a2)}
			newRgba.Set(x, y, tmp)
		}
	}
	return newRgba
}

// FillMask replaces alpha 0 pixels in an RGBA with corresponding
// pixels in a second RGBA.
func FillMask(rgba *image.RGBA, img image.RGBA) *image.RGBA {
	// Instead of static color it just two buffers melding
	bounds := rgba.Bounds()
	w := bounds.Max.X
	h := bounds.Max.Y
	newRgba := image.NewRGBA(image.Rect(0, 0, w, h))
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			newRgba.Set(x, y, rgba.At(x, y))
		}
	}
	bounds = img.Bounds()
	w = bounds.Max.X
	h = bounds.Max.Y
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			r1, g1, b1, a1 := rgba.At(x, y).RGBA()
			r2, g2, b2, a2 := img.At(x, y).RGBA()

			var tmp color.RGBA64

			if a1 == 0 {
				tmp = color.RGBA64{
					uint16(r2),
					uint16(g2),
					uint16(b2),
					uint16(a2),
				}
			} else {
				tmp = color.RGBA64{
					uint16(r1),
					uint16(g1),
					uint16(b1),
					uint16(a1),
				}
			}

			newRgba.Set(x, y, tmp)
		}
	}
	return newRgba
}

// ApplyMask mixes the rgba values of two images, according to
// their alpha levels, and returns that as a new rgba.
func ApplyMask(rgba *image.RGBA, img image.RGBA) *image.RGBA {
	// Instead of static color it just two buffers melding
	bounds := rgba.Bounds()
	w := bounds.Max.X
	h := bounds.Max.Y
	newRgba := image.NewRGBA(image.Rect(0, 0, w, h))
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			newRgba.Set(x, y, rgba.At(x, y))
		}
	}
	bounds = img.Bounds()
	w = bounds.Max.X
	h = bounds.Max.Y
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			r1, g1, b1, a1 := img.At(x, y).RGBA()
			r2, g2, b2, a2 := rgba.At(x, y).RGBA()

			var tmp color.RGBA64

			a3 := a1 + a2
			if a3 == 0 {
				tmp = color.RGBA64{
					0, 0, 0, 0,
				}
				newRgba.Set(x, y, tmp)
				continue
			}

			tmp = color.RGBA64{
				uint16(((a1 * r1) + (a2 * r2)) / a3),
				uint16(((a1 * g1) + (a2 * g2)) / a3),
				uint16(((a1 * b1) + (a2 * b2)) / a3),
				uint16(math.Max(float64(a1), float64(a2)))}

			newRgba.Set(x, y, tmp)
		}
	}
	return newRgba
}

// Rotate returns a rotated rgba.
func Rotate(rgba *image.RGBA, degrees int) *image.RGBA {
	filter := gift.New(
		gift.Rotate(float32(degrees), color.Black, gift.CubicInterpolation))
	dst := image.NewRGBA(filter.Bounds(rgba.Bounds()))
	filter.Draw(dst, rgba)
	return dst

}

// Scale returns a scaled rgba.
func Scale(rgba *image.RGBA, xRatio float64, yRatio float64) *image.RGBA {
	bounds := rgba.Bounds()
	w := int(math.Floor(float64(bounds.Max.X) * xRatio))
	h := int(math.Floor(float64(bounds.Max.Y) * yRatio))
	newRgba := image.NewRGBA(image.Rect(0, 0, w, h))
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			newRgba.Set(x, y, rgba.At(int(math.Floor(float64(x)/xRatio)), int(math.Floor(float64(y)/yRatio))))
		}
	}
	return newRgba
}

func round(f float64) int {
	if f < -0.5 {
		return int(f - 0.5)
	}
	if f > 0.5 {
		return int(f + 0.5)
	}
	return 0
}
<|endoftext|>


<commit_before>package server

import (
	"github.com/siddontang/ledisdb/ledis"
	"strconv"
	"strings"
)

func getCommand(c *client) error {
	args := c.args
	if len(args) != 1 {
		return ErrCmdParams
	}

	if v, err := c.db.Get(args[0]); err != nil {
		return err
	} else {
		c.resp.writeBulk(v)
	}
	return nil
}

func setCommand(c *client) error {
	args := c.args
	if len(args) != 2 {
		return ErrCmdParams
	}

	if err := c.db.Set(args[0], args[1]); err != nil {
		return err
	} else {
		c.resp.writeStatus(OK)
	}

	return nil
}

func getsetCommand(c *client) error {
	args := c.args
	if len(args) != 2 {
		return ErrCmdParams
	}

	if v, err := c.db.GetSet(args[0], args[1]); err != nil {
		return err
	} else {
		c.resp.writeBulk(v)
	}

	return nil
}

func setnxCommand(c *client) error {
	args := c.args
	if len(args) != 2 {
		return ErrCmdParams
	}

	if n, err := c.db.SetNX(args[0], args[1]); err != nil {
		return err
	} else {
		c.resp.writeInteger(n)
	}

	return nil
}

func existsCommand(c *client) error {
	args := c.args
	if len(args) != 1 {
		return ErrCmdParams
	}

	if n, err := c.db.Exists(args[0]); err != nil {
		return err
	} else {
		c.resp.writeInteger(n)
	}

	return nil
}

func incrCommand(c *client) error {
	args := c.args
	if len(args) != 1 {
		return ErrCmdParams
	}

	if n, err := c.db.Incr(c.args[0]); err != nil {
		return err
	} else {
		c.resp.writeInteger(n)
	}

	return nil
}

func decrCommand(c *client) error {
	args := c.args
	if len(args) != 1 {
		return ErrCmdParams
	}

	if n, err := c.db.Decr(c.args[0]); err != nil {
		return err
	} else {
		c.resp.writeInteger(n)
	}

	return nil
}

func incrbyCommand(c *client) error {
	args := c.args
	if len(args) != 2 {
		return ErrCmdParams
	}

	delta, err := ledis.StrInt64(args[1], nil)
	if err != nil {
		return ErrValue
	}

	if n, err := c.db.IncrBy(c.args[0], delta); err != nil {
		return err
	} else {
		c.resp.writeInteger(n)
	}

	return nil
}

func decrbyCommand(c *client) error {
	args := c.args
	if len(args) != 2 {
		return ErrCmdParams
	}

	delta, err := ledis.StrInt64(args[1], nil)
	if err != nil {
		return ErrValue
	}

	if n, err := c.db.DecrBy(c.args[0], delta); err != nil {
		return err
	} else {
		c.resp.writeInteger(n)
	}

	return nil
}

func delCommand(c *client) error {
	args := c.args
	if len(args) == 0 {
		return ErrCmdParams
	}

	if n, err := c.db.Del(args...); err != nil {
		return err
	} else {
		c.resp.writeInteger(n)
	}

	return nil
}

func msetCommand(c *client) error {
	args := c.args
	if len(args) == 0 || len(args)%2 != 0 {
		return ErrCmdParams
	}

	kvs := make([]ledis.KVPair, len(args)/2)
	for i := 0; i < len(kvs); i++ {
		kvs[i].Key = args[2*i]
		kvs[i].Value = args[2*i+1]
	}

	if err := c.db.MSet(kvs...); err != nil {
		return err
	} else {
		c.resp.writeStatus(OK)
	}

	return nil
}

// func setexCommand(c *client) error {
// 	return nil
// }

func mgetCommand(c *client) error {
	args := c.args
	if len(args) == 0 {
		return ErrCmdParams
	}

	if v, err := c.db.MGet(args...); err != nil {
		return err
	} else {
		c.resp.writeSliceArray(v)
	}

	return nil
}

func expireCommand(c *client) error {
	args := c.args
	if len(args) != 2 {
		return ErrCmdParams
	}

	duration, err := ledis.StrInt64(args[1], nil)
	if err != nil {
		return ErrValue
	}

	if v, err := c.db.Expire(args[0], duration); err != nil {
		return err
	} else {
		c.resp.writeInteger(v)
	}

	return nil
}

func expireAtCommand(c *client) error {
	args := c.args
	if len(args) != 2 {
		return ErrCmdParams
	}

	when, err := ledis.StrInt64(args[1], nil)
	if err != nil {
		return ErrValue
	}

	if v, err := c.db.ExpireAt(args[0], when); err != nil {
		return err
	} else {
		c.resp.writeInteger(v)
	}

	return nil
}

func ttlCommand(c *client) error {
	args := c.args
	if len(args) != 1 {
		return ErrCmdParams
	}

	if v, err := c.db.TTL(args[0]); err != nil {
		return err
	} else {
		c.resp.writeInteger(v)
	}

	return nil
}

func persistCommand(c *client) error {
	args := c.args
	if len(args) != 1 {
		return ErrCmdParams
	}

	if n, err := c.db.Persist(args[0]); err != nil {
		return err
	} else {
		c.resp.writeInteger(n)
	}

	return nil
}

func parseScanArgs(c *client) (key []byte, match string, count int, err error) {
	args := c.args
	count = 10

	switch len(args) {
	case 0:
		key = nil
		return
	case 1, 3, 5:
		key = args[0]
		break
	default:
		err = ErrCmdParams
		return
	}

	if len(args) == 3 {
		switch strings.ToLower(ledis.String(args[1])) {
		case "match":
			match = ledis.String(args[2])
			return
		case "count":
			count, err = strconv.Atoi(ledis.String(args[2]))
			return
		default:
			err = ErrCmdParams
			return
		}
	} else if len(args) == 5 {
		if strings.ToLower(ledis.String(args[1])) != "match" {
			err = ErrCmdParams
			return
		} else if strings.ToLower(ledis.String(args[3])) != "count" {
			err = ErrCmdParams
			return
		}

		match = ledis.String(args[2])
		count, err = strconv.Atoi(ledis.String(args[4]))
		return
	}

	return
}

func xscanCommand(c *client) error {
	key, match, count, err := parseScanArgs(c)
	if err != nil {
		return err
	}

	if ay, err := c.db.Scan(key, count, false, match); err != nil {
		return err
	} else {
		data := make([]interface{}, 2)
		if len(ay) < count {
			data[0] = []byte("")
		} else {
			data[0] = ay[len(ay)-1]
		}
		data[1] = ay
		c.resp.writeArray(data)
	}
	return nil
}

func init() {
	register("decr", decrCommand)
	register("decrby", decrbyCommand)
	register("del", delCommand)
	register("exists", existsCommand)
	register("get", getCommand)
	register("getset", getsetCommand)
	register("incr", incrCommand)
	register("incrby", incrbyCommand)
	register("mget", mgetCommand)
	register("mset", msetCommand)
	register("set", setCommand)
	register("setnx", setnxCommand)
	register("expire", expireCommand)
	register("expireat", expireAtCommand)
	register("ttl", ttlCommand)
	register("persist", persistCommand)
	register("xscan", xscanCommand)
}
<commit_msg>bugfix scan count 0<commit_after>package server

import (
	"github.com/siddontang/ledisdb/ledis"
	"strconv"
	"strings"
)

func getCommand(c *client) error {
	args := c.args
	if len(args) != 1 {
		return ErrCmdParams
	}

	if v, err := c.db.Get(args[0]); err != nil {
		return err
	} else {
		c.resp.writeBulk(v)
	}
	return nil
}

func setCommand(c *client) error {
	args := c.args
	if len(args) != 2 {
		return ErrCmdParams
	}

	if err := c.db.Set(args[0], args[1]); err != nil {
		return err
	} else {
		c.resp.writeStatus(OK)
	}

	return nil
}

func getsetCommand(c *client) error {
	args := c.args
	if len(args) != 2 {
		return ErrCmdParams
	}

	if v, err := c.db.GetSet(args[0], args[1]); err != nil {
		return err
	} else {
		c.resp.writeBulk(v)
	}

	return nil
}

func setnxCommand(c *client) error {
	args := c.args
	if len(args) != 2 {
		return ErrCmdParams
	}

	if n, err := c.db.SetNX(args[0], args[1]); err != nil {
		return err
	} else {
		c.resp.writeInteger(n)
	}

	return nil
}

func existsCommand(c *client) error {
	args := c.args
	if len(args) != 1 {
		return ErrCmdParams
	}

	if n, err := c.db.Exists(args[0]); err != nil {
		return err
	} else {
		c.resp.writeInteger(n)
	}

	return nil
}

func incrCommand(c *client) error {
	args := c.args
	if len(args) != 1 {
		return ErrCmdParams
	}

	if n, err := c.db.Incr(c.args[0]); err != nil {
		return err
	} else {
		c.resp.writeInteger(n)
	}

	return nil
}

func decrCommand(c *client) error {
	args := c.args
	if len(args) != 1 {
		return ErrCmdParams
	}

	if n, err := c.db.Decr(c.args[0]); err != nil {
		return err
	} else {
		c.resp.writeInteger(n)
	}

	return nil
}

func incrbyCommand(c *client) error {
	args := c.args
	if len(args) != 2 {
		return ErrCmdParams
	}

	delta, err := ledis.StrInt64(args[1], nil)
	if err != nil {
		return ErrValue
	}

	if n, err := c.db.IncrBy(c.args[0], delta); err != nil {
		return err
	} else {
		c.resp.writeInteger(n)
	}

	return nil
}

func decrbyCommand(c *client) error {
	args := c.args
	if len(args) != 2 {
		return ErrCmdParams
	}

	delta, err := ledis.StrInt64(args[1], nil)
	if err != nil {
		return ErrValue
	}

	if n, err := c.db.DecrBy(c.args[0], delta); err != nil {
		return err
	} else {
		c.resp.writeInteger(n)
	}

	return nil
}

func delCommand(c *client) error {
	args := c.args
	if len(args) == 0 {
		return ErrCmdParams
	}

	if n, err := c.db.Del(args...); err != nil {
		return err
	} else {
		c.resp.writeInteger(n)
	}

	return nil
}

func msetCommand(c *client) error {
	args := c.args
	if len(args) == 0 || len(args)%2 != 0 {
		return ErrCmdParams
	}

	kvs := make([]ledis.KVPair, len(args)/2)
	for i := 0; i < len(kvs); i++ {
		kvs[i].Key = args[2*i]
		kvs[i].Value = args[2*i+1]
	}

	if err := c.db.MSet(kvs...); err != nil {
		return err
	} else {
		c.resp.writeStatus(OK)
	}

	return nil
}

// func setexCommand(c *client) error {
// 	return nil
// }

func mgetCommand(c *client) error {
	args := c.args
	if len(args) == 0 {
		return ErrCmdParams
	}

	if v, err := c.db.MGet(args...); err != nil {
		return err
	} else {
		c.resp.writeSliceArray(v)
	}

	return nil
}

func expireCommand(c *client) error {
	args := c.args
	if len(args) != 2 {
		return ErrCmdParams
	}

	duration, err := ledis.StrInt64(args[1], nil)
	if err != nil {
		return ErrValue
	}

	if v, err := c.db.Expire(args[0], duration); err != nil {
		return err
	} else {
		c.resp.writeInteger(v)
	}

	return nil
}

func expireAtCommand(c *client) error {
	args := c.args
	if len(args) != 2 {
		return ErrCmdParams
	}

	when, err := ledis.StrInt64(args[1], nil)
	if err != nil {
		return ErrValue
	}

	if v, err := c.db.ExpireAt(args[0], when); err != nil {
		return err
	} else {
		c.resp.writeInteger(v)
	}

	return nil
}

func ttlCommand(c *client) error {
	args := c.args
	if len(args) != 1 {
		return ErrCmdParams
	}

	if v, err := c.db.TTL(args[0]); err != nil {
		return err
	} else {
		c.resp.writeInteger(v)
	}

	return nil
}

func persistCommand(c *client) error {
	args := c.args
	if len(args) != 1 {
		return ErrCmdParams
	}

	if n, err := c.db.Persist(args[0]); err != nil {
		return err
	} else {
		c.resp.writeInteger(n)
	}

	return nil
}

func parseScanArgs(c *client) (key []byte, match string, count int, err error) {
	args := c.args
	count = 10

	switch len(args) {
	case 0:
		key = nil
		return
	case 1, 3, 5:
		key = args[0]
		break
	default:
		err = ErrCmdParams
		return
	}

	if len(args) == 3 {
		switch strings.ToLower(ledis.String(args[1])) {
		case "match":
			match = ledis.String(args[2])
		case "count":
			count, err = strconv.Atoi(ledis.String(args[2]))
		default:
			err = ErrCmdParams
			return
		}
	} else if len(args) == 5 {
		if strings.ToLower(ledis.String(args[1])) != "match" {
			err = ErrCmdParams
			return
		} else if strings.ToLower(ledis.String(args[3])) != "count" {
			err = ErrCmdParams
			return
		}

		match = ledis.String(args[2])
		count, err = strconv.Atoi(ledis.String(args[4]))
	}

	if count <= 0 {
		err = ErrCmdParams
	}

	return
}

func xscanCommand(c *client) error {
	key, match, count, err := parseScanArgs(c)
	if err != nil {
		return err
	}

	if ay, err := c.db.Scan(key, count, false, match); err != nil {
		return err
	} else {
		data := make([]interface{}, 2)
		if len(ay) < count {
			data[0] = []byte("")
		} else {
			data[0] = ay[len(ay)-1]
		}
		data[1] = ay
		c.resp.writeArray(data)
	}
	return nil
}

func init() {
	register("decr", decrCommand)
	register("decrby", decrbyCommand)
	register("del", delCommand)
	register("exists", existsCommand)
	register("get", getCommand)
	register("getset", getsetCommand)
	register("incr", incrCommand)
	register("incrby", incrbyCommand)
	register("mget", mgetCommand)
	register("mset", msetCommand)
	register("set", setCommand)
	register("setnx", setnxCommand)
	register("expire", expireCommand)
	register("expireat", expireAtCommand)
	register("ttl", ttlCommand)
	register("persist", persistCommand)
	register("xscan", xscanCommand)
}
<|endoftext|>


<commit_before>package chunker

import (
	"errors"
	"hash"
	"io"
	"sync"
)

const (
	KiB = 1024
	MiB = 1024 * KiB

	// WindowSize is the size of the sliding window.
	windowSize = 64

	// aim to create chunks of 20 bits or about 1MiB on average.
	averageBits = 20

	// MinSize is the minimal size of a chunk.
	MinSize = 512 * KiB
	// MaxSize is the maximal size of a chunk.
	MaxSize = 8 * MiB

	splitmask = (1 << averageBits) - 1

	chunkerBufSize = 512 * KiB
)

var bufPool = sync.Pool{
	New: func() interface{} { return make([]byte, chunkerBufSize) },
}

type tables struct {
	out [256]Pol
	mod [256]Pol
}

// cache precomputed tables, these are read-only anyway
var cache struct {
	entries map[Pol]*tables
	sync.Mutex
}

func init() {
	cache.entries = make(map[Pol]*tables)
}

// Chunk is one content-dependent chunk of bytes whose end was cut when the
// Rabin Fingerprint had the value stored in Cut.
type Chunk struct {
	Start  uint
	Length uint
	Cut    uint64
	Digest []byte
}

func (c Chunk) Reader(r io.ReaderAt) io.Reader {
	return io.NewSectionReader(r, int64(c.Start), int64(c.Length))
}

// Chunker splits content with Rabin Fingerprints.
type Chunker struct {
	pol      Pol
	polShift uint
	tables   *tables

	rd     io.Reader
	closed bool

	window [windowSize]byte
	wpos   int

	buf  []byte
	bpos uint
	bmax uint

	start uint
	count uint
	pos   uint

	pre uint // wait for this many bytes before start calculating an new chunk

	digest uint64
	h      hash.Hash
}

// New returns a new Chunker based on polynomial p that reads from rd
// with bufsize and pass all data to hash along the way.
func New(rd io.Reader, pol Pol, h hash.Hash) *Chunker {
	c := &Chunker{
		buf: bufPool.Get().([]byte),
		h:   h,
		pol: pol,
		rd:  rd,
	}

	c.reset()

	return c
}

func (c *Chunker) reset() {
	c.polShift = uint(c.pol.Deg() - 8)
	c.fillTables()

	for i := 0; i < windowSize; i++ {
		c.window[i] = 0
	}

	c.closed = false
	c.digest = 0
	c.wpos = 0
	c.count = 0
	c.slide(1)
	c.start = c.pos

	if c.h != nil {
		c.h.Reset()
	}

	// do not start a new chunk unless at least MinSize bytes have been read
	c.pre = MinSize - windowSize
}

// Calculate out_table and mod_table for optimization. Must be called only
// once. This implementation uses a cache in the global variable cache.
func (c *Chunker) fillTables() {
	// if polynomial hasn't been specified, do not compute anything for now
	if c.pol == 0 {
		return
	}

	// test if the tables are cached for this polynomial
	cache.Lock()
	defer cache.Unlock()
	if t, ok := cache.entries[c.pol]; ok {
		c.tables = t
		return
	}

	// else create a new entry
	c.tables = &tables{}
	cache.entries[c.pol] = c.tables

	// calculate table for sliding out bytes. The byte to slide out is used as
	// the index for the table, the value contains the following:
	// out_table[b] = Hash(b || 0 ||        ...        || 0)
	//                          \ windowsize-1 zero bytes /
	// To slide out byte b_0 for window size w with known hash
	// H := H(b_0 || ... || b_w), it is sufficient to add out_table[b_0]:
	//    H(b_0 || ... || b_w) + H(b_0 || 0 || ... || 0)
	//  = H(b_0 + b_0 || b_1 + 0 || ... || b_w + 0)
	//  = H(    0     || b_1 || ...     || b_w)
	//
	// Afterwards a new byte can be shifted in.
	for b := 0; b < 256; b++ {
		var h Pol

		h = appendByte(h, byte(b), c.pol)
		for i := 0; i < windowSize-1; i++ {
			h = appendByte(h, 0, c.pol)
		}
		c.tables.out[b] = h
	}

	// calculate table for reduction mod Polynomial
	k := c.pol.Deg()
	for b := 0; b < 256; b++ {
		// mod_table[b] = A | B, where A = (b(x) * x^k mod pol) and  B = b(x) * x^k
		//
		// The 8 bits above deg(Polynomial) determine what happens next and so
		// these bits are used as a lookup to this table. The value is split in
		// two parts: Part A contains the result of the modulus operation, part
		// B is used to cancel out the 8 top bits so that one XOR operation is
		// enough to reduce modulo Polynomial
		c.tables.mod[b] = Pol(uint64(b)<<uint(k)).Mod(c.pol) | (Pol(b) << uint(k))
	}
}

// Next returns the position and length of the next chunk of data. If an error
// occurs while reading, the error is returned with a nil chunk. The state of
// the current chunk is undefined. When the last chunk has been returned, all
// subsequent calls yield a nil chunk and an io.EOF error.
func (c *Chunker) Next() (*Chunk, error) {
	if c.tables == nil {
		return nil, errors.New("polynomial is not set")
	}

	for {
		if c.bpos >= c.bmax {
			n, err := io.ReadFull(c.rd, c.buf[:])

			if err == io.ErrUnexpectedEOF {
				err = nil
			}

			// io.ReadFull only returns io.EOF when no bytes could be read. If
			// this is the case and we're in this branch, there are no more
			// bytes to buffer, so this was the last chunk. If a different
			// error has occurred, return that error and abandon the current
			// chunk.
			if err == io.EOF && !c.closed {
				c.closed = true

				// return the buffer to the pool
				bufPool.Put(c.buf)

				// return current chunk, if any bytes have been processed
				if c.count > 0 {
					return &Chunk{
						Start:  c.start,
						Length: c.count,
						Cut:    c.digest,
						Digest: c.hashDigest(),
					}, nil
				}
			}

			if err != nil {
				return nil, err
			}

			c.bpos = 0
			c.bmax = uint(n)
		}

		// check if bytes have to be dismissed before starting a new chunk
		if c.pre > 0 {
			n := c.bmax - c.bpos
			if c.pre > uint(n) {
				c.pre -= uint(n)
				c.updateHash(c.buf[c.bpos:c.bmax])

				c.count += uint(n)
				c.pos += uint(n)
				c.bpos = c.bmax

				continue
			}

			c.updateHash(c.buf[c.bpos : c.bpos+c.pre])

			c.bpos += c.pre
			c.count += c.pre
			c.pos += c.pre
			c.pre = 0
		}

		add := c.count
		for _, b := range c.buf[c.bpos:c.bmax] {
			// inline c.slide(b) and append(b) to increase performance
			out := c.window[c.wpos]
			c.window[c.wpos] = b
			c.digest ^= uint64(c.tables.out[out])
			c.wpos = (c.wpos + 1) % windowSize

			// c.append(b)
			index := c.digest >> c.polShift
			c.digest <<= 8
			c.digest |= uint64(b)

			c.digest ^= uint64(c.tables.mod[index])
			// end inline

			add++
			if add < MinSize {
				continue
			}

			if (c.digest&splitmask) == 0 || add >= MaxSize {
				i := add - c.count - 1
				c.updateHash(c.buf[c.bpos : c.bpos+uint(i)+1])
				c.count = add
				c.pos += uint(i) + 1
				c.bpos += uint(i) + 1

				chunk := &Chunk{
					Start:  c.start,
					Length: c.count,
					Cut:    c.digest,
					Digest: c.hashDigest(),
				}

				c.reset()

				return chunk, nil
			}
		}

		steps := c.bmax - c.bpos
		if steps > 0 {
			c.updateHash(c.buf[c.bpos : c.bpos+steps])
		}
		c.count += steps
		c.pos += steps
		c.bpos = c.bmax
	}
}

func (c *Chunker) updateHash(data []byte) {
	if c.h != nil {
		// the hashes from crypto/sha* do not return an error
		_, err := c.h.Write(data)
		if err != nil {
			panic(err)
		}
	}
}

func (c *Chunker) hashDigest() []byte {
	if c.h == nil {
		return nil
	}

	return c.h.Sum(nil)
}

func (c *Chunker) append(b byte) {
	index := c.digest >> c.polShift
	c.digest <<= 8
	c.digest |= uint64(b)

	c.digest ^= uint64(c.tables.mod[index])
}

func (c *Chunker) slide(b byte) {
	out := c.window[c.wpos]
	c.window[c.wpos] = b
	c.digest ^= uint64(c.tables.out[out])
	c.wpos = (c.wpos + 1) % windowSize

	c.append(b)
}

func appendByte(hash Pol, b byte, pol Pol) Pol {
	hash <<= 8
	hash |= Pol(b)

	return hash.Mod(pol)
}
<commit_msg>hack the chunker to work for ipfs chunking<commit_after>package chunker

import (
	"bytes"
	"errors"
	"hash"
	"io"
	"sync"
)

const (
	KiB = 1024
	MiB = 1024 * KiB

	// WindowSize is the size of the sliding window.
	windowSize = 16

	// MinSize is the minimal size of a chunk.
	MinSize = 128 * KiB
	// MaxSize is the maximal size of a chunk.
	MaxSize = 512 * KiB

	chunkerBufSize = 512 * KiB
)

var bufPool = sync.Pool{
	New: func() interface{} { return make([]byte, chunkerBufSize) },
}

type tables struct {
	out [256]Pol
	mod [256]Pol
}

// cache precomputed tables, these are read-only anyway
var cache struct {
	entries map[Pol]*tables
	sync.Mutex
}

func init() {
	cache.entries = make(map[Pol]*tables)
}

// Chunk is one content-dependent chunk of bytes whose end was cut when the
// Rabin Fingerprint had the value stored in Cut.
type Chunk struct {
	Start  uint
	Length uint
	Cut    uint64
	Digest []byte
	Data   []byte
}

func (c Chunk) Reader(r io.ReaderAt) io.Reader {
	return io.NewSectionReader(r, int64(c.Start), int64(c.Length))
}

// Chunker splits content with Rabin Fingerprints.
type Chunker struct {
	pol      Pol
	polShift uint
	tables   *tables

	rd     io.Reader
	closed bool

	chunkbuf *bytes.Buffer

	window [windowSize]byte
	wpos   int

	buf  []byte
	bpos uint
	bmax uint

	start uint
	count uint
	pos   uint

	pre uint // wait for this many bytes before start calculating an new chunk

	digest uint64
	h      hash.Hash

	chunkSize uint64
}

// New returns a new Chunker based on polynomial p that reads from rd
// with bufsize and pass all data to hash along the way.
func New(rd io.Reader, pol Pol, h hash.Hash, sizepow uint) *Chunker {
	c := &Chunker{
		buf:       bufPool.Get().([]byte),
		h:         h,
		pol:       pol,
		rd:        rd,
		chunkbuf:  new(bytes.Buffer),
		chunkSize: (1 << sizepow) - 1,
	}

	c.reset()

	return c
}

func (c *Chunker) reset() {
	c.polShift = uint(c.pol.Deg() - 8)
	c.fillTables()

	for i := 0; i < windowSize; i++ {
		c.window[i] = 0
	}

	c.closed = false
	c.digest = 0
	c.wpos = 0
	c.count = 0
	c.slide(1)
	c.start = c.pos

	if c.h != nil {
		c.h.Reset()
	}

	// do not start a new chunk unless at least MinSize bytes have been read
	c.pre = MinSize - windowSize
}

// Calculate out_table and mod_table for optimization. Must be called only
// once. This implementation uses a cache in the global variable cache.
func (c *Chunker) fillTables() {
	// if polynomial hasn't been specified, do not compute anything for now
	if c.pol == 0 {
		return
	}

	// test if the tables are cached for this polynomial
	cache.Lock()
	defer cache.Unlock()
	if t, ok := cache.entries[c.pol]; ok {
		c.tables = t
		return
	}

	// else create a new entry
	c.tables = &tables{}
	cache.entries[c.pol] = c.tables

	// calculate table for sliding out bytes. The byte to slide out is used as
	// the index for the table, the value contains the following:
	// out_table[b] = Hash(b || 0 ||        ...        || 0)
	//                          \ windowsize-1 zero bytes /
	// To slide out byte b_0 for window size w with known hash
	// H := H(b_0 || ... || b_w), it is sufficient to add out_table[b_0]:
	//    H(b_0 || ... || b_w) + H(b_0 || 0 || ... || 0)
	//  = H(b_0 + b_0 || b_1 + 0 || ... || b_w + 0)
	//  = H(    0     || b_1 || ...     || b_w)
	//
	// Afterwards a new byte can be shifted in.
	for b := 0; b < 256; b++ {
		var h Pol

		h = appendByte(h, byte(b), c.pol)
		for i := 0; i < windowSize-1; i++ {
			h = appendByte(h, 0, c.pol)
		}
		c.tables.out[b] = h
	}

	// calculate table for reduction mod Polynomial
	k := c.pol.Deg()
	for b := 0; b < 256; b++ {
		// mod_table[b] = A | B, where A = (b(x) * x^k mod pol) and  B = b(x) * x^k
		//
		// The 8 bits above deg(Polynomial) determine what happens next and so
		// these bits are used as a lookup to this table. The value is split in
		// two parts: Part A contains the result of the modulus operation, part
		// B is used to cancel out the 8 top bits so that one XOR operation is
		// enough to reduce modulo Polynomial
		c.tables.mod[b] = Pol(uint64(b)<<uint(k)).Mod(c.pol) | (Pol(b) << uint(k))
	}
}

// Next returns the position and length of the next chunk of data. If an error
// occurs while reading, the error is returned with a nil chunk. The state of
// the current chunk is undefined. When the last chunk has been returned, all
// subsequent calls yield a nil chunk and an io.EOF error.
func (c *Chunker) Next() (*Chunk, error) {
	if c.tables == nil {
		return nil, errors.New("polynomial is not set")
	}

	for {
		if c.bpos >= c.bmax {
			n, err := io.ReadFull(c.rd, c.buf[:])
			c.chunkbuf.Write(c.buf[:n])

			if err == io.ErrUnexpectedEOF {
				err = nil
			}

			// io.ReadFull only returns io.EOF when no bytes could be read. If
			// this is the case and we're in this branch, there are no more
			// bytes to buffer, so this was the last chunk. If a different
			// error has occurred, return that error and abandon the current
			// chunk.
			if err == io.EOF && !c.closed {
				c.closed = true

				// return the buffer to the pool
				bufPool.Put(c.buf)

				data := dupBytes(c.chunkbuf.Next(int(c.count)))
				// return current chunk, if any bytes have been processed
				if c.count > 0 {
					return &Chunk{
						Start:  c.start,
						Length: c.count,
						Cut:    c.digest,
						Digest: c.hashDigest(),
						Data:   data,
					}, nil
				}
			}

			if err != nil {
				return nil, err
			}

			c.bpos = 0
			c.bmax = uint(n)
		}

		// check if bytes have to be dismissed before starting a new chunk
		if c.pre > 0 {
			n := c.bmax - c.bpos
			if c.pre > uint(n) {
				c.pre -= uint(n)
				c.updateHash(c.buf[c.bpos:c.bmax])

				c.count += uint(n)
				c.pos += uint(n)
				c.bpos = c.bmax

				continue
			}

			c.updateHash(c.buf[c.bpos : c.bpos+c.pre])

			c.bpos += c.pre
			c.count += c.pre
			c.pos += c.pre
			c.pre = 0
		}

		add := c.count
		for _, b := range c.buf[c.bpos:c.bmax] {
			// inline c.slide(b) and append(b) to increase performance
			out := c.window[c.wpos]
			c.window[c.wpos] = b
			c.digest ^= uint64(c.tables.out[out])
			c.wpos = (c.wpos + 1) % windowSize

			// c.append(b)
			index := c.digest >> c.polShift
			c.digest <<= 8
			c.digest |= uint64(b)

			c.digest ^= uint64(c.tables.mod[index])
			// end inline

			add++
			if add < MinSize {
				continue
			}

			if (c.digest&c.chunkSize) == 0 || add >= MaxSize {
				i := add - c.count - 1
				c.updateHash(c.buf[c.bpos : c.bpos+uint(i)+1])
				c.count = add
				c.pos += uint(i) + 1
				c.bpos += uint(i) + 1

				data := dupBytes(c.chunkbuf.Next(int(c.count)))

				chunk := &Chunk{
					Start:  c.start,
					Length: c.count,
					Cut:    c.digest,
					Digest: c.hashDigest(),
					Data:   data,
				}

				c.reset()

				return chunk, nil
			}
		}

		steps := c.bmax - c.bpos
		if steps > 0 {
			c.updateHash(c.buf[c.bpos : c.bpos+steps])
		}
		c.count += steps
		c.pos += steps
		c.bpos = c.bmax
	}
}

func dupBytes(b []byte) []byte {
	out := make([]byte, len(b))
	copy(out, b)
	return out
}

func (c *Chunker) updateHash(data []byte) {
	if c.h != nil {
		// the hashes from crypto/sha* do not return an error
		_, err := c.h.Write(data)
		if err != nil {
			panic(err)
		}
	}
}

func (c *Chunker) hashDigest() []byte {
	if c.h == nil {
		return nil
	}

	return c.h.Sum(nil)
}

func (c *Chunker) append(b byte) {
	index := c.digest >> c.polShift
	c.digest <<= 8
	c.digest |= uint64(b)

	c.digest ^= uint64(c.tables.mod[index])
}

func (c *Chunker) slide(b byte) {
	out := c.window[c.wpos]
	c.window[c.wpos] = b
	c.digest ^= uint64(c.tables.out[out])
	c.wpos = (c.wpos + 1) % windowSize

	c.append(b)
}

func appendByte(hash Pol, b byte, pol Pol) Pol {
	hash <<= 8
	hash |= Pol(b)

	return hash.Mod(pol)
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package client

import (
	"fmt"

	"golang.org/x/net/context"

	"github.com/keybase/cli"
	"github.com/keybase/client/go/libcmdline"
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
	rpc "github.com/keybase/go-framed-msgpack-rpc"
)

type CmdDeviceRemove struct {
	idOrName string
	force    bool
	libkb.Contextified
}

func (c *CmdDeviceRemove) ParseArgv(ctx *cli.Context) error {
	if len(ctx.Args()) != 1 {
		return fmt.Errorf("Device remove only takes one argument: the device ID or name.")
	}
	c.idOrName = ctx.Args()[0]
	c.force = ctx.Bool("force")
	return nil
}

func (c *CmdDeviceRemove) Run() (err error) {
	var id keybase1.DeviceID
	id, err = keybase1.DeviceIDFromString(c.idOrName)
	if err != nil {
		id, err = c.lookup(c.idOrName)
		if err != nil {
			return err
		}
	}

	cli, err := GetRevokeClient()
	if err != nil {
		return err
	}

	protocols := []rpc.Protocol{
		NewSecretUIProtocol(c.G()),
	}
	if err = RegisterProtocols(protocols); err != nil {
		return err
	}

	return cli.RevokeDevice(context.TODO(), keybase1.RevokeDeviceArg{
		Force:    c.force,
		DeviceID: id,
	})
}

func (c *CmdDeviceRemove) lookup(name string) (keybase1.DeviceID, error) {
	cli, err := GetDeviceClient()
	if err != nil {
		return "", err
	}
	if err := RegisterProtocols(nil); err != nil {
		return "", err
	}
	devs, err := cli.DeviceList(context.TODO(), 0)
	if err != nil {
		return "", err
	}

	for _, dev := range devs {
		if dev.Name == name {
			return dev.DeviceID, nil
		}
	}
	return "", fmt.Errorf("Unknown Device ID or Name")
}

func NewCmdDeviceRemove(cl *libcmdline.CommandLine, g *libkb.GlobalContext) cli.Command {
	return cli.Command{
		Name:         "remove",
		ArgumentHelp: "<id|name>",
		Usage:        "Remove a device",
		Flags: []cli.Flag{
			cli.BoolFlag{
				Name:  "f, force",
				Usage: "Override warning about removing the current device.",
			},
		},
		Action: func(c *cli.Context) {
			cl.ChooseCommand(&CmdDeviceRemove{Contextified: libkb.NewContextified(g)}, "remove", c)
		},
	}
}

func (c *CmdDeviceRemove) GetUsage() libkb.Usage {
	return libkb.Usage{
		Config:     true,
		GpgKeyring: true,
		KbKeyring:  true,
		API:        true,
	}
}
<commit_msg>only register protocols once<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package client

import (
	"fmt"

	"golang.org/x/net/context"

	"github.com/keybase/cli"
	"github.com/keybase/client/go/libcmdline"
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
	rpc "github.com/keybase/go-framed-msgpack-rpc"
)

type CmdDeviceRemove struct {
	idOrName string
	force    bool
	libkb.Contextified
}

func (c *CmdDeviceRemove) ParseArgv(ctx *cli.Context) error {
	if len(ctx.Args()) != 1 {
		return fmt.Errorf("Device remove only takes one argument: the device ID or name.")
	}
	c.idOrName = ctx.Args()[0]
	c.force = ctx.Bool("force")
	return nil
}

func (c *CmdDeviceRemove) Run() (err error) {
	protocols := []rpc.Protocol{
		NewSecretUIProtocol(c.G()),
	}
	if err = RegisterProtocolsWithContext(protocols, c.G()); err != nil {
		return err
	}

	var id keybase1.DeviceID
	id, err = keybase1.DeviceIDFromString(c.idOrName)
	if err != nil {
		id, err = c.lookup(c.idOrName)
		if err != nil {
			return err
		}
	}

	cli, err := GetRevokeClient()
	if err != nil {
		return err
	}

	return cli.RevokeDevice(context.TODO(), keybase1.RevokeDeviceArg{
		Force:    c.force,
		DeviceID: id,
	})
}

func (c *CmdDeviceRemove) lookup(name string) (keybase1.DeviceID, error) {
	cli, err := GetDeviceClient()
	if err != nil {
		return "", err
	}
	devs, err := cli.DeviceList(context.TODO(), 0)
	if err != nil {
		return "", err
	}

	for _, dev := range devs {
		if dev.Name == name {
			return dev.DeviceID, nil
		}
	}
	return "", fmt.Errorf("Unknown Device ID or Name")
}

func NewCmdDeviceRemove(cl *libcmdline.CommandLine, g *libkb.GlobalContext) cli.Command {
	return cli.Command{
		Name:         "remove",
		ArgumentHelp: "<id|name>",
		Usage:        "Remove a device",
		Flags: []cli.Flag{
			cli.BoolFlag{
				Name:  "f, force",
				Usage: "Override warning about removing the current device.",
			},
		},
		Action: func(c *cli.Context) {
			cl.ChooseCommand(&CmdDeviceRemove{Contextified: libkb.NewContextified(g)}, "remove", c)
		},
	}
}

func (c *CmdDeviceRemove) GetUsage() libkb.Usage {
	return libkb.Usage{
		Config:     true,
		GpgKeyring: true,
		KbKeyring:  true,
		API:        true,
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 The Bazel Authors. All rights reserved.
//
// 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.

// stdlib builds the standard library in the appropriate mode into a new goroot.
package main

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
)

type replicateMode int

const (
	copyMode replicateMode = iota
	hardlinkMode
	softlinkMode
)

type replicateOption func(*replicateConfig)
type replicateConfig struct {
	removeFirst bool
	fileMode    replicateMode
	dirMode     replicateMode
	paths       []string
}

func replicatePaths(paths ...string) replicateOption {
	return func(config *replicateConfig) {
		config.paths = append(config.paths, paths...)
	}
}

// replicatePrepare is the common preparation steps for a replication entry
func replicatePrepare(dst string, config *replicateConfig) error {
	dir := filepath.Dir(dst)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return fmt.Errorf("Failed to make %s: %v", dir, err)
	}
	if config.removeFirst {
		_ = os.Remove(dst)
	}
	return nil
}

// replicateFile is called internally by replicate to map a single file from src into dst.
func replicateFile(src, dst string, config *replicateConfig) error {
	if err := replicatePrepare(dst, config); err != nil {
		return err
	}
	switch config.fileMode {
	case copyMode:
		in, err := os.Open(src)
		if err != nil {
			return err
		}
		defer in.Close()
		out, err := os.Create(dst)
		if err != nil {
			return err
		}
		_, err = io.Copy(out, in)
		closeerr := out.Close()
		if err != nil {
			return err
		}
		if closeerr != nil {
			return closeerr
		}
		s, err := os.Stat(src)
		if err != nil {
			return err
		}
		if err := os.Chmod(dst, s.Mode()); err != nil {
			return err
		}
		return nil
	case hardlinkMode:
		return os.Link(src, dst)
	case softlinkMode:
		return os.Symlink(src, dst)
	default:
		return fmt.Errorf("Invalid replication mode %d", config.fileMode)
	}
}

// replicateDir makes a tree of files visible in a new location.
// It is allowed to take any efficient method of doing so.
func replicateDir(src, dst string, config *replicateConfig) error {
	if err := replicatePrepare(dst, config); err != nil {
		return err
	}
	switch config.dirMode {
	case copyMode:
		return filepath.Walk(src, func(path string, f os.FileInfo, err error) error {
			if f.IsDir() {
				return nil
			}
			relative, err := filepath.Rel(src, path)
			if err != nil {
				return err
			}
			return replicateFile(path, filepath.Join(dst, relative), config)
		})
	case hardlinkMode:
		return os.Link(src, dst)
	case softlinkMode:
		return os.Symlink(src, dst)
	default:
		return fmt.Errorf("Invalid replication mode %d", config.fileMode)
	}
}

// replicateTree is called for each single src dst pair.
func replicateTree(src, dst string, config *replicateConfig) error {
	if err := os.RemoveAll(dst); err != nil {
		return fmt.Errorf("Failed to remove file at destination %s: %v", dst, err)
	}
	if l, err := os.Readlink(src); err == nil {
		src = l
	}
	s, err := os.Stat(src)
	if err != nil && !os.IsNotExist(err) {
		return err
	}
	if s.IsDir() {
		return replicateDir(src, dst, config)
	}
	return replicateFile(src, dst, config)
}

// replicate makes a tree of files visible in a new location.
// You control how it does so using options, by default it presumes the entire tree
// of files rooted at src must be visible at dst, and that it should do so by copying.
// src is allowed to be a file, in which case just the one file is copied.
func replicate(src, dst string, options ...replicateOption) error {
	config := replicateConfig{
		removeFirst: true,
	}
	for _, option := range options {
		option(&config)
	}
	if len(config.paths) == 0 {
		return replicateTree(src, dst, &config)
	}
	for _, base := range config.paths {
		from := filepath.Join(src, base)
		to := filepath.Join(dst, base)
		if err := replicateTree(from, to, &config); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>stdlib: dereference all symlinks before copying a directory (#1348)<commit_after>// Copyright 2018 The Bazel Authors. All rights reserved.
//
// 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.

// stdlib builds the standard library in the appropriate mode into a new goroot.
package main

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
)

type replicateMode int

const (
	copyMode replicateMode = iota
	hardlinkMode
	softlinkMode
)

type replicateOption func(*replicateConfig)
type replicateConfig struct {
	removeFirst bool
	fileMode    replicateMode
	dirMode     replicateMode
	paths       []string
}

func replicatePaths(paths ...string) replicateOption {
	return func(config *replicateConfig) {
		config.paths = append(config.paths, paths...)
	}
}

// replicatePrepare is the common preparation steps for a replication entry
func replicatePrepare(dst string, config *replicateConfig) error {
	dir := filepath.Dir(dst)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return fmt.Errorf("Failed to make %s: %v", dir, err)
	}
	if config.removeFirst {
		_ = os.Remove(dst)
	}
	return nil
}

// replicateFile is called internally by replicate to map a single file from src into dst.
func replicateFile(src, dst string, config *replicateConfig) error {
	if err := replicatePrepare(dst, config); err != nil {
		return err
	}
	switch config.fileMode {
	case copyMode:
		in, err := os.Open(src)
		if err != nil {
			return err
		}
		defer in.Close()
		out, err := os.Create(dst)
		if err != nil {
			return err
		}
		_, err = io.Copy(out, in)
		closeerr := out.Close()
		if err != nil {
			return err
		}
		if closeerr != nil {
			return closeerr
		}
		s, err := os.Stat(src)
		if err != nil {
			return err
		}
		if err := os.Chmod(dst, s.Mode()); err != nil {
			return err
		}
		return nil
	case hardlinkMode:
		return os.Link(src, dst)
	case softlinkMode:
		return os.Symlink(src, dst)
	default:
		return fmt.Errorf("Invalid replication mode %d", config.fileMode)
	}
}

// replicateDir makes a tree of files visible in a new location.
// It is allowed to take any efficient method of doing so.
func replicateDir(src, dst string, config *replicateConfig) error {
	if err := replicatePrepare(dst, config); err != nil {
		return err
	}
	switch config.dirMode {
	case copyMode:
		return filepath.Walk(src, func(path string, f os.FileInfo, err error) error {
			if f.IsDir() {
				return nil
			}
			relative, err := filepath.Rel(src, path)
			if err != nil {
				return err
			}
			return replicateFile(path, filepath.Join(dst, relative), config)
		})
	case hardlinkMode:
		return os.Link(src, dst)
	case softlinkMode:
		return os.Symlink(src, dst)
	default:
		return fmt.Errorf("Invalid replication mode %d", config.fileMode)
	}
}

// replicateTree is called for each single src dst pair.
func replicateTree(src, dst string, config *replicateConfig) error {
	if err := os.RemoveAll(dst); err != nil {
		return fmt.Errorf("Failed to remove file at destination %s: %v", dst, err)
	}
	if l, err := filepath.EvalSymlinks(src); err == nil {
		src = l
	}
	s, err := os.Stat(src)
	if err != nil && !os.IsNotExist(err) {
		return err
	}
	if s.IsDir() {
		return replicateDir(src, dst, config)
	}
	return replicateFile(src, dst, config)
}

// replicate makes a tree of files visible in a new location.
// You control how it does so using options, by default it presumes the entire tree
// of files rooted at src must be visible at dst, and that it should do so by copying.
// src is allowed to be a file, in which case just the one file is copied.
func replicate(src, dst string, options ...replicateOption) error {
	config := replicateConfig{
		removeFirst: true,
	}
	for _, option := range options {
		option(&config)
	}
	if len(config.paths) == 0 {
		return replicateTree(src, dst, &config)
	}
	for _, base := range config.paths {
		from := filepath.Join(src, base)
		to := filepath.Join(dst, base)
		if err := replicateTree(from, to, &config); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"code.google.com/p/go.net/websocket"
	"log"
	"math"
)

type payload struct {
	Robots      []robot      `json:"robots"`
	Projectiles []projectile `json:"projectiles"`
	Splosions   []splosion   `json:"splosions"`
}

type player struct {
	ws          *websocket.Conn
	Robot       robot
	send        chan *payload
	Instruction instruction
}

type instruction struct {
	MoveTo position `json:"move_to"`
	FireAt position `json:"fire_at"`
}

func (p *player) sender() {
	for things := range p.send {

		err := websocket.JSON.Send(p.ws, *things)
		if err != nil {
			break
		}
	}
	p.ws.Close()
}

func (p *player) recv() {
	for {
		var msg instruction
		err := websocket.JSON.Receive(p.ws, &msg)
		if err != nil {
			log.Print(err)
			break
		}
		p.Robot.MoveTo = msg.MoveTo
		p.Robot.FireAt = msg.FireAt
	}
	p.ws.Close()
}

const threshold = 1e-7

func distance(p1, p2 position) float64 {
	deltaX := float64(p2.X - p1.X)
	deltaY := float64(p2.Y - p1.Y)
	mag := math.Abs(math.Sqrt(deltaX*deltaX + deltaY*deltaY))
	return mag
}

func move(d1, d2 position, velocity float64, timeDelta float64) position {
	deltaX := float64(d2.X - d1.X)
	deltaY := float64(d2.Y - d1.Y)
	mag := math.Sqrt(deltaX*deltaX + deltaY*deltaY)
	if math.Abs(mag) < threshold {
		return d1
	}

	distance_this_frame := velocity * timeDelta
	// log.Printf("%+v  ->  %+v   :  %+v", d1, d2, distance_this_frame)

	if math.Abs(mag) < distance_this_frame {
		return d2
	}

	moveX := 0.0
	moveY := 0.0

	if deltaY > 0 {
		theta := math.Atan(deltaX / deltaY)
		moveX = math.Sin(theta) * distance_this_frame
		moveY = math.Cos(theta) * distance_this_frame
	} else {
		theta := math.Atan(deltaX / deltaY)
		moveX = -math.Sin(theta) * distance_this_frame
		moveY = -math.Cos(theta) * distance_this_frame
	}

	// log.Printf("%+v , %+v", moveX, moveY)
	return position{
		d1.X + moveX,
		d1.Y + moveY,
	}
}

func (p *player) nudge() {
	newPos := move(p.Robot.Position, p.Robot.MoveTo, *velocity, *delta)
	p.Robot.Position.X = newPos.X
	p.Robot.Position.Y = newPos.Y
}

func (p *player) scan() {
	p.Robot.Scanners = make([]scanner, 0)
	for player := range g.players {
		if player.Robot.Id == p.Robot.Id {
			continue
		}
		dist := distance(player.Robot.Position, p.Robot.Position)
		if dist < 200.0 {
			s := scanner{
				Position: position{
					X: player.Robot.Position.X,
					Y: player.Robot.Position.Y,
				},
			}
			p.Robot.Scanners = append(p.Robot.Scanners, s)
		}
	}
}

func (s *splosion) tick() {
	s.Lifespan--
	if s.Lifespan <= 0 {
		delete(g.splosions, s)
	}
}

func (p *projectile) nudge() {
	newPos := move(p.Position, p.MoveTo, *velocity*5, *delta)

	hit_player := false
	for player := range g.players {
		if player.Robot.Id == p.Id {
			continue
		}
		dist := distance(player.Robot.Position, p.Position)
		if dist < 5.0 {
			hit_player = true
		}
	}

	if distance(p.Position, p.MoveTo) < 5 || hit_player {
		delete(g.projectiles, p)

		// Spawn a splosion
		splo := &splosion{
			Id:        p.Id,
			Position:  p.Position,
			Radius:    *weapon_radius,
			MaxDamage: 10,
			MinDamage: 5,
			Lifespan:  8,
		}
		g.splosions[splo] = true

		for player := range g.players {
			dist := distance(player.Robot.Position, p.Position)
			if dist < float64(*weapon_radius) {

				// TODO map damage Max to Min based on distance from explosion
				if player.Robot.Health > 0 {
					player.Robot.Health -= p.Damage
					// log.Printf("Robot %+v is injured", player.Robot)
					if player.Robot.Health <= 0 {
						log.Printf("Robot %+v is dead", player.Robot)
					}
				}
			}
		}
	}
	p.Position.X = newPos.X
	p.Position.Y = newPos.Y
}

func (p *player) fire() {

	for proj := range g.projectiles {
		if proj.Id == p.Robot.Id {
			return
		}
	}

	// log.Printf("%v Fired at %v %v", p.Robot.Id, p.Robot.FireAt.X, p.Robot.FireAt.Y)

	proj := &projectile{
		Id:       p.Robot.Id,
		Position: p.Robot.Position,
		MoveTo:   p.Robot.FireAt,
		Damage:   10,
	}
	g.projectiles[proj] = true
}
<commit_msg>Dont show dead robots in scans<commit_after>package main

import (
	"code.google.com/p/go.net/websocket"
	"log"
	"math"
)

type payload struct {
	Robots      []robot      `json:"robots"`
	Projectiles []projectile `json:"projectiles"`
	Splosions   []splosion   `json:"splosions"`
}

type player struct {
	ws          *websocket.Conn
	Robot       robot
	send        chan *payload
	Instruction instruction
}

type instruction struct {
	MoveTo position `json:"move_to"`
	FireAt position `json:"fire_at"`
}

func (p *player) sender() {
	for things := range p.send {

		err := websocket.JSON.Send(p.ws, *things)
		if err != nil {
			break
		}
	}
	p.ws.Close()
}

func (p *player) recv() {
	for {
		var msg instruction
		err := websocket.JSON.Receive(p.ws, &msg)
		if err != nil {
			log.Print(err)
			break
		}
		p.Robot.MoveTo = msg.MoveTo
		p.Robot.FireAt = msg.FireAt
	}
	p.ws.Close()
}

const threshold = 1e-7

func distance(p1, p2 position) float64 {
	deltaX := float64(p2.X - p1.X)
	deltaY := float64(p2.Y - p1.Y)
	mag := math.Abs(math.Sqrt(deltaX*deltaX + deltaY*deltaY))
	return mag
}

func move(d1, d2 position, velocity float64, timeDelta float64) position {
	deltaX := float64(d2.X - d1.X)
	deltaY := float64(d2.Y - d1.Y)
	mag := math.Sqrt(deltaX*deltaX + deltaY*deltaY)
	if math.Abs(mag) < threshold {
		return d1
	}

	distance_this_frame := velocity * timeDelta
	// log.Printf("%+v  ->  %+v   :  %+v", d1, d2, distance_this_frame)

	if math.Abs(mag) < distance_this_frame {
		return d2
	}

	moveX := 0.0
	moveY := 0.0

	if deltaY > 0 {
		theta := math.Atan(deltaX / deltaY)
		moveX = math.Sin(theta) * distance_this_frame
		moveY = math.Cos(theta) * distance_this_frame
	} else {
		theta := math.Atan(deltaX / deltaY)
		moveX = -math.Sin(theta) * distance_this_frame
		moveY = -math.Cos(theta) * distance_this_frame
	}

	// log.Printf("%+v , %+v", moveX, moveY)
	return position{
		d1.X + moveX,
		d1.Y + moveY,
	}
}

func (p *player) nudge() {
	newPos := move(p.Robot.Position, p.Robot.MoveTo, *velocity, *delta)
	p.Robot.Position.X = newPos.X
	p.Robot.Position.Y = newPos.Y
}

func (p *player) scan() {
	p.Robot.Scanners = make([]scanner, 0)
	for player := range g.players {
		if player.Robot.Id == p.Robot.Id || player.Robot.Health <= 0 {
			continue
		}
		dist := distance(player.Robot.Position, p.Robot.Position)
		if dist < 200.0 {
			s := scanner{
				Position: position{
					X: player.Robot.Position.X,
					Y: player.Robot.Position.Y,
				},
			}
			p.Robot.Scanners = append(p.Robot.Scanners, s)
		}
	}
}

func (s *splosion) tick() {
	s.Lifespan--
	if s.Lifespan <= 0 {
		delete(g.splosions, s)
	}
}

func (p *projectile) nudge() {
	newPos := move(p.Position, p.MoveTo, *velocity*5, *delta)

	hit_player := false
	for player := range g.players {
		if player.Robot.Id == p.Id {
			continue
		}
		dist := distance(player.Robot.Position, p.Position)
		if dist < 5.0 {
			hit_player = true
		}
	}

	if distance(p.Position, p.MoveTo) < 5 || hit_player {
		delete(g.projectiles, p)

		// Spawn a splosion
		splo := &splosion{
			Id:        p.Id,
			Position:  p.Position,
			Radius:    *weapon_radius,
			MaxDamage: 10,
			MinDamage: 5,
			Lifespan:  8,
		}
		g.splosions[splo] = true

		for player := range g.players {
			dist := distance(player.Robot.Position, p.Position)
			if dist < float64(*weapon_radius) {

				// TODO map damage Max to Min based on distance from explosion
				if player.Robot.Health > 0 {
					player.Robot.Health -= p.Damage
					// log.Printf("Robot %+v is injured", player.Robot)
					if player.Robot.Health <= 0 {
						log.Printf("Robot %+v is dead", player.Robot)
					}
				}
			}
		}
	}
	p.Position.X = newPos.X
	p.Position.Y = newPos.Y
}

func (p *player) fire() {

	for proj := range g.projectiles {
		if proj.Id == p.Robot.Id {
			return
		}
	}

	// log.Printf("%v Fired at %v %v", p.Robot.Id, p.Robot.FireAt.X, p.Robot.FireAt.Y)

	proj := &projectile{
		Id:       p.Robot.Id,
		Position: p.Robot.Position,
		MoveTo:   p.Robot.FireAt,
		Damage:   10,
	}
	g.projectiles[proj] = true
}
<|endoftext|>


<commit_before>package zeusmaster

import (
	"fmt"
	"math/rand"
	"net"
	"path/filepath"
	"strconv"
	"time"

	"github.com/burke/zeus/go/unixsocket"
)

const zeusSockName string = ".zeus.sock"

func StartClientHandler(tree *ProcessTree, quit chan bool) {
	path, _ := filepath.Abs(zeusSockName)
	addr, err := net.ResolveUnixAddr("unix", path)
	if err != nil {
		panic("Can't open socket.")
	}
	listener, err := net.ListenUnix("unix", addr)
	if err != nil {
		ErrorCantCreateListener()
	}
	defer println("CLOSING")
	defer listener.Close()

	connections := make(chan *unixsocket.Usock)
	go func() {
		for {
			if conn, err := listener.AcceptUnix(); err != nil {
				errorUnableToAcceptSocketConnection()
				time.Sleep(500 * time.Millisecond)
			} else {
				connections <- unixsocket.NewUsock(conn)
			}
		}
	}()

	for {
		select {
		case <-quit:
			quit <- true
			return
		case conn := <-connections:
			go handleClientConnection(tree, conn)
		}
	}
}

// see docs/client_master_handshake.md
func handleClientConnection(tree *ProcessTree, usock *unixsocket.Usock) {
	defer usock.Close()

	// we have established first contact to the client.

	// we first read the command and arguments specified from the connection. (step 1)
	msg, err := usock.ReadMessage()
	if err != nil {
		fmt.Println(err)
		return
	}
	command, arguments, err := ParseClientCommandRequestMessage(msg)
	if err != nil {
		fmt.Println(err)
		return
	}

	commandNode := tree.FindCommand(command)
	if commandNode == nil {
		fmt.Println("ERROR: Node not found!: ", command)
		return
	}
	command = commandNode.Name // resolve aliases
	slaveNode := commandNode.Parent

	// Now we read the terminal IO socket to use for raw IO (step 2)
	clientFd, err := usock.ReadFD()
	if err != nil {
		fmt.Println("Expected FD, none received!")
		return
	}
	fileName := strconv.Itoa(rand.Int())
	clientFile := unixsocket.FdToFile(clientFd, fileName)
	defer clientFile.Close()

	// We now need to fork a new command process.
	// For now, we naively assume it's running...

	if slaveNode.Error != "" {
		// we can skip steps 3-5 as they deal with the command process we're not spawning.
		// Write a fake pid (step 6)
		usock.WriteMessage("0")
		// Write the error message to the terminal
		clientFile.Write([]byte(slaveNode.Error))
		// Skip step 7, and write an exit code to the client (step 8)
		usock.WriteMessage("1")
		return
	}

	// boot a command process and establish a socket connection to it.
	slaveNode.WaitUntilBooted()

	msg = CreateSpawnCommandMessage(command)
	slaveNode.mu.Lock()
	slaveNode.Socket.Write([]byte(msg))
	slaveNode.mu.Unlock()

	// TODO: deadline? how to respond if this is never sent?
	commandFd := <-slaveNode.ClientCommandPTYFileDescriptor
	if err != nil {
		fmt.Println("Couldn't start command process!", err)
	}
	fileName = strconv.Itoa(rand.Int())
	commandFile := unixsocket.FdToFile(commandFd, fileName)
	defer commandFile.Close()

	// Send the arguments to the command process (step 3)
	commandFile.Write([]byte(arguments)) // TODO: What if they're too long?

	commandSocket, err := unixsocket.MakeUnixSocket(commandFile)
	if err != nil {
		fmt.Println("MakeUnixSocket", err)
	}
	defer commandSocket.Close()

	// Send the client terminal connection to the command process (step 4)
	commandUsock := unixsocket.NewUsock(commandSocket)
	commandUsock.WriteFD(clientFd)

	// Receive the pid from the command process (step 5)
	msg, err = commandUsock.ReadMessage()
	if err != nil {
		println(err)
	}
	intPid, _, _ := ParsePidMessage(msg)

	// Send the pid to the client process (step 6)
	strPid := strconv.Itoa(intPid)
	usock.WriteMessage(strPid)

	// Receive the exit status from the command (step 7)
	msg, err = commandUsock.ReadMessage()
	if err != nil {
		println(err)
	}

	// Forward the exit status to the Client (step 8)
	usock.WriteMessage(msg)

	// Done! Hooray!

}
<commit_msg>Print errors, not addresses of errors<commit_after>package zeusmaster

import (
	"fmt"
	"math/rand"
	"net"
	"path/filepath"
	"strconv"
	"time"

	"github.com/burke/zeus/go/unixsocket"
)

const zeusSockName string = ".zeus.sock"

func StartClientHandler(tree *ProcessTree, quit chan bool) {
	path, _ := filepath.Abs(zeusSockName)
	addr, err := net.ResolveUnixAddr("unix", path)
	if err != nil {
		panic("Can't open socket.")
	}
	listener, err := net.ListenUnix("unix", addr)
	if err != nil {
		ErrorCantCreateListener()
	}
	defer listener.Close()

	connections := make(chan *unixsocket.Usock)
	go func() {
		for {
			if conn, err := listener.AcceptUnix(); err != nil {
				errorUnableToAcceptSocketConnection()
				time.Sleep(500 * time.Millisecond)
			} else {
				connections <- unixsocket.NewUsock(conn)
			}
		}
	}()

	for {
		select {
		case <-quit:
			quit <- true
			return
		case conn := <-connections:
			go handleClientConnection(tree, conn)
		}
	}
}

// see docs/client_master_handshake.md
func handleClientConnection(tree *ProcessTree, usock *unixsocket.Usock) {
	defer usock.Close()

	// we have established first contact to the client.

	// we first read the command and arguments specified from the connection. (step 1)
	msg, err := usock.ReadMessage()
	if err != nil {
		fmt.Println(err)
		return
	}
	command, arguments, err := ParseClientCommandRequestMessage(msg)
	if err != nil {
		fmt.Println(err)
		return
	}

	commandNode := tree.FindCommand(command)
	if commandNode == nil {
		fmt.Println("ERROR: Node not found!: ", command)
		return
	}
	command = commandNode.Name // resolve aliases
	slaveNode := commandNode.Parent

	// Now we read the terminal IO socket to use for raw IO (step 2)
	clientFd, err := usock.ReadFD()
	if err != nil {
		fmt.Println("Expected FD, none received!")
		return
	}
	fileName := strconv.Itoa(rand.Int())
	clientFile := unixsocket.FdToFile(clientFd, fileName)
	defer clientFile.Close()

	// We now need to fork a new command process.
	// For now, we naively assume it's running...

	if slaveNode.Error != "" {
		// we can skip steps 3-5 as they deal with the command process we're not spawning.
		// Write a fake pid (step 6)
		usock.WriteMessage("0")
		// Write the error message to the terminal
		clientFile.Write([]byte(slaveNode.Error))
		// Skip step 7, and write an exit code to the client (step 8)
		usock.WriteMessage("1")
		return
	}

	// boot a command process and establish a socket connection to it.
	slaveNode.WaitUntilBooted()

	msg = CreateSpawnCommandMessage(command)
	slaveNode.mu.Lock()
	slaveNode.Socket.Write([]byte(msg))
	slaveNode.mu.Unlock()

	// TODO: deadline? how to respond if this is never sent?
	commandFd := <-slaveNode.ClientCommandPTYFileDescriptor
	if err != nil {
		fmt.Println("Couldn't start command process!", err)
	}
	fileName = strconv.Itoa(rand.Int())
	commandFile := unixsocket.FdToFile(commandFd, fileName)
	defer commandFile.Close()

	// Send the arguments to the command process (step 3)
	commandFile.Write([]byte(arguments)) // TODO: What if they're too long?

	commandSocket, err := unixsocket.MakeUnixSocket(commandFile)
	if err != nil {
		fmt.Println("MakeUnixSocket", err)
	}
	defer commandSocket.Close()

	// Send the client terminal connection to the command process (step 4)
	commandUsock := unixsocket.NewUsock(commandSocket)
	commandUsock.WriteFD(clientFd)

	// Receive the pid from the command process (step 5)
	msg, err = commandUsock.ReadMessage()
	if err != nil {
		fmt.Println(err)
	}
	intPid, _, _ := ParsePidMessage(msg)

	// Send the pid to the client process (step 6)
	strPid := strconv.Itoa(intPid)
	usock.WriteMessage(strPid)

	// Receive the exit status from the command (step 7)
	msg, err = commandUsock.ReadMessage()
	if err != nil {
		fmt.Println(err)
	}

	// Forward the exit status to the Client (step 8)
	usock.WriteMessage(msg)

	// Done! Hooray!

}
<|endoftext|>


<commit_before>package main

import (
    "bytes"
    "encoding/binary"
    "errors"
    "flag"
    "fmt"
    "github.com/cyfdecyf/leakybuf"
    ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
    "io"
    "log"
    "net"
    "os"
    "os/signal"
    "runtime"
    "strconv"
    "syscall"
    "bufio"
    "net/http"
)

var debug ss.DebugLog

var storage *Storage

const dnsGoroutineNum = 64

func getUser(conn net.Conn) (user User, err error) {
    const (
        idVersion = 0 // \x01
        version   = 1
        idUser    = 1 // \x02
    )
    buf := make([]byte, 2)
    if _, err = io.ReadFull(conn, buf); err != nil {
        return
    }
    switch buf[idVersion] {
    case version:
        break
    default:
        err = errors.New(fmt.Sprintf("version %s not supported", buf[idVersion]))
        return
    }
    userLen := buf[idUser]
    username := make([]byte, userLen)
    if _, err = io.ReadFull(conn, username); err != nil {
        return
    }
    user, err = storage.Get(string(username))
    return
}

func getRequest(conn *ss.Conn) (host string, extra []byte, err error) {
    const (
        idType  = 0 // address type index
        idIP0   = 1 // ip addres start index
        idDmLen = 1 // domain address length index
        idDm0   = 2 // domain address start index

        typeIPv4 = 1 // type is ipv4 address
        typeDm   = 3 // type is domain address
        typeIPv6 = 4 // type is ipv6 address

        lenIPv4   = 1 + net.IPv4len + 2 // 1addrType + ipv4 + 2port
        lenIPv6   = 1 + net.IPv6len + 2 // 1addrType + ipv6 + 2port
        lenDmBase = 1 + 1 + 2           // 1addrType + 1addrLen + 2port, plus addrLen
    )

    // buf size should at least have the same size with the largest possible
    // request size (when addrType is 3, domain name has at most 256 bytes)
    // 1(addrType) + 1(lenByte) + 256(max length address) + 2(port)
    buf := make([]byte, 260)
    var n int
    // read till we get possible domain length field
    ss.SetReadTimeout(conn)
    if n, err = io.ReadAtLeast(conn, buf, idDmLen+1); err != nil {
        return
    }

    reqLen := -1
    switch buf[idType] {
    case typeIPv4:
        reqLen = lenIPv4
    case typeIPv6:
        reqLen = lenIPv6
    case typeDm:
        reqLen = int(buf[idDmLen]) + lenDmBase
    default:
        err = errors.New(fmt.Sprintf("addr type %d not supported", buf[idType]))
        return
    }

    if n < reqLen { // rare case
        ss.SetReadTimeout(conn)
        if _, err = io.ReadFull(conn, buf[n:reqLen]); err != nil {
            return
        }
    } else if n > reqLen {
        // it's possible to read more than just the request head
        extra = buf[reqLen:n]
    }

    // Return string for typeIP is not most efficient, but browsers (Chrome,
    // Safari, Firefox) all seems using typeDm exclusively. So this is not a
    // big problem.
    switch buf[idType] {
    case typeIPv4:
        host = net.IP(buf[idIP0 : idIP0+net.IPv4len]).String()
    case typeIPv6:
        host = net.IP(buf[idIP0 : idIP0+net.IPv6len]).String()
    case typeDm:
        host = string(buf[idDm0 : idDm0+buf[idDmLen]])
    }
    // parse port
    port := binary.BigEndian.Uint16(buf[reqLen-2 : reqLen])
    host = net.JoinHostPort(host, strconv.Itoa(int(port)))
    return
}

const logCntDelta = 100

var connCnt int
var nextLogConnCnt int = logCntDelta

func handleConnection(user User, conn *ss.Conn) {
    var host string
    var size = 0
    var raw_req_header, raw_res_header []byte
    var is_http = false
    var res_size = 0
    var req_chan = make(chan []byte)

    connCnt++ // this maybe not accurate, but should be enough
    if connCnt-nextLogConnCnt >= 0 {
        // XXX There's no xadd in the atomic package, so it's difficult to log
        // the message only once with low cost. Also note nextLogConnCnt maybe
        // added twice for current peak connection number level.
        log.Printf("Number of client connections reaches %d\n", nextLogConnCnt)
        nextLogConnCnt += logCntDelta
    }

    // function arguments are always evaluated, so surround debug statement
    // with if statement
    if debug {
        debug.Printf("new client %s->%s\n", conn.RemoteAddr().String(), conn.LocalAddr())
    }
    closed := false
    defer func() {
        if debug {
            debug.Printf("closed pipe %s<->%s\n", conn.RemoteAddr(), host)
        }
        connCnt--
        if !closed {
            conn.Close()
        }
    }()

    host, extra, err := getRequest(conn)
    if err != nil {
        log.Println("error getting request", conn.RemoteAddr(), conn.LocalAddr(), err)
        return
    }
    debug.Println("connecting", host)
    remote, err := net.Dial("tcp", host)
    if err != nil {
        if ne, ok := err.(*net.OpError); ok && (ne.Err == syscall.EMFILE || ne.Err == syscall.ENFILE) {
            // log too many open file error
            // EMFILE is process reaches open file limits, ENFILE is system limit
            log.Println("dial error:", err)
        } else {
            log.Println("error connecting to:", host, err)
        }
        return
    }
    defer func() {
        if is_http{
            tmp_req_header := <-req_chan
            buffer := bytes.NewBuffer(raw_req_header)
            buffer.Write(tmp_req_header)
            raw_req_header = buffer.Bytes()
        }
        showConn(raw_req_header, raw_res_header, host, user, size, is_http)
        close(req_chan)
        if !closed {
            remote.Close()
        }
    }()
    // write extra bytes read from

    is_http, extra, _ = checkHttp(extra, conn)
    raw_req_header = extra
    res_size, err = remote.Write(extra)
    storage.IncrSize("flow:" + user.Name, res_size)
    size += res_size
    if err != nil {
        debug.Println("write request extra error:", err)
        return
    }

    if debug {
        debug.Printf("piping %s<->%s", conn.RemoteAddr(), host)
    }

    go func() {
        _, raw_header := PipeThenClose(conn, remote, ss.SET_TIMEOUT, is_http, false, user)
        if is_http {
            req_chan<-raw_header
        }
    }()

    res_size, raw_res_header = PipeThenClose(remote, conn, ss.NO_TIMEOUT, is_http, true, user)
    size += res_size
    closed = true
    return
}

func showConn(raw_req_header, raw_res_header []byte, host string, user User, size int, is_http bool) {
    if is_http {
        req, _ := http.ReadRequest(bufio.NewReader(bytes.NewReader(raw_req_header)))
        if req == nil {
            lines := bytes.SplitN(raw_req_header, []byte(" "), 2)
            fmt.Printf("%s http://%s/ \"Unknow\" HTTP/1.1 unknow %s %d\n", lines[0], host, user.Name, size)
            return
        }
        res, _ := http.ReadResponse(bufio.NewReader(bytes.NewReader(raw_res_header)), req)
        statusCode := 200
        if res != nil {
            statusCode = res.StatusCode
        }
        fmt.Printf("%s http://%s%s \"%s\" %s %d %s %d\n", req.Method, req.Host, req.URL.String(), req.Header.Get("user-agent"), req.Proto, statusCode, user.Name, size)
    } else {
        fmt.Printf("CONNECT %s \"Unknow\" HTTPS unknow %s %d\n", host, user.Name, size)
    }
}

func checkHttp(extra []byte, conn *ss.Conn) (is_http bool, data []byte, err error) {
    var buf []byte
    var methods = []string{"GET", "HEAD", "POST", "PUT", "TRACE", "OPTIONS", "DELETE"}
    is_http = false
    if extra == nil || len(extra) < 10 {
        buf = make([]byte, 10)
        if _, err = io.ReadFull(conn, buf); err != nil {
            return
        }
    }

    if buf == nil {
        data = extra
    } else if extra == nil {
        data = buf
    }else {
        buffer := bytes.NewBuffer(extra)
        buffer.Write(buf)
        data = buffer.Bytes()
    }

    for _, method := range methods {
        if bytes.HasPrefix(data, []byte(method)) {
            is_http = true
            break
        }
    }
    return
}

const bufSize = 4096
const nBuf = 2048

func PipeThenClose(src, dst net.Conn, timeoutOpt int, is_http bool, is_res bool, user User) (total int, raw_header []byte) {
    var pipeBuf = leakybuf.NewLeakyBuf(nBuf, bufSize)
    defer dst.Close()
    buf := pipeBuf.Get()
    // defer pipeBuf.Put(buf)
    var buffer = bytes.NewBuffer(nil)
    var is_end = false
    var size int

    for {
        if timeoutOpt == ss.SET_TIMEOUT {
            ss.SetReadTimeout(src)
        }
        n, err := src.Read(buf)
        // read may return EOF with n > 0
        // should always process n > 0 bytes before handling error
        if n > 0 {
            if is_http && !is_end {
                buffer.Write(buf)
                raw_header = buffer.Bytes()
                lines := bytes.SplitN(raw_header, []byte("\r\n\r\n"), 2)
                if len(lines) == 2 {
                    is_end = true
                }
            }

            size, err = dst.Write(buf[0:n])
            if is_res {
                total_size, _ := storage.IncrSize("flow:" + user.Name, size)
                if total_size > user.Limit {
                    return
                }
            }
            total += size
            if err != nil {
                ss.Debug.Println("write:", err)
                break
            }
        }
        if err != nil || n == 0 {
            // Always "use of closed network connection", but no easy way to
            // identify this specific error. So just leave the error along for now.
            // More info here: https://code.google.com/p/go/issues/detail?id=4373
            break
        }
    }
    return
}

func waitSignal() {
    var sigChan = make(chan os.Signal, 1)
    signal.Notify(sigChan, syscall.SIGHUP)
    for sig := range sigChan {
        if sig == syscall.SIGHUP {
        } else {
            // is this going to happen?
            log.Printf("caught signal %v, exit", sig)
            os.Exit(0)
        }
    }
}

func run(port string) {
    ln, err := net.Listen("tcp", ":"+port)
    if err != nil {
        log.Printf("error listening port %v: %v\n", port, err)
        return
    }
    var cipher *ss.Cipher
    log.Printf("server listening port %v ...\n", port)
    for {
        conn, err := ln.Accept()
        if err != nil {
            debug.Printf("accept error: %v\n", err)
            return
        }

        user, err := getUser(conn)
        if err != nil {
            log.Printf("Error get passeoed: %s %v\n", port, err)
            conn.Close()
            continue
        }
        size, err := storage.GetSize("flow:" + user.Name)
        if user.Limit < size {
            log.Printf("Error user runover: %s\n", size)
            conn.Close()
            continue
        }

        // log.Println("creating cipher for user:", user.Name)
        cipher, err = ss.NewCipher(user.Method, user.Password)
        if err != nil {
            log.Printf("Error generating cipher for user: %s %v\n", user.Name, err)
            conn.Close()
            continue
        }
        go handleConnection(user, ss.NewConn(conn, cipher.Copy()))
    }
}

func main() {
    log.SetOutput(os.Stdout)

    var printVer bool
    var core int
    var serverPort string
    var redisServer string

    flag.BoolVar(&printVer, "version", false, "print version")
    flag.StringVar(&serverPort, "p", "8388", "server port")
    flag.StringVar(&redisServer, "redis", ":6379", "redis server")
    flag.IntVar(&core, "core", 0, "maximum number of CPU cores to use, default is determinied by Go runtime")
    flag.BoolVar((*bool)(&debug), "d", false, "print debug message")

    flag.Parse()

    if printVer {
        ss.PrintVersion()
        os.Exit(0)
    }

    ss.SetDebug(debug)

    storage = NewStorage(redisServer)
    if core > 0 {
        runtime.GOMAXPROCS(core)
    }

    go run(serverPort)

    waitSignal()
}
<commit_msg>use log rather than fmt<commit_after>package main

import (
    "bytes"
    "encoding/binary"
    "errors"
    "flag"
    "fmt"
    "github.com/cyfdecyf/leakybuf"
    ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
    "io"
    "log"
    "net"
    "os"
    "os/signal"
    "runtime"
    "strconv"
    "syscall"
    "bufio"
    "net/http"
)

var debug ss.DebugLog

var storage *Storage

const dnsGoroutineNum = 64

func getUser(conn net.Conn) (user User, err error) {
    const (
        idVersion = 0 // \x01
        version   = 1
        idUser    = 1 // \x02
    )
    buf := make([]byte, 2)
    if _, err = io.ReadFull(conn, buf); err != nil {
        return
    }
    switch buf[idVersion] {
    case version:
        break
    default:
        err = errors.New(fmt.Sprintf("version %s not supported", buf[idVersion]))
        return
    }
    userLen := buf[idUser]
    username := make([]byte, userLen)
    if _, err = io.ReadFull(conn, username); err != nil {
        return
    }
    user, err = storage.Get(string(username))
    return
}

func getRequest(conn *ss.Conn) (host string, extra []byte, err error) {
    const (
        idType  = 0 // address type index
        idIP0   = 1 // ip addres start index
        idDmLen = 1 // domain address length index
        idDm0   = 2 // domain address start index

        typeIPv4 = 1 // type is ipv4 address
        typeDm   = 3 // type is domain address
        typeIPv6 = 4 // type is ipv6 address

        lenIPv4   = 1 + net.IPv4len + 2 // 1addrType + ipv4 + 2port
        lenIPv6   = 1 + net.IPv6len + 2 // 1addrType + ipv6 + 2port
        lenDmBase = 1 + 1 + 2           // 1addrType + 1addrLen + 2port, plus addrLen
    )

    // buf size should at least have the same size with the largest possible
    // request size (when addrType is 3, domain name has at most 256 bytes)
    // 1(addrType) + 1(lenByte) + 256(max length address) + 2(port)
    buf := make([]byte, 260)
    var n int
    // read till we get possible domain length field
    ss.SetReadTimeout(conn)
    if n, err = io.ReadAtLeast(conn, buf, idDmLen+1); err != nil {
        return
    }

    reqLen := -1
    switch buf[idType] {
    case typeIPv4:
        reqLen = lenIPv4
    case typeIPv6:
        reqLen = lenIPv6
    case typeDm:
        reqLen = int(buf[idDmLen]) + lenDmBase
    default:
        err = errors.New(fmt.Sprintf("addr type %d not supported", buf[idType]))
        return
    }

    if n < reqLen { // rare case
        ss.SetReadTimeout(conn)
        if _, err = io.ReadFull(conn, buf[n:reqLen]); err != nil {
            return
        }
    } else if n > reqLen {
        // it's possible to read more than just the request head
        extra = buf[reqLen:n]
    }

    // Return string for typeIP is not most efficient, but browsers (Chrome,
    // Safari, Firefox) all seems using typeDm exclusively. So this is not a
    // big problem.
    switch buf[idType] {
    case typeIPv4:
        host = net.IP(buf[idIP0 : idIP0+net.IPv4len]).String()
    case typeIPv6:
        host = net.IP(buf[idIP0 : idIP0+net.IPv6len]).String()
    case typeDm:
        host = string(buf[idDm0 : idDm0+buf[idDmLen]])
    }
    // parse port
    port := binary.BigEndian.Uint16(buf[reqLen-2 : reqLen])
    host = net.JoinHostPort(host, strconv.Itoa(int(port)))
    return
}

const logCntDelta = 100

var connCnt int
var nextLogConnCnt int = logCntDelta

func handleConnection(user User, conn *ss.Conn) {
    var host string
    var size = 0
    var raw_req_header, raw_res_header []byte
    var is_http = false
    var res_size = 0
    var req_chan = make(chan []byte)

    connCnt++ // this maybe not accurate, but should be enough
    if connCnt-nextLogConnCnt >= 0 {
        // XXX There's no xadd in the atomic package, so it's difficult to log
        // the message only once with low cost. Also note nextLogConnCnt maybe
        // added twice for current peak connection number level.
        log.Printf("Number of client connections reaches %d\n", nextLogConnCnt)
        nextLogConnCnt += logCntDelta
    }

    // function arguments are always evaluated, so surround debug statement
    // with if statement
    if debug {
        debug.Printf("new client %s->%s\n", conn.RemoteAddr().String(), conn.LocalAddr())
    }
    closed := false
    defer func() {
        if debug {
            debug.Printf("closed pipe %s<->%s\n", conn.RemoteAddr(), host)
        }
        connCnt--
        if !closed {
            conn.Close()
        }
    }()

    host, extra, err := getRequest(conn)
    if err != nil {
        log.Println("error getting request", conn.RemoteAddr(), conn.LocalAddr(), err)
        return
    }
    debug.Println("connecting", host)
    remote, err := net.Dial("tcp", host)
    if err != nil {
        if ne, ok := err.(*net.OpError); ok && (ne.Err == syscall.EMFILE || ne.Err == syscall.ENFILE) {
            // log too many open file error
            // EMFILE is process reaches open file limits, ENFILE is system limit
            log.Println("dial error:", err)
        } else {
            log.Println("error connecting to:", host, err)
        }
        return
    }
    defer func() {
        if is_http{
            tmp_req_header := <-req_chan
            buffer := bytes.NewBuffer(raw_req_header)
            buffer.Write(tmp_req_header)
            raw_req_header = buffer.Bytes()
        }
        showConn(raw_req_header, raw_res_header, host, user, size, is_http)
        close(req_chan)
        if !closed {
            remote.Close()
        }
    }()
    // write extra bytes read from

    is_http, extra, _ = checkHttp(extra, conn)
    raw_req_header = extra
    res_size, err = remote.Write(extra)
    storage.IncrSize("flow:" + user.Name, res_size)
    size += res_size
    if err != nil {
        debug.Println("write request extra error:", err)
        return
    }

    if debug {
        debug.Printf("piping %s<->%s", conn.RemoteAddr(), host)
    }

    go func() {
        _, raw_header := PipeThenClose(conn, remote, ss.SET_TIMEOUT, is_http, false, user)
        if is_http {
            req_chan<-raw_header
        }
    }()

    res_size, raw_res_header = PipeThenClose(remote, conn, ss.NO_TIMEOUT, is_http, true, user)
    size += res_size
    closed = true
    return
}

func showConn(raw_req_header, raw_res_header []byte, host string, user User, size int, is_http bool) {
    if is_http {
        req, _ := http.ReadRequest(bufio.NewReader(bytes.NewReader(raw_req_header)))
        if req == nil {
            lines := bytes.SplitN(raw_req_header, []byte(" "), 2)
            log.Printf("%s http://%s/ \"Unknow\" HTTP/1.1 unknow %s %d\n", lines[0], host, user.Name, size)
            return
        }
        res, _ := http.ReadResponse(bufio.NewReader(bytes.NewReader(raw_res_header)), req)
        statusCode := 200
        if res != nil {
            statusCode = res.StatusCode
        }
        log.Printf("%s http://%s%s \"%s\" %s %d %s %d\n", req.Method, req.Host, req.URL.String(), req.Header.Get("user-agent"), req.Proto, statusCode, user.Name, size)
    } else {
        log.Printf("CONNECT %s \"Unknow\" HTTPS unknow %s %d\n", host, user.Name, size)
    }
}

func checkHttp(extra []byte, conn *ss.Conn) (is_http bool, data []byte, err error) {
    var buf []byte
    var methods = []string{"GET", "HEAD", "POST", "PUT", "TRACE", "OPTIONS", "DELETE"}
    is_http = false
    if extra == nil || len(extra) < 10 {
        buf = make([]byte, 10)
        if _, err = io.ReadFull(conn, buf); err != nil {
            return
        }
    }

    if buf == nil {
        data = extra
    } else if extra == nil {
        data = buf
    }else {
        buffer := bytes.NewBuffer(extra)
        buffer.Write(buf)
        data = buffer.Bytes()
    }

    for _, method := range methods {
        if bytes.HasPrefix(data, []byte(method)) {
            is_http = true
            break
        }
    }
    return
}

const bufSize = 4096
const nBuf = 2048

func PipeThenClose(src, dst net.Conn, timeoutOpt int, is_http bool, is_res bool, user User) (total int, raw_header []byte) {
    var pipeBuf = leakybuf.NewLeakyBuf(nBuf, bufSize)
    defer dst.Close()
    buf := pipeBuf.Get()
    // defer pipeBuf.Put(buf)
    var buffer = bytes.NewBuffer(nil)
    var is_end = false
    var size int

    for {
        if timeoutOpt == ss.SET_TIMEOUT {
            ss.SetReadTimeout(src)
        }
        n, err := src.Read(buf)
        // read may return EOF with n > 0
        // should always process n > 0 bytes before handling error
        if n > 0 {
            if is_http && !is_end {
                buffer.Write(buf)
                raw_header = buffer.Bytes()
                lines := bytes.SplitN(raw_header, []byte("\r\n\r\n"), 2)
                if len(lines) == 2 {
                    is_end = true
                }
            }

            size, err = dst.Write(buf[0:n])
            if is_res {
                total_size, _ := storage.IncrSize("flow:" + user.Name, size)
                if total_size > user.Limit {
                    return
                }
            }
            total += size
            if err != nil {
                ss.Debug.Println("write:", err)
                break
            }
        }
        if err != nil || n == 0 {
            // Always "use of closed network connection", but no easy way to
            // identify this specific error. So just leave the error along for now.
            // More info here: https://code.google.com/p/go/issues/detail?id=4373
            break
        }
    }
    return
}

func waitSignal() {
    var sigChan = make(chan os.Signal, 1)
    signal.Notify(sigChan, syscall.SIGHUP)
    for sig := range sigChan {
        if sig == syscall.SIGHUP {
        } else {
            // is this going to happen?
            log.Printf("caught signal %v, exit", sig)
            os.Exit(0)
        }
    }
}

func run(port string) {
    ln, err := net.Listen("tcp", ":"+port)
    if err != nil {
        log.Printf("error listening port %v: %v\n", port, err)
        return
    }
    var cipher *ss.Cipher
    log.Printf("server listening port %v ...\n", port)
    for {
        conn, err := ln.Accept()
        if err != nil {
            debug.Printf("accept error: %v\n", err)
            return
        }

        user, err := getUser(conn)
        if err != nil {
            log.Printf("Error get passeoed: %s %v\n", port, err)
            conn.Close()
            continue
        }
        size, err := storage.GetSize("flow:" + user.Name)
        if user.Limit < size {
            log.Printf("Error user runover: %s\n", size)
            conn.Close()
            continue
        }

        // log.Println("creating cipher for user:", user.Name)
        cipher, err = ss.NewCipher(user.Method, user.Password)
        if err != nil {
            log.Printf("Error generating cipher for user: %s %v\n", user.Name, err)
            conn.Close()
            continue
        }
        go handleConnection(user, ss.NewConn(conn, cipher.Copy()))
    }
}

func main() {
    log.SetOutput(os.Stdout)

    var printVer bool
    var core int
    var serverPort string
    var redisServer string

    flag.BoolVar(&printVer, "version", false, "print version")
    flag.StringVar(&serverPort, "p", "8388", "server port")
    flag.StringVar(&redisServer, "redis", ":6379", "redis server")
    flag.IntVar(&core, "core", 0, "maximum number of CPU cores to use, default is determinied by Go runtime")
    flag.BoolVar((*bool)(&debug), "d", false, "print debug message")

    flag.Parse()

    if printVer {
        ss.PrintVersion()
        os.Exit(0)
    }

    ss.SetDebug(debug)

    storage = NewStorage(redisServer)
    if core > 0 {
        runtime.GOMAXPROCS(core)
    }

    go run(serverPort)

    waitSignal()
}
<|endoftext|>