<commit_before>package golang

import (
	"encoding/json"
	"fmt"
	"path/filepath"

	"strings"

	"github.com/golang/gddo/gosrc"
	"sourcegraph.com/sourcegraph"
	"sourcegraph.com/sourcegraph/srcgraph/config"
	"sourcegraph.com/sourcegraph/srcgraph/container"
	"sourcegraph.com/sourcegraph/srcgraph/dep2"
	"sourcegraph.com/sourcegraph/srcgraph/task2"
	"sourcegraph.com/sourcegraph/srcgraph/unit"
)

func init() {
	dep2.RegisterLister(&Package{}, dep2.DockerLister{defaultGoVersion})
	dep2.RegisterResolver(goImportPathTargetType, defaultGoVersion)
}

func (v *goVersion) BuildLister(dir string, unit unit.SourceUnit, c *config.Repository, x *task2.Context) (*container.Command, error) {
	goConfig := v.goConfig(c)
	pkg := unit.(*Package)

	dockerfile, err := v.baseDockerfile()
	if err != nil {
		return nil, err
	}

	containerDir := filepath.Join(containerGOPATH, "src", goConfig.BaseImportPath)
	cmd := container.Command{
		Container: container.Container{
			Dockerfile: dockerfile,
			RunOptions: []string{"-v", dir + ":" + containerDir},
			// TODO(sqs): include TestImports and XTestImports
			Cmd: []string{"go", "list", "-e", "-f", `[{{if .Imports}}"{{join .Imports "\",\""}}"{{end}}]`, pkg.ImportPath},
		},
		Transform: func(orig []byte) ([]byte, error) {
			var importPaths []string
			err := json.Unmarshal(orig, &importPaths)
			if err != nil {
				return nil, err
			}

			deps := make([]*dep2.RawDependency, len(importPaths))
			for i, importPath := range importPaths {
				deps[i] = &dep2.RawDependency{
					TargetType: goImportPathTargetType,
					Target:     goImportPath(importPath),
				}
			}

			return json.Marshal(deps)
		},
	}
	return &cmd, nil
}

// goImportPath represents a Go import path, such as "github.com/user/repo" or
// "net/http".
type goImportPath string

const goImportPathTargetType = "go-import-path"

func (v *goVersion) Resolve(dep *dep2.RawDependency, c *config.Repository, x *task2.Context) (*dep2.ResolvedTarget, error) {
	importPath := dep.Target.(string)
	return v.resolveGoImportDep(importPath, c, x)
}

func (v *goVersion) resolveGoImportDep(importPath string, c *config.Repository, x *task2.Context) (*dep2.ResolvedTarget, error) {
	// Look up in cache.
	resolvedTarget := func() *dep2.ResolvedTarget {
		v.resolveCacheMu.Lock()
		defer v.resolveCacheMu.Unlock()
		return v.resolveCache[importPath]
	}()
	if resolvedTarget != nil {
		return resolvedTarget, nil
	}

	// Check if this importPath is in this repository.
	goConfig := v.goConfig(c)
	if strings.HasPrefix(importPath, goConfig.BaseImportPath) {
		dir, err := filepath.Rel(goConfig.BaseImportPath, importPath)
		if err != nil {
			return nil, err
		}
		toUnit := &Package{Dir: dir, ImportPath: importPath}
		return &dep2.ResolvedTarget{
			// TODO(sqs): this is a URI not a clone URL
			ToRepoCloneURL: string(c.URI),
			ToUnit:         toUnit.Name(),
			ToUnitType:     unit.Type(toUnit),
		}, nil
	}

	// Special-case the cgo package "C".
	if importPath == "C" {
		return nil, nil
	}

	if gosrc.IsGoRepoPath(importPath) {
		toUnit := &Package{ImportPath: importPath, Dir: "src/pkg/" + importPath}
		return &dep2.ResolvedTarget{
			ToRepoCloneURL:  v.RepositoryCloneURL,
			ToVersionString: v.VersionString,
			ToRevSpec:       v.VCSRevision,
			ToUnit:          toUnit.Name(),
			ToUnitType:      unit.Type(toUnit),
		}, nil
	}

	x.Log.Printf("Resolving Go dep: %s", importPath)

	dir, err := gosrc.Get(sourcegraph.AuthenticatingAsNeededHTTPClient, string(importPath), "")
	if err != nil {
		return nil, fmt.Errorf("unable to fetch information about Go package %q", importPath)
	}

	// gosrc returns code.google.com URLs ending in a slash. Remove it.
	dir.ProjectURL = strings.TrimSuffix(dir.ProjectURL, "/")

	toUnit := &Package{ImportPath: dir.ImportPath}
	toUnit.Dir, err = filepath.Rel(dir.ProjectRoot, dir.ImportPath)
	if err != nil {
		return nil, err
	}

	resolvedTarget = &dep2.ResolvedTarget{
		ToRepoCloneURL: dir.ProjectURL,
		ToUnit:         toUnit.Name(),
		ToUnitType:     unit.Type(toUnit),
	}

	if gosrc.IsGoRepoPath(dir.ImportPath) {
		resolvedTarget.ToVersionString = v.VersionString
		resolvedTarget.ToRevSpec = v.VCSRevision
		resolvedTarget.ToUnit = "src/pkg/" + resolvedTarget.ToUnit
	}

	// Save in cache.
	v.resolveCacheMu.Lock()
	defer v.resolveCacheMu.Unlock()
	if v.resolveCache == nil {
		v.resolveCache = make(map[string]*dep2.ResolvedTarget)
	}
	v.resolveCache[importPath] = resolvedTarget

	return resolvedTarget, nil
}
<commit_msg>use client that doesnt require oauth creds...might hit rate limit, we will see<commit_after>package golang

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"path/filepath"

	"strings"

	"github.com/golang/gddo/gosrc"
	"github.com/peterbourgon/diskv"
	"github.com/sourcegraph/httpcache"
	"github.com/sourcegraph/httpcache/diskcache"
	"sourcegraph.com/sourcegraph/srcgraph/config"
	"sourcegraph.com/sourcegraph/srcgraph/container"
	"sourcegraph.com/sourcegraph/srcgraph/dep2"
	"sourcegraph.com/sourcegraph/srcgraph/task2"
	"sourcegraph.com/sourcegraph/srcgraph/unit"
)

func init() {
	dep2.RegisterLister(&Package{}, dep2.DockerLister{defaultGoVersion})
	dep2.RegisterResolver(goImportPathTargetType, defaultGoVersion)
}

func (v *goVersion) BuildLister(dir string, unit unit.SourceUnit, c *config.Repository, x *task2.Context) (*container.Command, error) {
	goConfig := v.goConfig(c)
	pkg := unit.(*Package)

	dockerfile, err := v.baseDockerfile()
	if err != nil {
		return nil, err
	}

	containerDir := filepath.Join(containerGOPATH, "src", goConfig.BaseImportPath)
	cmd := container.Command{
		Container: container.Container{
			Dockerfile: dockerfile,
			RunOptions: []string{"-v", dir + ":" + containerDir},
			// TODO(sqs): include TestImports and XTestImports
			Cmd: []string{"go", "list", "-e", "-f", `[{{if .Imports}}"{{join .Imports "\",\""}}"{{end}}]`, pkg.ImportPath},
		},
		Transform: func(orig []byte) ([]byte, error) {
			var importPaths []string
			err := json.Unmarshal(orig, &importPaths)
			if err != nil {
				return nil, err
			}

			deps := make([]*dep2.RawDependency, len(importPaths))
			for i, importPath := range importPaths {
				deps[i] = &dep2.RawDependency{
					TargetType: goImportPathTargetType,
					Target:     goImportPath(importPath),
				}
			}

			return json.Marshal(deps)
		},
	}
	return &cmd, nil
}

// goImportPath represents a Go import path, such as "github.com/user/repo" or
// "net/http".
type goImportPath string

const goImportPathTargetType = "go-import-path"

func (v *goVersion) Resolve(dep *dep2.RawDependency, c *config.Repository, x *task2.Context) (*dep2.ResolvedTarget, error) {
	importPath := dep.Target.(string)
	return v.resolveGoImportDep(importPath, c, x)
}

func (v *goVersion) resolveGoImportDep(importPath string, c *config.Repository, x *task2.Context) (*dep2.ResolvedTarget, error) {
	// Look up in cache.
	resolvedTarget := func() *dep2.ResolvedTarget {
		v.resolveCacheMu.Lock()
		defer v.resolveCacheMu.Unlock()
		return v.resolveCache[importPath]
	}()
	if resolvedTarget != nil {
		return resolvedTarget, nil
	}

	// Check if this importPath is in this repository.
	goConfig := v.goConfig(c)
	if strings.HasPrefix(importPath, goConfig.BaseImportPath) {
		dir, err := filepath.Rel(goConfig.BaseImportPath, importPath)
		if err != nil {
			return nil, err
		}
		toUnit := &Package{Dir: dir, ImportPath: importPath}
		return &dep2.ResolvedTarget{
			// TODO(sqs): this is a URI not a clone URL
			ToRepoCloneURL: string(c.URI),
			ToUnit:         toUnit.Name(),
			ToUnitType:     unit.Type(toUnit),
		}, nil
	}

	// Special-case the cgo package "C".
	if importPath == "C" {
		return nil, nil
	}

	if gosrc.IsGoRepoPath(importPath) {
		toUnit := &Package{ImportPath: importPath, Dir: "src/pkg/" + importPath}
		return &dep2.ResolvedTarget{
			ToRepoCloneURL:  v.RepositoryCloneURL,
			ToVersionString: v.VersionString,
			ToRevSpec:       v.VCSRevision,
			ToUnit:          toUnit.Name(),
			ToUnitType:      unit.Type(toUnit),
		}, nil
	}

	x.Log.Printf("Resolving Go dep: %s", importPath)

	dir, err := gosrc.Get(cachingHTTPClient, string(importPath), "")
	if err != nil {
		return nil, fmt.Errorf("unable to fetch information about Go package %q", importPath)
	}

	// gosrc returns code.google.com URLs ending in a slash. Remove it.
	dir.ProjectURL = strings.TrimSuffix(dir.ProjectURL, "/")

	toUnit := &Package{ImportPath: dir.ImportPath}
	toUnit.Dir, err = filepath.Rel(dir.ProjectRoot, dir.ImportPath)
	if err != nil {
		return nil, err
	}

	resolvedTarget = &dep2.ResolvedTarget{
		ToRepoCloneURL: dir.ProjectURL,
		ToUnit:         toUnit.Name(),
		ToUnitType:     unit.Type(toUnit),
	}

	if gosrc.IsGoRepoPath(dir.ImportPath) {
		resolvedTarget.ToVersionString = v.VersionString
		resolvedTarget.ToRevSpec = v.VCSRevision
		resolvedTarget.ToUnit = "src/pkg/" + resolvedTarget.ToUnit
	}

	// Save in cache.
	v.resolveCacheMu.Lock()
	defer v.resolveCacheMu.Unlock()
	if v.resolveCache == nil {
		v.resolveCache = make(map[string]*dep2.ResolvedTarget)
	}
	v.resolveCache[importPath] = resolvedTarget

	return resolvedTarget, nil
}

var cachingHTTPClient = &http.Client{
	Transport: &httpcache.Transport{
		Cache: diskcache.NewWithDiskv(diskv.New(diskv.Options{
			BasePath:     filepath.Join(os.TempDir(), "sg-golang-toolchain-cache"),
			CacheSizeMax: 5000 * 1024 * 100, // 500 MB
		})),
	},
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"sync"
)

var (
	addPrefix bool
	version   string
)

type strslice []string

func (s *strslice) String() string {
	return fmt.Sprintf("%v", *s)
}

func (s *strslice) Set(v string) error {
	*s = append(*s, v)
	return nil
}

func main() {
	var targets strslice
	var showVersion bool
	flag.Var(&targets, "t", "target hostname")
	flag.BoolVar(&addPrefix, "p", false, "add hostname to line prefix")
	flag.BoolVar(&showVersion, "v", false, "show version")
	flag.Parse()
	if showVersion {
		fmt.Println("version:", version)
		return
	}

	command := flag.Args()
	if len(command) < 1 {
		flag.PrintDefaults()
		return
	}
	wg := &sync.WaitGroup{}
	for _, host := range targets {
		wg.Add(1)
		go func(h string) {
			remoteCommand(h, command)
			wg.Done()
		}(host)
	}
	wg.Wait()
}

func remoteCommand(host string, command []string) {
	args := []string{"-t", "-t"} // man ssh(1): Multiple -t options force tty allocation, even if ssh has no local tty.
	args = append(args, host)
	args = append(args, command...)

	cmd := exec.Command("ssh", args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		fmt.Fprintln(os.Stderr, "[error]", host, err)
		return
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		fmt.Fprintln(os.Stderr, "[error]", host, err)
		return
	}
	cmd.Start()
	go scanLines(stderr, os.Stderr, host)
	go scanLines(stdout, os.Stdout, host)
	cmd.Wait()
}

func scanLines(src io.ReadCloser, dest io.Writer, prefix string) {
	scanner := bufio.NewScanner(src)
	for scanner.Scan() {
		if addPrefix {
			fmt.Fprintln(dest, prefix, scanner.Text())
		} else {
			fmt.Fprintln(dest, scanner.Text())
		}
	}
}
<commit_msg>handle stdin stream<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"sync"
)

var (
	addPrefix   bool
	version     string
	handleStdin bool
)

type strslice []string

func (s *strslice) String() string {
	return fmt.Sprintf("%v", *s)
}

func (s *strslice) Set(v string) error {
	*s = append(*s, v)
	return nil
}

func main() {
	var targets strslice
	var showVersion bool
	flag.Var(&targets, "t", "target hostname")
	flag.BoolVar(&addPrefix, "p", false, "add hostname to line prefix")
	flag.BoolVar(&showVersion, "v", false, "show version")
	flag.BoolVar(&handleStdin, "i", false, "handle STDIN")
	flag.Parse()
	if showVersion {
		fmt.Println("version:", version)
		return
	}

	command := flag.Args()
	if len(command) < 1 {
		flag.PrintDefaults()
		return
	}
	wg := &sync.WaitGroup{}
	stdinChs := make([]chan ([]byte), len(targets))
	if handleStdin {
		for i, _ := range targets {
			stdinChs[i] = make(chan []byte, 256)
		}
		go func() {
			wg.Add(1)
			processStdin(stdinChs)
			wg.Done()
		}()
	}
	for i, host := range targets {
		wg.Add(1)
		go func(h string, ch chan []byte) {
			remoteCommand(h, ch, command)
			wg.Done()
		}(host, stdinChs[i])
	}
	wg.Wait()
}

func processStdin(chs []chan []byte) {
	buf := make([]byte, 1024)
	for {
		n, err := io.ReadAtLeast(os.Stdin, buf, 1)
		if err != nil {
			if err != io.EOF {
				log.Println("[error]", err)
			}
			break
		}
		for _, ch := range chs {
			ch <- buf[0:n]
		}
	}
	for _, ch := range chs {
		close(ch)
	}
}

func remoteCommand(host string, src chan []byte, command []string) {
	args := []string{}
	if src == nil {
		// man ssh(1): Multiple -t options force tty allocation, even if ssh has no local tty.
		args = append(args, "-t", "-t")
	}
	args = append(args, host)
	args = append(args, command...)

	cmd := exec.Command("ssh", args...)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		fmt.Fprintln(os.Stderr, "[error]", host, err)
		return
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		fmt.Fprintln(os.Stderr, "[error]", host, err)
		return
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		fmt.Fprintln(os.Stderr, "[error]", host, err)
		return
	}
	cmd.Start()
	if src != nil {
		go writeInput(src, stdin, host)
	}
	go scanLines(stderr, os.Stderr, host)
	go scanLines(stdout, os.Stdout, host)
	cmd.Wait()
}

func scanLines(src io.ReadCloser, dest io.Writer, prefix string) {
	scanner := bufio.NewScanner(src)
	for scanner.Scan() {
		if addPrefix {
			fmt.Fprintln(dest, prefix, scanner.Text())
		} else {
			fmt.Fprintln(dest, scanner.Text())
		}
	}
}

func writeInput(src chan []byte, dest io.WriteCloser, host string) {
	for {
		b, more := <-src
		if more {
			_, err := dest.Write(b)
			if err != nil {
				log.Println("[error]", host, err)
				dest.Close()
				break
			}
		} else {
			dest.Close()
			break
		}
	}
}
<|endoftext|>


<commit_before>package nude

import (
	"errors"
	"fmt"
	"image"
	"image/gif"
	"image/jpeg"
	"image/png"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

func IsNude(imageFilePath string) (result bool, err error) {
	path, err := filepath.Abs(imageFilePath)
	if err != nil {
		return false, err
	}
	n := New(path)
	result, err = n.Parse()

	return
}

type Nude struct {
	filePath        string
	image           image.Image
	width           int
	height          int
	totalPixels     int
	skinMap         SkinMap
	SkinRegions     SkinMapList
	detectedRegions SkinMapList
	mergeRegions    [][]int
	lastFrom        int
	lastTo          int
	message         string
	result          bool
}

// experimental
func DecodeImage(filePath string) (img image.Image, err error) {
	return decodeImage(filePath)
}

func New(path string) *Nude {
	nude := &Nude{
		filePath: path,
	}
	return nude
}

func decodeImage(filePath string) (img image.Image, err error) {
	reader, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}
	defer reader.Close()

	last3Strings := strings.ToLower(filePath[len(filePath)-3:])
	last4Strings := strings.ToLower(filePath[len(filePath)-4:])
	if last3Strings == "jpg" || last4Strings == "jpeg" {
		img, err = jpeg.Decode(reader)
	} else if last3Strings == "gif" {
		img, err = gif.Decode(reader)
	} else if last3Strings == "png" {
		img, err = png.Decode(reader)
	} else {
		img = nil
		err = errors.New("unknown format")
	}
	return
}

func (nude *Nude) Parse() (result bool, err error) {
	img, err := decodeImage(nude.filePath)
	if err != nil {
		return false, err
	}
	bounds := img.Bounds()
	nude.image = img
	nude.width = bounds.Size().X
	nude.height = bounds.Size().Y
	nude.lastFrom = -1
	nude.lastTo = -1
	nude.totalPixels = nude.width * nude.height

	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		width := bounds.Size().X
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			r, g, b, _ := nude.image.At(x, y).RGBA()
			normR := r / 256
			normG := g / 256
			normB := b / 256
			currentIndex := x + y*width
			nextIndex := currentIndex + 1

			if !classifySkin(normR, normG, normB) {
				nude.skinMap = append(nude.skinMap, &Skin{currentIndex, false, 0, x, y, false})
			} else {
				nude.skinMap = append(nude.skinMap, &Skin{currentIndex, true, 0, x, y, false})

				region := -1
				checkIndexes := []int{
					nextIndex - 2,
					nextIndex - width - 2,
					nextIndex - width - 1,
					nextIndex - width,
				}
				checker := false

				for _, checkIndex := range checkIndexes {
					if checkIndex < 0 {
						continue
					}
					skin := nude.skinMap[checkIndex]
					if skin != nil && skin.skin {
						if skin.region != region &&
							region != -1 &&
							nude.lastFrom != region &&
							nude.lastTo != skin.region {
							nude.addMerge(region, skin.region)
						}
						region = nude.skinMap[checkIndex].region
						checker = true
					}
				}

				if !checker {
					nude.skinMap[currentIndex].region = len(nude.detectedRegions)
					nude.detectedRegions = append(nude.detectedRegions, []*Skin{nude.skinMap[currentIndex]})
					continue
				} else {
					if region > -1 {
						if len(nude.detectedRegions) >= region {
							nude.detectedRegions = append(nude.detectedRegions, SkinMap{})
						}
						nude.skinMap[currentIndex].region = region
						nude.detectedRegions[region] = append(nude.detectedRegions[region], nude.skinMap[currentIndex])
					}
				}
			}
		}
	}

	nude.merge(nude.detectedRegions, nude.mergeRegions)
	nude.analyzeRegions()

	return nude.result, err
}

func (nude *Nude) addMerge(from, to int) {
	nude.lastFrom = from
	nude.lastTo = to

	fromIndex := -1
	toIndex := -1

	for index, region := range nude.mergeRegions {
		for _, regionIndex := range region {
			if regionIndex == from {
				fromIndex = index
			}
			if regionIndex == to {
				toIndex = index
			}
		}
	}

	if fromIndex != -1 && toIndex != -1 {
		if fromIndex != toIndex {
			fromRegion := nude.mergeRegions[fromIndex]
			toRegion := nude.mergeRegions[toIndex]
			region := append(fromRegion, toRegion...)
			nude.mergeRegions[fromIndex] = region
			nude.mergeRegions = append(nude.mergeRegions[:toIndex], nude.mergeRegions[toIndex+1:]...)
		}
		return
	}

	if fromIndex == -1 && toIndex == -1 {
		nude.mergeRegions = append(nude.mergeRegions, []int{from, to})
		return
	}

	if fromIndex != -1 && toIndex == -1 {
		nude.mergeRegions[fromIndex] = append(nude.mergeRegions[fromIndex], to)
		return
	}

	if fromIndex == -1 && toIndex != -1 {
		nude.mergeRegions[toIndex] = append(nude.mergeRegions[toIndex], from)
		return
	}

}

// function for merging detected regions
func (nude *Nude) merge(detectedRegions SkinMapList, mergeRegions [][]int) {
	var newDetectedRegions SkinMapList

	// merging detected regions
	for index, region := range mergeRegions {
		if len(newDetectedRegions) >= index {
			newDetectedRegions = append(newDetectedRegions, SkinMap{})
		}
		for _, r := range region {
			newDetectedRegions[index] = append(newDetectedRegions[index], detectedRegions[r]...)
			detectedRegions[r] = SkinMap{}
		}
	}

	// push the rest of the regions to the newDetectedRegions array
	// (regions without merging)
	for _, region := range detectedRegions {
		if len(region) > 0 {
			newDetectedRegions = append(newDetectedRegions, region)
		}
	}

	// clean up
	nude.clearRegions(newDetectedRegions)
}

// clean up function
// only push regions which are bigger than a specific amount to the final resul
func (nude *Nude) clearRegions(detectedRegions SkinMapList) {
	for _, region := range detectedRegions {
		if len(region) > 30 {
			nude.SkinRegions = append(nude.SkinRegions, region)
		}
	}
}

func (nude *Nude) analyzeRegions() bool {
	skinRegionLength := len(nude.SkinRegions)

	// if there are less than 3 regions
	if skinRegionLength < 3 {
		nude.message = fmt.Sprintf("Less than 3 skin regions (%v)", skinRegionLength)
		nude.result = false
		return nude.result
	}

	// sort the skinRegions
	sort.Sort(nude.SkinRegions)
	//sort.Reverse(nude.SkinRegions)

	// count total skin pixels
	var totalSkin float64
	for _, region := range nude.SkinRegions {
		totalSkin += float64(len(region))
	}

	// check if there are more than 15% skin pixel in the image
	totalSkinParcentage := totalSkin / float64(nude.totalPixels) * 100
	if totalSkinParcentage < 15 {
		// if the parcentage lower than 15, it's not nude!
		nude.message = fmt.Sprintf("Total skin parcentage lower than 15 (%v%%)", totalSkinParcentage)
		nude.result = false
		return nude.result
	}

	// check if the largest skin region is less than 35% of the total skin count
	// AND if the second largest region is less than 30% of the total skin count
	// AND if the third largest region is less than 30% of the total skin count
	biggestRegionParcentage := float64(len(nude.SkinRegions[0])) / totalSkin * 100
	secondLargeRegionParcentage := float64(len(nude.SkinRegions[1])) / totalSkin * 100
	thirdLargesRegionParcentage := float64(len(nude.SkinRegions[2])) / totalSkin * 100
	if biggestRegionParcentage < 35 &&
		secondLargeRegionParcentage < 30 &&
		thirdLargesRegionParcentage < 30 {
		nude.message = "Less than 35%, 30%, 30% skin in the biggest regions"
		nude.result = false
		return nude.result
	}

	// check if the number of skin pixels in the largest region is less than 45% of the total skin count
	if biggestRegionParcentage < 45 {
		nude.message = fmt.Sprintf("The biggest region contains less than 45%% (%v)", biggestRegionParcentage)
		nude.result = false
		return nude.result
	}

	// TODO:
	// build the bounding polygon by the regions edge values:
	// Identify the leftmost, the uppermost, the rightmost, and the lowermost skin pixels of the three largest skin regions.
	// Use these points as the corner points of a bounding polygon.

	// TODO:
	// check if the total skin count is less than 30% of the total number of pixels
	// AND the number of skin pixels within the bounding polygon is less than 55% of the size of the polygon
	// if this condition is true, it's not nude.

	// TODO: include bounding polygon functionality
	// if there are more than 60 skin regions and the average intensity within the polygon is less than 0.25
	// the image is not nude
	if skinRegionLength > 60 {
		nude.message = fmt.Sprintf("More than 60 skin regions (%v)", skinRegionLength)
		nude.result = false
		return nude.result
	}

	// otherwise it is nude
	nude.result = true
	return nude.result
}

// A Survey on Pixel-Based Skin Color Detection Techniques
func classifySkin(r, g, b uint32) bool {
	rgbClassifier := r > 95 &&
		g > 40 && g < 100 &&
		b > 20 &&
		maxRgb(r, g, b)-minRgb(r, g, b) > 15 &&
		math.Abs(float64(r-g)) > 15 &&
		r > g &&
		r > b

	nr, ng, _ := toNormalizedRgb(r, g, b)
	normalizedRgbClassifier := nr/ng > 1.185 &&
		(float64(r*b))/math.Pow(float64(r+g+b), 2) > 0.107 &&
		(float64(r*g))/math.Pow(float64(r+g+b), 2) > 0.112

	h, s, _ := toHsv(r, g, b)
	hsvClassifier := h > 0 &&
		h < 35 &&
		s > 0.23 &&
		s < 0.68

	// ycc doesnt work

	result := rgbClassifier || normalizedRgbClassifier || hsvClassifier
	return result
}

func maxRgb(r, g, b uint32) float64 {
	return math.Max(math.Max(float64(r), float64(g)), float64(b))
}

func minRgb(r, g, b uint32) float64 {
	return math.Min(math.Min(float64(r), float64(g)), float64(b))
}

func toNormalizedRgb(r, g, b uint32) (nr, ng, nb float64) {
	sum := float64(r + g + b)
	nr = float64(r) / sum
	ng = float64(g) / sum
	nb = float64(b) / sum

	return nr, ng, nb
}

func toHsv(r, g, b uint32) (h, s, v float64) {
	h = 0.0
	sum := float64(r + g + b)
	max := maxRgb(r, g, b)
	min := minRgb(r, g, b)
	diff := max - min

	if max == float64(r) {
		h = float64(g-b) / diff
	} else if max == float64(g) {
		h = 2 + float64(g-r)/diff
	} else {
		h = 4 + float64(r-g)/diff
	}

	h *= 60
	if h < 0 {
		h += 360
	}

	s = 1.0 - 3.0*(min/sum)
	v = (1.0 / 3.0) * max

	return h, s, v
}
<commit_msg>rename methods's receiver name.<commit_after>package nude

import (
	"errors"
	"fmt"
	"image"
	"image/gif"
	"image/jpeg"
	"image/png"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

func IsNude(imageFilePath string) (result bool, err error) {
	path, err := filepath.Abs(imageFilePath)
	if err != nil {
		return false, err
	}
	a := New(path)
	result, err = a.Parse()

	return
}

// experimental
func DecodeImage(filePath string) (img image.Image, err error) {
	return decodeImage(filePath)
}

func decodeImage(filePath string) (img image.Image, err error) {
	reader, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}
	defer reader.Close()

	last3Strings := strings.ToLower(filePath[len(filePath)-3:])
	last4Strings := strings.ToLower(filePath[len(filePath)-4:])
	if last3Strings == "jpg" || last4Strings == "jpeg" {
		img, err = jpeg.Decode(reader)
	} else if last3Strings == "gif" {
		img, err = gif.Decode(reader)
	} else if last3Strings == "png" {
		img, err = png.Decode(reader)
	} else {
		img = nil
		err = errors.New("unknown format")
	}
	return
}

type Analyzer struct {
	filePath        string
	image           image.Image
	width           int
	height          int
	totalPixels     int
	skinMap         SkinMap
	SkinRegions     SkinMapList
	detectedRegions SkinMapList
	mergeRegions    [][]int
	lastFrom        int
	lastTo          int
	message         string
	result          bool
}

func New(path string) *Analyzer {
	analyzer := &Analyzer{
		filePath: path,
	}
	return analyzer
}

func (a *Analyzer) Parse() (result bool, err error) {
	img, err := decodeImage(a.filePath)
	if err != nil {
		return false, err
	}
	bounds := img.Bounds()
	a.image = img
	a.width = bounds.Size().X
	a.height = bounds.Size().Y
	a.lastFrom = -1
	a.lastTo = -1
	a.totalPixels = a.width * a.height

	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		width := bounds.Size().X
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			r, g, b, _ := a.image.At(x, y).RGBA()
			normR := r / 256
			normG := g / 256
			normB := b / 256
			currentIndex := x + y*width
			nextIndex := currentIndex + 1

			if !classifySkin(normR, normG, normB) {
				a.skinMap = append(a.skinMap, &Skin{currentIndex, false, 0, x, y, false})
			} else {
				a.skinMap = append(a.skinMap, &Skin{currentIndex, true, 0, x, y, false})

				region := -1
				checkIndexes := []int{
					nextIndex - 2,
					nextIndex - width - 2,
					nextIndex - width - 1,
					nextIndex - width,
				}
				checker := false

				for _, checkIndex := range checkIndexes {
					if checkIndex < 0 {
						continue
					}
					skin := a.skinMap[checkIndex]
					if skin != nil && skin.skin {
						if skin.region != region &&
							region != -1 &&
							a.lastFrom != region &&
							a.lastTo != skin.region {
							a.addMerge(region, skin.region)
						}
						region = a.skinMap[checkIndex].region
						checker = true
					}
				}

				if !checker {
					a.skinMap[currentIndex].region = len(a.detectedRegions)
					a.detectedRegions = append(a.detectedRegions, []*Skin{a.skinMap[currentIndex]})
					continue
				} else {
					if region > -1 {
						if len(a.detectedRegions) >= region {
							a.detectedRegions = append(a.detectedRegions, SkinMap{})
						}
						a.skinMap[currentIndex].region = region
						a.detectedRegions[region] = append(a.detectedRegions[region], a.skinMap[currentIndex])
					}
				}
			}
		}
	}

	a.merge(a.detectedRegions, a.mergeRegions)
	a.analyzeRegions()

	return a.result, err
}

func (a *Analyzer) addMerge(from, to int) {
	a.lastFrom = from
	a.lastTo = to

	fromIndex := -1
	toIndex := -1

	for index, region := range a.mergeRegions {
		for _, regionIndex := range region {
			if regionIndex == from {
				fromIndex = index
			}
			if regionIndex == to {
				toIndex = index
			}
		}
	}

	if fromIndex != -1 && toIndex != -1 {
		if fromIndex != toIndex {
			fromRegion := a.mergeRegions[fromIndex]
			toRegion := a.mergeRegions[toIndex]
			region := append(fromRegion, toRegion...)
			a.mergeRegions[fromIndex] = region
			a.mergeRegions = append(a.mergeRegions[:toIndex], a.mergeRegions[toIndex+1:]...)
		}
		return
	}

	if fromIndex == -1 && toIndex == -1 {
		a.mergeRegions = append(a.mergeRegions, []int{from, to})
		return
	}

	if fromIndex != -1 && toIndex == -1 {
		a.mergeRegions[fromIndex] = append(a.mergeRegions[fromIndex], to)
		return
	}

	if fromIndex == -1 && toIndex != -1 {
		a.mergeRegions[toIndex] = append(a.mergeRegions[toIndex], from)
		return
	}

}

// function for merging detected regions
func (a *Analyzer) merge(detectedRegions SkinMapList, mergeRegions [][]int) {
	var newDetectedRegions SkinMapList

	// merging detected regions
	for index, region := range mergeRegions {
		if len(newDetectedRegions) >= index {
			newDetectedRegions = append(newDetectedRegions, SkinMap{})
		}
		for _, r := range region {
			newDetectedRegions[index] = append(newDetectedRegions[index], detectedRegions[r]...)
			detectedRegions[r] = SkinMap{}
		}
	}

	// push the rest of the regions to the newDetectedRegions array
	// (regions without merging)
	for _, region := range detectedRegions {
		if len(region) > 0 {
			newDetectedRegions = append(newDetectedRegions, region)
		}
	}

	// clean up
	a.clearRegions(newDetectedRegions)
}

// clean up function
// only push regions which are bigger than a specific amount to the final resul
func (a *Analyzer) clearRegions(detectedRegions SkinMapList) {
	for _, region := range detectedRegions {
		if len(region) > 30 {
			a.SkinRegions = append(a.SkinRegions, region)
		}
	}
}

func (a *Analyzer) analyzeRegions() bool {
	skinRegionLength := len(a.SkinRegions)

	// if there are less than 3 regions
	if skinRegionLength < 3 {
		a.message = fmt.Sprintf("Less than 3 skin regions (%v)", skinRegionLength)
		a.result = false
		return a.result
	}

	// sort the skinRegions
	sort.Sort(a.SkinRegions)
	//sort.Reverse(a.SkinRegions)

	// count total skin pixels
	var totalSkin float64
	for _, region := range a.SkinRegions {
		totalSkin += float64(len(region))
	}

	// check if there are more than 15% skin pixel in the image
	totalSkinParcentage := totalSkin / float64(a.totalPixels) * 100
	if totalSkinParcentage < 15 {
		// if the parcentage lower than 15, it's not nude!
		a.message = fmt.Sprintf("Total skin parcentage lower than 15 (%v%%)", totalSkinParcentage)
		a.result = false
		return a.result
	}

	// check if the largest skin region is less than 35% of the total skin count
	// AND if the second largest region is less than 30% of the total skin count
	// AND if the third largest region is less than 30% of the total skin count
	biggestRegionParcentage := float64(len(a.SkinRegions[0])) / totalSkin * 100
	secondLargeRegionParcentage := float64(len(a.SkinRegions[1])) / totalSkin * 100
	thirdLargesRegionParcentage := float64(len(a.SkinRegions[2])) / totalSkin * 100
	if biggestRegionParcentage < 35 &&
		secondLargeRegionParcentage < 30 &&
		thirdLargesRegionParcentage < 30 {
		a.message = "Less than 35%, 30%, 30% skin in the biggest regions"
		a.result = false
		return a.result
	}

	// check if the number of skin pixels in the largest region is less than 45% of the total skin count
	if biggestRegionParcentage < 45 {
		a.message = fmt.Sprintf("The biggest region contains less than 45%% (%v)", biggestRegionParcentage)
		a.result = false
		return a.result
	}

	// TODO:
	// build the bounding polygon by the regions edge values:
	// Identify the leftmost, the uppermost, the rightmost, and the lowermost skin pixels of the three largest skin regions.
	// Use these points as the corner points of a bounding polygon.

	// TODO:
	// check if the total skin count is less than 30% of the total number of pixels
	// AND the number of skin pixels within the bounding polygon is less than 55% of the size of the polygon
	// if this condition is true, it's not nude.

	// TODO: include bounding polygon functionality
	// if there are more than 60 skin regions and the average intensity within the polygon is less than 0.25
	// the image is not nude
	if skinRegionLength > 60 {
		a.message = fmt.Sprintf("More than 60 skin regions (%v)", skinRegionLength)
		a.result = false
		return a.result
	}

	// otherwise it is nude
	a.result = true
	return a.result
}

// A Survey on Pixel-Based Skin Color Detection Techniques
func classifySkin(r, g, b uint32) bool {
	rgbClassifier := r > 95 &&
		g > 40 && g < 100 &&
		b > 20 &&
		maxRgb(r, g, b)-minRgb(r, g, b) > 15 &&
		math.Abs(float64(r-g)) > 15 &&
		r > g &&
		r > b

	nr, ng, _ := toNormalizedRgb(r, g, b)
	normalizedRgbClassifier := nr/ng > 1.185 &&
		(float64(r*b))/math.Pow(float64(r+g+b), 2) > 0.107 &&
		(float64(r*g))/math.Pow(float64(r+g+b), 2) > 0.112

	h, s, _ := toHsv(r, g, b)
	hsvClassifier := h > 0 &&
		h < 35 &&
		s > 0.23 &&
		s < 0.68

	// ycc doesnt work

	result := rgbClassifier || normalizedRgbClassifier || hsvClassifier
	return result
}

func maxRgb(r, g, b uint32) float64 {
	return math.Max(math.Max(float64(r), float64(g)), float64(b))
}

func minRgb(r, g, b uint32) float64 {
	return math.Min(math.Min(float64(r), float64(g)), float64(b))
}

func toNormalizedRgb(r, g, b uint32) (nr, ng, nb float64) {
	sum := float64(r + g + b)
	nr = float64(r) / sum
	ng = float64(g) / sum
	nb = float64(b) / sum

	return nr, ng, nb
}

func toHsv(r, g, b uint32) (h, s, v float64) {
	h = 0.0
	sum := float64(r + g + b)
	max := maxRgb(r, g, b)
	min := minRgb(r, g, b)
	diff := max - min

	if max == float64(r) {
		h = float64(g-b) / diff
	} else if max == float64(g) {
		h = 2 + float64(g-r)/diff
	} else {
		h = 4 + float64(r-g)/diff
	}

	h *= 60
	if h < 0 {
		h += 360
	}

	s = 1.0 - 3.0*(min/sum)
	v = (1.0 / 3.0) * max

	return h, s, v
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime"
	"mime/multipart"
	"net/http"
	"strings"
)

type inboundParse struct {
	Test1 string
	Test2 string
}

type configuration struct {
	Endpoint string `json:"endpoint"`
	Port     string `json:"port"`
}

func loadConfig(path string) configuration {
	file, err := ioutil.ReadFile(path)
	if err != nil {
		log.Fatal("Config File Missing. ", err)
	}
	var conf configuration
	err = json.Unmarshal(file, &conf)
	if err != nil {
		log.Fatal("Config Parse Error: ", err)
	}

	return conf
}

func indexHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "%s", "Hello World")
}

func inboundHandler(w http.ResponseWriter, r *http.Request) {
	mediaType, params, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
	if err != nil {
		log.Fatal(err)
	}
	if strings.HasPrefix(mediaType, "multipart/") {
		fmt.Println("================MESSAGE RECEIVED===============")
		mr := multipart.NewReader(r.Body, params["boundary"])
		parsedEmail := make(map[string]string)
		emailHeader := make(map[string]string)
		for {
			p, err := mr.NextPart()
			// We have found an attachment, binary data
			if err == nil && p.FileName() != "" {
				fmt.Println("FileName: ", p.FileName())
				contentType := p.Header.Get("Content-Type")
				fmt.Println("Content-Type: ", contentType)
				contents, err := ioutil.ReadAll(p)
				if err != nil {
					log.Fatal(err)
				}
				// Binary file contents
				fmt.Println("Contents: ", contents)
				// Only works with text files, this is just for testing
				fmt.Println("Contents Decoded: ", string(contents))
			}
			if err == io.EOF {
				// We have finished parsing
				for key, value := range parsedEmail {
					fmt.Println("Key:", key, " Value:", value)
				}
				for key, value := range emailHeader {
					fmt.Println("eKey:", key, " eValue:", value)
				}
				w.WriteHeader(http.StatusOK)
				return
			}
			if err != nil {
				log.Fatal(err)
			}
			value, err := ioutil.ReadAll(p)
			if err != nil {
				log.Fatal(err)
			}
			header := p.Header.Get("Content-Disposition")

			if strings.Contains(header, "filename") != true {
				header = header[17 : len(header)-1]
				parsedEmail[header] = string(value)
			} else {
				header = header[11:]
				f := strings.Split(header, "=")
				parsedEmail[f[1][1:len(f[1])-11]] = f[2][1 : len(f[2])-1]
			}
			if header == "headers" {
				s := strings.Split(string(value), "\n")
				var a []string
				for _, v := range s {
					t := strings.Split(string(v), ": ")
					a = append(a, t...)
				}
				for i := 0; i < len(a)-1; i += 2 {
					emailHeader[a[i]] = a[i+1]
				}
			}
			delete(parsedEmail, "headers")

			// We have a raw message
			if header == "email" {
				body := strings.NewReader(string(value))
				bodySplit := strings.Split(string(value), "Content-Type: multipart/mixed; ")
				scanner := bufio.NewScanner(strings.NewReader(bodySplit[1]))
				var lines []string
				for scanner.Scan() {
					lines = append(lines, scanner.Text())
					break
				}
				//fmt.Println("Split:", bodySplit[0])
				//fmt.Println("Split:", bodySplit[1])
				boundary := lines[0][9:]
				//fmt.Println("Split:", params["boundary"])

				// Get the attachments, base64 encoded
				raw := multipart.NewReader(body, boundary)
				for {
					q, err := raw.NextPart()
					if err == io.EOF {
						// We have finished parsing
						break
					}
					v, err := ioutil.ReadAll(q)
					if err != nil {
						log.Fatal(err)
					}
					header := q.Header.Get("Content-Type")

					// Parse the headers
					if strings.Contains(header, "multipart/alternative") {
						body := strings.NewReader(string(value))
						bodySplit := strings.Split(string(value), "Content-Type: multipart/alternative; ")
						scanner := bufio.NewScanner(strings.NewReader(bodySplit[1]))
						var lines []string
						for scanner.Scan() {
							lines = append(lines, scanner.Text())
							break
						}
						boundary := lines[0][9:]
						raw := multipart.NewReader(body, boundary)
						for {
							q, err := raw.NextPart()
							if err == io.EOF {
								// We have finished parsing
								break
							}
							v, err := ioutil.ReadAll(q)
							if err != nil {
								log.Fatal(err)
							}
							header := q.Header.Get("Content-Type")
							fmt.Println("iEmail:", header, "iValue:", string(v))
						}
					} else {
						fmt.Println("Email:", header, "Value:", string(v))
					}
				}
			}
			delete(parsedEmail, "email")
		}
	}
}

func main() {
	conf := loadConfig("./conf.json")
	http.HandleFunc("/", indexHandler)
	http.HandleFunc(conf.Endpoint, inboundHandler)
	http.ListenAndServe(conf.Port, nil)
}
<commit_msg>Clean up/Refactor<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime"
	"mime/multipart"
	"net/http"
	"strings"
)

type configuration struct {
	Endpoint string `json:"endpoint"`
	Port     string `json:"port"`
}

func loadConfig(path string) configuration {
	file, err := ioutil.ReadFile(path)
	if err != nil {
		log.Fatal("Config File Missing. ", err)
	}
	var conf configuration
	err = json.Unmarshal(file, &conf)
	if err != nil {
		log.Fatal("Config Parse Error: ", err)
	}
	return conf
}

func indexHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "%s", "Hello World")
}

func getBoundary(value string, contentType string) (string, *strings.Reader) {
	body := strings.NewReader(value)
	bodySplit := strings.Split(string(value), contentType)
	scanner := bufio.NewScanner(strings.NewReader(bodySplit[1]))
	var lines []string
	for scanner.Scan() {
		lines = append(lines, scanner.Text())
		break
	}
	boundary := lines[0][9:]
	return boundary, body
}

func inboundHandler(w http.ResponseWriter, r *http.Request) {
	mediaType, params, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
	if err != nil {
		log.Fatal(err)
	}
	if strings.HasPrefix(mediaType, "multipart/") {
		mr := multipart.NewReader(r.Body, params["boundary"])
		parsedEmail := make(map[string]string)
		emailHeader := make(map[string]string)
		//TODO: this needs to be a struct
		binaryFiles := make(map[string][]byte)
		rawEmail := make(map[string]string)
		rawFiles := make(map[string]string)
		for {
			p, err := mr.NextPart()
			// We have found an attachment, binary data
			if err == nil && p.FileName() != "" {
				//contentType := p.Header.Get("Content-Type")
				contents, err := ioutil.ReadAll(p)
				if err != nil {
					log.Fatal(err)
				}
				binaryFiles[p.FileName()] = contents
			}
			if err == io.EOF {
				// We have finished parsing
				for key, value := range parsedEmail {
					fmt.Println("Key:", key, " Value:", value)
				}
				for key, value := range emailHeader {
					fmt.Println("eKey:", key, " eValue:", value)
				}
				for key, value := range binaryFiles {
					fmt.Println("bKey:", key, " bValue:", value)
				}
				for key, value := range rawEmail {
					fmt.Println("rKey:", key, " rValue:", value)
				}
				for key, value := range rawFiles {
					fmt.Println("rfKey:", key, " rfValue:", value)
				}
				w.WriteHeader(http.StatusOK)
				return
			}
			if err != nil {
				log.Fatal(err)
			}
			value, err := ioutil.ReadAll(p)
			if err != nil {
				log.Fatal(err)
			}
			header := p.Header.Get("Content-Disposition")

			if strings.Contains(header, "filename") != true {
				header = header[17 : len(header)-1]
				parsedEmail[header] = string(value)
			} else {
				header = header[11:]
				f := strings.Split(header, "=")
				parsedEmail[f[1][1:len(f[1])-11]] = f[2][1 : len(f[2])-1]
			}
			if header == "headers" {
				s := strings.Split(string(value), "\n")
				var a []string
				for _, v := range s {
					t := strings.Split(string(v), ": ")
					a = append(a, t...)
				}
				for i := 0; i < len(a)-1; i += 2 {
					emailHeader[a[i]] = a[i+1]
				}
			}
			delete(parsedEmail, "headers")

			// We have a raw message
			if header == "email" {
				boundary, body := getBoundary(string(value), "Content-Type: multipart/mixed; ")
				raw := multipart.NewReader(body, boundary)
				for {
					// TODO: function1
					q, err := raw.NextPart()
					if err == io.EOF {
						// We have finished parsing
						break
					}
					v, err := ioutil.ReadAll(q)
					if err != nil {
						log.Fatal(err)
					}
					header := q.Header.Get("Content-Type")
					//--end

					// Parse the headers
					if strings.Contains(header, "multipart/alternative") {
						boundary, body := getBoundary(string(value), "Content-Type: multipart/alternative; ")
						raw := multipart.NewReader(body, boundary)
						for {
							// TODO: function1
							q, err := raw.NextPart()
							if err == io.EOF {
								// We have finished parsing
								break
							}
							v, err := ioutil.ReadAll(q)
							if err != nil {
								log.Fatal(err)
							}
							header := q.Header.Get("Content-Type")
							//--end
							rawEmail[header] = string(v)
						}
					} else {
						// Get the attachments, base64 encoded
						rawFiles[header] = string(v)
					}
				}
			}
			delete(parsedEmail, "email")
		}
	}
}

func main() {
	conf := loadConfig("./conf.json")
	http.HandleFunc("/", indexHandler)
	http.HandleFunc(conf.Endpoint, inboundHandler)
	http.ListenAndServe(conf.Port, nil)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"unicode"
)

type line struct {
	label     int
	statement string
}
type labelMap map[int]int
type corrector func(string) string

var (
	usageDescription     = makeUsageDescription()
	linePattern          = regexp.MustCompile(`^\s*(\d+)?(.*?)\s*$`)
	jumpPattern          = regexp.MustCompile(`\bGOTO\s+\d+$`)
	conditionJumpPattern = regexp.MustCompile(`\bTHEN\s+\d+$`)
	labelPattern         = regexp.MustCompile(`\d+$`)
)

func main() {
	filename := processArguments()
	code := readFile(filename)
	cuttedCode := removeEndingWhitespaces(code)
	rawLines := splitLines(cuttedCode)
	parsedLines := parseLines(rawLines)
	labels := makelabels(parsedLines)
	renumberedLines := renumberLines(parsedLines, labels)
	renumberedCode := combineLines(renumberedLines)
	updateFile(filename, renumberedCode)
}

func makeUsageDescription() string {
	_, scriptPath, _, ok := runtime.Caller(0)
	if !ok {
		return ""
	}

	return fmt.Sprintf(
		"Usage:\n"+
			"\tgo run %s [options] <filename>\n"+
			"\n"+
			"Options:\n"+
			"\t-h, --help - show help.\n",
		filepath.Base(scriptPath),
	)
}

func processArguments() string {
	testArguments()

	firstArgument := os.Args[1]
	processHelpOption(firstArgument)

	return firstArgument
}

func testArguments() {
	numberOfArguments := len(os.Args)
	if numberOfArguments < 2 {
		fmt.Print(
			"Error: filename not specified.\n" +
				"\n" +
				usageDescription,
		)

		os.Exit(1)
	}
}

func processHelpOption(firstArgument string) {
	if firstArgument == "-h" || firstArgument == "--help" {
		fmt.Print(usageDescription)
		os.Exit(0)
	}
}

func readFile(filename string) string {
	code, error := ioutil.ReadFile(filename)
	if error != nil {
		fmt.Printf(
			"Error: unable to read file \"%s\" (%v).\n",
			filename,
			error,
		)

		os.Exit(1)
	}

	return string(code)
}

func removeEndingWhitespaces(code string) string {
	return strings.TrimRightFunc(code, unicode.IsSpace)
}

func splitLines(code string) []string {
	return strings.Split(code, "\n")
}

func parseLines(lines []string) []line {
	var parsedLines []line
	for index, content := range lines {
		parsedLine := parseLine(index, content)
		parsedLines = append(parsedLines, parsedLine)
	}

	return parsedLines
}

func parseLine(index int, content string) line {
	lineParts := linePattern.FindStringSubmatch(content)
	if len(lineParts) != 3 {
		fmt.Printf("Warning: invalid line #%d.\n", index+1)
		return line{}
	}

	label := parseLabel(index, lineParts[1])
	return line{label, lineParts[2]}
}

func parseLabel(index int, stringLabel string) int {
	integralLabel, error := strconv.Atoi(stringLabel)
	if error != nil && len(stringLabel) != 0 {
		fmt.Printf(
			"Warning: invalid label \"%s\" on line #%d.\n",
			stringLabel,
			index+1,
		)
	}

	return integralLabel
}

func makelabels(lines []line) labelMap {
	labels := make(labelMap)
	for index, parsedLine := range lines {
		labels[parsedLine.label] = (index + 1) * 10
	}

	return labels
}

func renumberLines(lines []line, labels labelMap) []line {
	var renumberedLines []line
	for index, parsedLine := range lines {
		renumberedLine := renumberLine(index, parsedLine, labels)
		renumberedLines = append(renumberedLines, renumberedLine)
	}

	return renumberedLines
}

func renumberLine(index int, parsedLine line, labels labelMap) line {
	correctedLabel := labels[parsedLine.label]
	correctedStatement := correctStatement(
		index,
		parsedLine.statement,
		labels,
	)
	return line{correctedLabel, correctedStatement}
}

func correctStatement(index int, statement string, labels labelMap) string {
	statement = jumpPattern.ReplaceAllStringFunc(
		statement,
		jumpCorrector(index, labels),
	)
	statement = conditionJumpPattern.ReplaceAllStringFunc(
		statement,
		jumpCorrector(index, labels),
	)
	return statement
}

func jumpCorrector(index int, labels labelMap) corrector {
	return func(match string) string {
		return labelPattern.ReplaceAllStringFunc(
			match,
			labelCorrector(index, labels),
		)
	}
}

func labelCorrector(index int, labels labelMap) corrector {
	return func(match string) string {
		oldLabel := parseLabel(index, match)
		newLabel := labels[oldLabel]
		return strconv.Itoa(newLabel)
	}
}

func combineLines(lines []line) string {
	code := ""
	for _, renumberedLine := range lines {
		code += fmt.Sprintf(
			"%d%s\n",
			renumberedLine.label,
			renumberedLine.statement,
		)
	}

	return code
}

func updateFile(filename string, code string) {
	error := ioutil.WriteFile(filename, []byte(code), 0644)
	if error != nil {
		fmt.Printf(
			"Error: unable to update file \"%s\" (%v).\n",
			filename,
			error,
		)

		os.Exit(1)
	}
}
<commit_msg>Fixes #6: small correct.<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"unicode"
)

type line struct {
	label     int
	statement string
}
type labelMap map[int]int
type corrector func(string) string

var (
	usageDescription     = makeUsageDescription()
	linePattern          = regexp.MustCompile(`^\s*(\d+)?(.*?)\s*$`)
	jumpPattern          = regexp.MustCompile(`\bGOTO\s+\d+$`)
	conditionJumpPattern = regexp.MustCompile(`\bTHEN\s+\d+$`)
	labelPattern         = regexp.MustCompile(`\d+$`)
)

func main() {
	filename := processArguments()
	code := readFile(filename)
	cuttedCode := removeEndingWhitespaces(code)
	rawLines := splitLines(cuttedCode)
	parsedLines := parseLines(rawLines)
	labels := makeLabelMap(parsedLines)
	renumberedLines := renumberLines(parsedLines, labels)
	renumberedCode := combineLines(renumberedLines)
	updateFile(filename, renumberedCode)
}

func makeUsageDescription() string {
	_, scriptPath, _, ok := runtime.Caller(0)
	if !ok {
		return ""
	}

	return fmt.Sprintf(
		"Usage:\n"+
			"\tgo run %s [options] <filename>\n"+
			"\n"+
			"Options:\n"+
			"\t-h, --help - show help.\n",
		filepath.Base(scriptPath),
	)
}

func processArguments() string {
	testArguments()

	firstArgument := os.Args[1]
	processHelpOption(firstArgument)

	return firstArgument
}

func testArguments() {
	numberOfArguments := len(os.Args)
	if numberOfArguments < 2 {
		fmt.Print(
			"Error: filename not specified.\n" +
				"\n" +
				usageDescription,
		)

		os.Exit(1)
	}
}

func processHelpOption(firstArgument string) {
	if firstArgument == "-h" || firstArgument == "--help" {
		fmt.Print(usageDescription)
		os.Exit(0)
	}
}

func readFile(filename string) string {
	code, error := ioutil.ReadFile(filename)
	if error != nil {
		fmt.Printf(
			"Error: unable to read file \"%s\" (%v).\n",
			filename,
			error,
		)

		os.Exit(1)
	}

	return string(code)
}

func removeEndingWhitespaces(code string) string {
	return strings.TrimRightFunc(code, unicode.IsSpace)
}

func splitLines(code string) []string {
	return strings.Split(code, "\n")
}

func parseLines(lines []string) []line {
	var parsedLines []line
	for index, content := range lines {
		parsedLine := parseLine(index, content)
		parsedLines = append(parsedLines, parsedLine)
	}

	return parsedLines
}

func parseLine(index int, content string) line {
	lineParts := linePattern.FindStringSubmatch(content)
	if len(lineParts) != 3 {
		fmt.Printf("Warning: invalid line #%d.\n", index+1)
		return line{}
	}

	label := parseLabel(index, lineParts[1])
	return line{label, lineParts[2]}
}

func parseLabel(index int, stringLabel string) int {
	integralLabel, error := strconv.Atoi(stringLabel)
	if error != nil && len(stringLabel) != 0 {
		fmt.Printf(
			"Warning: invalid label \"%s\" on line #%d.\n",
			stringLabel,
			index+1,
		)
	}

	return integralLabel
}

func makeLabelMap(lines []line) labelMap {
	labels := make(labelMap)
	for index, parsedLine := range lines {
		labels[parsedLine.label] = (index + 1) * 10
	}

	return labels
}

func renumberLines(lines []line, labels labelMap) []line {
	var renumberedLines []line
	for index, parsedLine := range lines {
		renumberedLine := renumberLine(index, parsedLine, labels)
		renumberedLines = append(renumberedLines, renumberedLine)
	}

	return renumberedLines
}

func renumberLine(index int, parsedLine line, labels labelMap) line {
	correctedLabel := labels[parsedLine.label]
	correctedStatement := correctStatement(
		index,
		parsedLine.statement,
		labels,
	)
	return line{correctedLabel, correctedStatement}
}

func correctStatement(index int, statement string, labels labelMap) string {
	statement = jumpPattern.ReplaceAllStringFunc(
		statement,
		jumpCorrector(index, labels),
	)
	statement = conditionJumpPattern.ReplaceAllStringFunc(
		statement,
		jumpCorrector(index, labels),
	)
	return statement
}

func jumpCorrector(index int, labels labelMap) corrector {
	return func(match string) string {
		return labelPattern.ReplaceAllStringFunc(
			match,
			labelCorrector(index, labels),
		)
	}
}

func labelCorrector(index int, labels labelMap) corrector {
	return func(match string) string {
		oldLabel := parseLabel(index, match)
		newLabel := labels[oldLabel]
		return strconv.Itoa(newLabel)
	}
}

func combineLines(lines []line) string {
	code := ""
	for _, renumberedLine := range lines {
		code += fmt.Sprintf(
			"%d%s\n",
			renumberedLine.label,
			renumberedLine.statement,
		)
	}

	return code
}

func updateFile(filename string, code string) {
	error := ioutil.WriteFile(filename, []byte(code), 0644)
	if error != nil {
		fmt.Printf(
			"Error: unable to update file \"%s\" (%v).\n",
			filename,
			error,
		)

		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"html/template"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"strings"
	"time"

	"v.io/core/veyron/services/identity"
	"v.io/core/veyron/services/identity/oauth"
	"v.io/core/veyron2/context"
	"v.io/core/veyron2/options"
	"v.io/core/veyron2/security"
	"v.io/core/veyron2/vlog"
)

func exchangeMacaroonForBlessing(ctx *context.T, macaroonChan <-chan string) (security.Blessings, error) {
	service, macaroon, rootKey, err := prepareBlessArgs(ctx, macaroonChan)
	if err != nil {
		return nil, err
	}

	ctx, cancel := context.WithTimeout(ctx, time.Minute)
	defer cancel()

	var reply security.WireBlessings
	reply, err = identity.MacaroonBlesserClient(service).Bless(ctx, macaroon, options.ServerPublicKey{rootKey})
	if err != nil {
		return nil, fmt.Errorf("failed to get blessing from %q: %v", service, err)
	}
	blessings, err := security.NewBlessings(reply)
	if err != nil {
		return nil, fmt.Errorf("failed to construct Blessings object from response: %v", err)
	}
	return blessings, nil
}

func prepareBlessArgs(ctx *context.T, macaroonChan <-chan string) (service, macaroon string, root security.PublicKey, err error) {
	macaroon = <-macaroonChan
	service = <-macaroonChan

	marshalKey, err := base64.URLEncoding.DecodeString(<-macaroonChan)
	if err != nil {
		return "", "", nil, fmt.Errorf("failed to decode root key: %v", err)
	}
	root, err = security.UnmarshalPublicKey(marshalKey)
	if err != nil {
		return "", "", nil, fmt.Errorf("failed to unmarshal root key: %v", err)
	}

	return service, macaroon, root, nil
}

func getMacaroonForBlessRPC(blessServerURL string, blessedChan <-chan string, browser bool) (<-chan string, error) {
	// Setup a HTTP server to recieve a blessing macaroon from the identity server.
	// Steps:
	// 1. Generate a state token to be included in the HTTP request
	//    (though, arguably, the random port assigment for the HTTP server is enough
	//    for XSRF protection)
	// 2. Setup a HTTP server which will receive the final blessing macaroon from the id server.
	// 3. Print out the link (to start the auth flow) for the user to click.
	// 4. Return the macaroon and the rpc object name(where to make the MacaroonBlesser.Bless RPC call)
	//    in the "result" channel.
	var stateBuf [32]byte
	if _, err := rand.Read(stateBuf[:]); err != nil {
		return nil, fmt.Errorf("failed to generate state token for OAuth: %v", err)
	}
	state := base64.URLEncoding.EncodeToString(stateBuf[:])

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, fmt.Errorf("failed to setup authorization code interception server: %v", err)
	}
	result := make(chan string)

	redirectURL := fmt.Sprintf("http://%s/macaroon", ln.Addr())
	http.HandleFunc("/macaroon", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/html")
		tmplArgs := struct {
			Blessings, ErrShort, ErrLong string
		}{}
		defer func() {
			if len(tmplArgs.ErrShort) > 0 {
				w.WriteHeader(http.StatusBadRequest)
			}
			if err := tmpl.Execute(w, tmplArgs); err != nil {
				vlog.Info("Failed to render template:", err)
			}
		}()

		toolState := r.FormValue("state")
		if toolState != state {
			tmplArgs.ErrShort = "Unexpected request"
			tmplArgs.ErrLong = "Mismatched state parameter. Possible cross-site-request-forgery?"
			return
		}
		result <- r.FormValue("macaroon")
		result <- r.FormValue("object_name")
		result <- r.FormValue("root_key")
		defer close(result)
		blessed, ok := <-blessedChan
		if !ok {
			tmplArgs.ErrShort = "No blessings received"
			tmplArgs.ErrLong = "Unable to obtain blessings from the Veyron service"
			return
		}
		tmplArgs.Blessings = blessed
		ln.Close()
	})
	go http.Serve(ln, nil)

	// Print the link to start the flow.
	url, err := seekBlessingsURL(blessServerURL, redirectURL, state)
	if err != nil {
		return nil, fmt.Errorf("failed to create seekBlessingsURL: %s", err)
	}
	fmt.Fprintln(os.Stdout, "Please visit the following URL to seek blessings:")
	fmt.Fprintln(os.Stdout, url)
	// Make an attempt to start the browser as a convenience.
	// If it fails, doesn't matter - the client can see the URL printed above.
	// Use exec.Command().Start instead of exec.Command().Run since there is no
	// need to wait for the command to return (and indeed on some window managers,
	// the command will not exit until the browser is closed).
	if len(openCommand) != 0 && browser {
		exec.Command(openCommand, url).Start()
	}
	return result, nil
}

func seekBlessingsURL(blessServerURL, redirectURL, state string) (string, error) {
	baseURL, err := url.Parse(joinURL(blessServerURL, oauth.SeekBlessingsRoute))
	if err != nil {
		return "", fmt.Errorf("failed to parse url: %v", err)
	}
	params := url.Values{}
	params.Add("redirect_url", redirectURL)
	params.Add("state", state)
	baseURL.RawQuery = params.Encode()
	return baseURL.String(), nil
}

func joinURL(baseURL, suffix string) string {
	if !strings.HasSuffix(baseURL, "/") {
		baseURL += "/"
	}
	return baseURL + suffix
}

var tmpl = template.Must(template.New("name").Parse(`<!doctype html>

<meta charset="UTF-8">
<title>Veyron Identity: Google</title>
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<link rel="stylesheet" href="//netdna.bootstrapcdn.com/bootstrap/3.2.0/css/bootstrap.min.css">
{{if .Blessings}}

{{end}}

{{if .ErrShort}}

error{{.ErrShort}}

{{.ErrLong}}

{{else}}

Received blessings: <tt>{{.Blessings}}</tt>

{{end}}

`))
<commit_msg>TBR: tools/principal: Fix missing change in commit 799aa054d3c65bd34c48f98fcf8ecc3844731162<commit_after>package main

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"html/template"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"strings"
	"time"

	"v.io/core/veyron/services/identity"
	"v.io/core/veyron/services/identity/oauth"
	"v.io/core/veyron2/context"
	"v.io/core/veyron2/options"
	"v.io/core/veyron2/security"
	"v.io/core/veyron2/vlog"
)

func exchangeMacaroonForBlessing(ctx *context.T, macaroonChan <-chan string) (security.Blessings, error) {
	service, macaroon, rootKey, err := prepareBlessArgs(ctx, macaroonChan)
	if err != nil {
		return nil, err
	}

	ctx, cancel := context.WithTimeout(ctx, time.Minute)
	defer cancel()

	var reply security.WireBlessings
	// Authorize the server by its public key (obtained from macaroonChan).
	// Must skip authorization during name resolution because the identity
	// service is not a trusted root yet.
	reply, err = identity.MacaroonBlesserClient(service).Bless(ctx, macaroon, options.SkipResolveAuthorization{}, options.ServerPublicKey{rootKey})
	if err != nil {
		return nil, fmt.Errorf("failed to get blessing from %q: %v", service, err)
	}
	blessings, err := security.NewBlessings(reply)
	if err != nil {
		return nil, fmt.Errorf("failed to construct Blessings object from response: %v", err)
	}
	return blessings, nil
}

func prepareBlessArgs(ctx *context.T, macaroonChan <-chan string) (service, macaroon string, root security.PublicKey, err error) {
	macaroon = <-macaroonChan
	service = <-macaroonChan

	marshalKey, err := base64.URLEncoding.DecodeString(<-macaroonChan)
	if err != nil {
		return "", "", nil, fmt.Errorf("failed to decode root key: %v", err)
	}
	root, err = security.UnmarshalPublicKey(marshalKey)
	if err != nil {
		return "", "", nil, fmt.Errorf("failed to unmarshal root key: %v", err)
	}

	return service, macaroon, root, nil
}

func getMacaroonForBlessRPC(blessServerURL string, blessedChan <-chan string, browser bool) (<-chan string, error) {
	// Setup a HTTP server to recieve a blessing macaroon from the identity server.
	// Steps:
	// 1. Generate a state token to be included in the HTTP request
	//    (though, arguably, the random port assigment for the HTTP server is enough
	//    for XSRF protection)
	// 2. Setup a HTTP server which will receive the final blessing macaroon from the id server.
	// 3. Print out the link (to start the auth flow) for the user to click.
	// 4. Return the macaroon and the rpc object name(where to make the MacaroonBlesser.Bless RPC call)
	//    in the "result" channel.
	var stateBuf [32]byte
	if _, err := rand.Read(stateBuf[:]); err != nil {
		return nil, fmt.Errorf("failed to generate state token for OAuth: %v", err)
	}
	state := base64.URLEncoding.EncodeToString(stateBuf[:])

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, fmt.Errorf("failed to setup authorization code interception server: %v", err)
	}
	result := make(chan string)

	redirectURL := fmt.Sprintf("http://%s/macaroon", ln.Addr())
	http.HandleFunc("/macaroon", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/html")
		tmplArgs := struct {
			Blessings, ErrShort, ErrLong string
		}{}
		defer func() {
			if len(tmplArgs.ErrShort) > 0 {
				w.WriteHeader(http.StatusBadRequest)
			}
			if err := tmpl.Execute(w, tmplArgs); err != nil {
				vlog.Info("Failed to render template:", err)
			}
		}()

		toolState := r.FormValue("state")
		if toolState != state {
			tmplArgs.ErrShort = "Unexpected request"
			tmplArgs.ErrLong = "Mismatched state parameter. Possible cross-site-request-forgery?"
			return
		}
		result <- r.FormValue("macaroon")
		result <- r.FormValue("object_name")
		result <- r.FormValue("root_key")
		defer close(result)
		blessed, ok := <-blessedChan
		if !ok {
			tmplArgs.ErrShort = "No blessings received"
			tmplArgs.ErrLong = "Unable to obtain blessings from the Veyron service"
			return
		}
		tmplArgs.Blessings = blessed
		ln.Close()
	})
	go http.Serve(ln, nil)

	// Print the link to start the flow.
	url, err := seekBlessingsURL(blessServerURL, redirectURL, state)
	if err != nil {
		return nil, fmt.Errorf("failed to create seekBlessingsURL: %s", err)
	}
	fmt.Fprintln(os.Stdout, "Please visit the following URL to seek blessings:")
	fmt.Fprintln(os.Stdout, url)
	// Make an attempt to start the browser as a convenience.
	// If it fails, doesn't matter - the client can see the URL printed above.
	// Use exec.Command().Start instead of exec.Command().Run since there is no
	// need to wait for the command to return (and indeed on some window managers,
	// the command will not exit until the browser is closed).
	if len(openCommand) != 0 && browser {
		exec.Command(openCommand, url).Start()
	}
	return result, nil
}

func seekBlessingsURL(blessServerURL, redirectURL, state string) (string, error) {
	baseURL, err := url.Parse(joinURL(blessServerURL, oauth.SeekBlessingsRoute))
	if err != nil {
		return "", fmt.Errorf("failed to parse url: %v", err)
	}
	params := url.Values{}
	params.Add("redirect_url", redirectURL)
	params.Add("state", state)
	baseURL.RawQuery = params.Encode()
	return baseURL.String(), nil
}

func joinURL(baseURL, suffix string) string {
	if !strings.HasSuffix(baseURL, "/") {
		baseURL += "/"
	}
	return baseURL + suffix
}

var tmpl = template.Must(template.New("name").Parse(`<!doctype html>

<meta charset="UTF-8">
<title>Veyron Identity: Google</title>
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<link rel="stylesheet" href="//netdna.bootstrapcdn.com/bootstrap/3.2.0/css/bootstrap.min.css">
{{if .Blessings}}

{{end}}

{{if .ErrShort}}

error{{.ErrShort}}

{{.ErrLong}}

{{else}}

Received blessings: <tt>{{.Blessings}}</tt>

{{end}}

`))
<|endoftext|>


<commit_before>package rain

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"io"
	"io/ioutil"
	"net"
	"sync"
	"time"

	"github.com/cenkalti/log"

	"github.com/cenkalti/rain/bitfield"
	"github.com/cenkalti/rain/bt"
)

const connReadTimeout = 3 * time.Minute

// Reject requests larger than this size.
const maxAllowedBlockSize = 32 * 1024

type Peer struct {
	conn     net.Conn
	id       bt.PeerID
	transfer *transfer

	disconnected bool
	amInterested bool
	peerChoking  bool

	// pieces that the peer has
	bitfield *bitfield.Bitfield

	cond *sync.Cond
	log  log.Logger
}

type Request struct {
	Peer *Peer
	requestMessage
}
type requestMessage struct {
	Index, Begin, Length uint32
}

type PieceMessage struct {
	pieceMessage
	Data chan []byte
}
type pieceMessage struct {
	Index, Begin uint32
}

func (t *transfer) newPeer(conn net.Conn, id bt.PeerID, l log.Logger) *Peer {
	p := &Peer{
		conn:        conn,
		id:          id,
		transfer:    t,
		peerChoking: true,
		bitfield:    bitfield.New(t.bitfield.Len()),
		log:         l,
	}
	p.cond = sync.NewCond(&t.m)
	return p
}

func (p *Peer) String() string { return p.conn.RemoteAddr().String() }
func (p *Peer) Close() error   { return p.conn.Close() }

// Run reads and processes incoming messages after handshake.
func (p *Peer) Run() {
	p.log.Debugln("Communicating peer", p.conn.RemoteAddr())

	go p.downloader()

	defer func() {
		for i := uint32(0); i < p.bitfield.Len(); i++ {
			if p.bitfield.Test(i) {
				delete(p.transfer.pieces[i].peers, p.id)
			}
		}
	}()

	defer func() {
		p.transfer.m.Lock()
		p.disconnected = true
		p.transfer.m.Unlock()
		p.cond.Broadcast()
	}()

	first := true
	for {
		err := p.conn.SetReadDeadline(time.Now().Add(connReadTimeout))
		if err != nil {
			p.log.Error(err)
			return
		}

		var length uint32
		p.log.Debug("Reading message...")
		err = binary.Read(p.conn, binary.BigEndian, &length)
		if err != nil {
			if err == io.EOF {
				p.log.Warning("Remote peer has closed the connection")
				return
			}
			p.log.Error(err)
			return
		}
		p.log.Debugf("Received message of length: %d", length)

		if length == 0 { // keep-alive message
			p.log.Debug("Received message of type \"keep alive\"")
			continue
		}

		var id messageID
		err = binary.Read(p.conn, binary.BigEndian, &id)
		if err != nil {
			p.log.Error(err)
			return
		}
		length--

		p.log.Debugf("Received message of type: %q", id)

		switch id {
		case chokeID:
			p.transfer.m.Lock()
			// Discard all pending requests. TODO
			p.peerChoking = true
			p.transfer.m.Unlock()
			p.cond.Broadcast()
		case unchokeID:
			p.transfer.m.Lock()
			p.peerChoking = false
			p.transfer.m.Unlock()
			p.cond.Broadcast()
		case interestedID:
			// TODO this should not be here
			if err := p.Unchoke(); err != nil {
				p.log.Error(err)
				return
			}
		case notInterestedID:
		case haveID:
			var i uint32
			err = binary.Read(p.conn, binary.BigEndian, &i)
			if err != nil {
				p.log.Error(err)
				return
			}
			if i >= p.transfer.torrent.Info.NumPieces {
				p.log.Error("unexpected piece index")
				return
			}
			p.log.Debug("Peer ", p.conn.RemoteAddr(), " has piece #", i)
			p.bitfield.Set(i)
			p.handleHave(i)
		case bitfieldID:
			if !first {
				p.log.Error("bitfield can only be sent after handshake")
				return
			}

			if length != uint32(len(p.transfer.bitfield.Bytes())) {
				p.log.Error("invalid bitfield length")
				return
			}

			p.transfer.m.Lock()
			_, err = p.conn.Read(p.bitfield.Bytes())
			p.transfer.m.Unlock()
			if err != nil {
				p.log.Error(err)
				return
			}
			p.log.Debugln("Received bitfield:", p.bitfield.Hex())

			for i := uint32(0); i < p.bitfield.Len(); i++ {
				if p.bitfield.Test(i) {
					p.handleHave(i)
				}
			}
		case requestID:
			var req requestMessage
			err = binary.Read(p.conn, binary.BigEndian, &req)
			if err != nil {
				p.log.Error(err)
				return
			}
			p.log.Debugf("Request: %+v", req)

			if req.Index >= p.transfer.torrent.Info.NumPieces {
				p.log.Error("invalid request: index")
				return
			}
			if req.Length > maxAllowedBlockSize {
				p.log.Error("received a request with block size larger than allowed")
				return
			}
			if req.Begin+req.Length > p.transfer.pieces[req.Index].Length {
				p.log.Error("invalid request: length")
			}

			p.transfer.requestC <- &Request{p, req}
		case pieceID:
			var msg pieceMessage
			err = binary.Read(p.conn, binary.BigEndian, &msg)
			if err != nil {
				p.log.Error(err)
				return
			}
			length -= 8

			if msg.Index >= p.transfer.torrent.Info.NumPieces {
				p.log.Error("invalid request: index")
				return
			}
			piece := p.transfer.pieces[msg.Index]

			// We request only in blockSize length
			blockIndex, mod := divMod32(msg.Begin, blockSize)
			if mod != 0 {
				p.log.Error("unexpected block begin")
				return
			}
			if blockIndex >= uint32(len(piece.Blocks)) {
				p.log.Error("invalid block begin")
				return
			}
			block := p.transfer.pieces[msg.Index].Blocks[blockIndex]
			if length != block.Length {
				p.log.Error("invalid piece block length")
				return
			}

			p.transfer.m.Lock()
			active := piece.getActiveRequest(p.id)
			if active == nil {
				p.transfer.m.Unlock()
				p.log.Warning("received a piece that is not activeed")
				continue
			}

			if active.blocksReceiving.Test(block.Index) {
				p.log.Warningf("Receiving duplicate block: Piece #%d Block #%d", piece.Index, block.Index)
			} else {
				active.blocksReceiving.Set(block.Index)
			}
			p.transfer.m.Unlock()

			if _, err = io.ReadFull(p.conn, active.data[msg.Begin:msg.Begin+length]); err != nil {
				p.log.Error(err)
				return
			}

			p.transfer.m.Lock()
			active.blocksReceived.Set(block.Index)
			if !active.blocksReceived.All() {
				p.transfer.m.Unlock()
				continue
			}
			p.transfer.m.Unlock()

			p.log.Debugf("Writing piece to disk: #%d", piece.Index)
			if _, err = piece.Write(active.data); err != nil {
				p.log.Error(err)
				p.conn.Close()
				return
			}

			p.transfer.m.Lock()
			p.transfer.bitfield.Set(piece.Index)
			p.transfer.m.Unlock()
			p.cond.Broadcast()
		case cancelID:
		case portID:
		default:
			p.log.Debugf("Unknown message type: %d", id)
			p.log.Debugln("Discarding", length, "bytes...")
			io.CopyN(ioutil.Discard, p.conn, int64(length))
			p.log.Debug("Discarding finished.")
		}

		first = false
	}
}

func (p *Peer) handleHave(i uint32) {
	p.transfer.m.Lock()
	p.transfer.pieces[i].peers[p.id] = struct{}{}
	p.transfer.m.Unlock()
	p.cond.Broadcast()
}

func (p *Peer) SendBitfield() error {
	// Do not send a bitfield message if we don't have any pieces.
	if p.transfer.bitfield.Count() == 0 {
		return nil
	}
	return p.sendMessage(bitfieldID, p.transfer.bitfield.Bytes())
}

func (p *Peer) BeInterested() error {
	if p.amInterested {
		return nil
	}
	p.amInterested = true
	return p.sendMessage(interestedID, nil)
}

func (p *Peer) BeNotInterested() error {
	if !p.amInterested {
		return nil
	}
	p.amInterested = false
	return p.sendMessage(notInterestedID, nil)
}

func (p *Peer) Choke() error   { return p.sendMessage(chokeID, nil) }
func (p *Peer) Unchoke() error { return p.sendMessage(unchokeID, nil) }

func (p *Peer) Request(b *Block) error {
	req := requestMessage{b.Piece.Index, b.Begin, b.Length}
	buf := bytes.NewBuffer(make([]byte, 0, 12))
	binary.Write(buf, binary.BigEndian, &req)
	return p.sendMessage(requestID, buf.Bytes())
}

func (p *Peer) SendPiece(index, begin uint32, block []byte) error {
	msg := &pieceMessage{index, begin}
	buf := bytes.NewBuffer(make([]byte, 0, 8))
	binary.Write(buf, binary.BigEndian, msg)
	buf.Write(block)
	return p.sendMessage(pieceID, buf.Bytes())
}

func (p *Peer) sendMessage(id messageID, payload []byte) error {
	p.log.Debugf("Sending message of type: %q", id)
	buf := bufio.NewWriterSize(p.conn, 4+1+len(payload))
	var header = struct {
		Length uint32
		ID     messageID
	}{
		uint32(1 + len(payload)),
		id,
	}
	binary.Write(buf, binary.BigEndian, &header)
	buf.Write(payload)
	return buf.Flush()
}
<commit_msg>fix data race<commit_after>package rain

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"io"
	"io/ioutil"
	"net"
	"sync"
	"time"

	"github.com/cenkalti/log"

	"github.com/cenkalti/rain/bitfield"
	"github.com/cenkalti/rain/bt"
)

const connReadTimeout = 3 * time.Minute

// Reject requests larger than this size.
const maxAllowedBlockSize = 32 * 1024

type Peer struct {
	conn     net.Conn
	id       bt.PeerID
	transfer *transfer

	disconnected bool
	amInterested bool
	peerChoking  bool

	amInterestedM sync.Mutex

	// pieces that the peer has
	bitfield *bitfield.Bitfield

	cond *sync.Cond
	log  log.Logger
}

type Request struct {
	Peer *Peer
	requestMessage
}
type requestMessage struct {
	Index, Begin, Length uint32
}

type PieceMessage struct {
	pieceMessage
	Data chan []byte
}
type pieceMessage struct {
	Index, Begin uint32
}

func (t *transfer) newPeer(conn net.Conn, id bt.PeerID, l log.Logger) *Peer {
	p := &Peer{
		conn:        conn,
		id:          id,
		transfer:    t,
		peerChoking: true,
		bitfield:    bitfield.New(t.bitfield.Len()),
		log:         l,
	}
	p.cond = sync.NewCond(&t.m)
	return p
}

func (p *Peer) String() string { return p.conn.RemoteAddr().String() }
func (p *Peer) Close() error   { return p.conn.Close() }

// Run reads and processes incoming messages after handshake.
func (p *Peer) Run() {
	p.log.Debugln("Communicating peer", p.conn.RemoteAddr())

	go p.downloader()

	defer func() {
		for i := uint32(0); i < p.bitfield.Len(); i++ {
			if p.bitfield.Test(i) {
				delete(p.transfer.pieces[i].peers, p.id)
			}
		}
	}()

	defer func() {
		p.transfer.m.Lock()
		p.disconnected = true
		p.transfer.m.Unlock()
		p.cond.Broadcast()
	}()

	first := true
	for {
		err := p.conn.SetReadDeadline(time.Now().Add(connReadTimeout))
		if err != nil {
			p.log.Error(err)
			return
		}

		var length uint32
		p.log.Debug("Reading message...")
		err = binary.Read(p.conn, binary.BigEndian, &length)
		if err != nil {
			if err == io.EOF {
				p.log.Warning("Remote peer has closed the connection")
				return
			}
			p.log.Error(err)
			return
		}
		p.log.Debugf("Received message of length: %d", length)

		if length == 0 { // keep-alive message
			p.log.Debug("Received message of type \"keep alive\"")
			continue
		}

		var id messageID
		err = binary.Read(p.conn, binary.BigEndian, &id)
		if err != nil {
			p.log.Error(err)
			return
		}
		length--

		p.log.Debugf("Received message of type: %q", id)

		switch id {
		case chokeID:
			p.transfer.m.Lock()
			// Discard all pending requests. TODO
			p.peerChoking = true
			p.transfer.m.Unlock()
			p.cond.Broadcast()
		case unchokeID:
			p.transfer.m.Lock()
			p.peerChoking = false
			p.transfer.m.Unlock()
			p.cond.Broadcast()
		case interestedID:
			// TODO this should not be here
			if err := p.Unchoke(); err != nil {
				p.log.Error(err)
				return
			}
		case notInterestedID:
		case haveID:
			var i uint32
			err = binary.Read(p.conn, binary.BigEndian, &i)
			if err != nil {
				p.log.Error(err)
				return
			}
			if i >= p.transfer.torrent.Info.NumPieces {
				p.log.Error("unexpected piece index")
				return
			}
			p.log.Debug("Peer ", p.conn.RemoteAddr(), " has piece #", i)
			p.bitfield.Set(i)
			p.handleHave(i)
		case bitfieldID:
			if !first {
				p.log.Error("bitfield can only be sent after handshake")
				return
			}

			if length != uint32(len(p.transfer.bitfield.Bytes())) {
				p.log.Error("invalid bitfield length")
				return
			}

			p.transfer.m.Lock()
			_, err = p.conn.Read(p.bitfield.Bytes())
			p.transfer.m.Unlock()
			if err != nil {
				p.log.Error(err)
				return
			}
			p.log.Debugln("Received bitfield:", p.bitfield.Hex())

			for i := uint32(0); i < p.bitfield.Len(); i++ {
				if p.bitfield.Test(i) {
					p.handleHave(i)
				}
			}
		case requestID:
			var req requestMessage
			err = binary.Read(p.conn, binary.BigEndian, &req)
			if err != nil {
				p.log.Error(err)
				return
			}
			p.log.Debugf("Request: %+v", req)

			if req.Index >= p.transfer.torrent.Info.NumPieces {
				p.log.Error("invalid request: index")
				return
			}
			if req.Length > maxAllowedBlockSize {
				p.log.Error("received a request with block size larger than allowed")
				return
			}
			if req.Begin+req.Length > p.transfer.pieces[req.Index].Length {
				p.log.Error("invalid request: length")
			}

			p.transfer.requestC <- &Request{p, req}
		case pieceID:
			var msg pieceMessage
			err = binary.Read(p.conn, binary.BigEndian, &msg)
			if err != nil {
				p.log.Error(err)
				return
			}
			length -= 8

			if msg.Index >= p.transfer.torrent.Info.NumPieces {
				p.log.Error("invalid request: index")
				return
			}
			piece := p.transfer.pieces[msg.Index]

			// We request only in blockSize length
			blockIndex, mod := divMod32(msg.Begin, blockSize)
			if mod != 0 {
				p.log.Error("unexpected block begin")
				return
			}
			if blockIndex >= uint32(len(piece.Blocks)) {
				p.log.Error("invalid block begin")
				return
			}
			block := p.transfer.pieces[msg.Index].Blocks[blockIndex]
			if length != block.Length {
				p.log.Error("invalid piece block length")
				return
			}

			p.transfer.m.Lock()
			active := piece.getActiveRequest(p.id)
			if active == nil {
				p.transfer.m.Unlock()
				p.log.Warning("received a piece that is not activeed")
				continue
			}

			if active.blocksReceiving.Test(block.Index) {
				p.log.Warningf("Receiving duplicate block: Piece #%d Block #%d", piece.Index, block.Index)
			} else {
				active.blocksReceiving.Set(block.Index)
			}
			p.transfer.m.Unlock()

			if _, err = io.ReadFull(p.conn, active.data[msg.Begin:msg.Begin+length]); err != nil {
				p.log.Error(err)
				return
			}

			p.transfer.m.Lock()
			active.blocksReceived.Set(block.Index)
			if !active.blocksReceived.All() {
				p.transfer.m.Unlock()
				continue
			}
			p.transfer.m.Unlock()

			p.log.Debugf("Writing piece to disk: #%d", piece.Index)
			if _, err = piece.Write(active.data); err != nil {
				p.log.Error(err)
				p.conn.Close()
				return
			}

			p.transfer.m.Lock()
			p.transfer.bitfield.Set(piece.Index)
			p.transfer.m.Unlock()
			p.cond.Broadcast()
		case cancelID:
		case portID:
		default:
			p.log.Debugf("Unknown message type: %d", id)
			p.log.Debugln("Discarding", length, "bytes...")
			io.CopyN(ioutil.Discard, p.conn, int64(length))
			p.log.Debug("Discarding finished.")
		}

		first = false
	}
}

func (p *Peer) handleHave(i uint32) {
	p.transfer.m.Lock()
	p.transfer.pieces[i].peers[p.id] = struct{}{}
	p.transfer.m.Unlock()
	p.cond.Broadcast()
}

func (p *Peer) SendBitfield() error {
	// Do not send a bitfield message if we don't have any pieces.
	if p.transfer.bitfield.Count() == 0 {
		return nil
	}
	return p.sendMessage(bitfieldID, p.transfer.bitfield.Bytes())
}

func (p *Peer) BeInterested() error {
	p.amInterestedM.Lock()
	defer p.amInterestedM.Unlock()
	if p.amInterested {
		return nil
	}
	p.amInterested = true
	return p.sendMessage(interestedID, nil)
}

func (p *Peer) BeNotInterested() error {
	p.amInterestedM.Lock()
	defer p.amInterestedM.Unlock()
	if !p.amInterested {
		return nil
	}
	p.amInterested = false
	return p.sendMessage(notInterestedID, nil)
}

func (p *Peer) Choke() error   { return p.sendMessage(chokeID, nil) }
func (p *Peer) Unchoke() error { return p.sendMessage(unchokeID, nil) }

func (p *Peer) Request(b *Block) error {
	req := requestMessage{b.Piece.Index, b.Begin, b.Length}
	buf := bytes.NewBuffer(make([]byte, 0, 12))
	binary.Write(buf, binary.BigEndian, &req)
	return p.sendMessage(requestID, buf.Bytes())
}

func (p *Peer) SendPiece(index, begin uint32, block []byte) error {
	msg := &pieceMessage{index, begin}
	buf := bytes.NewBuffer(make([]byte, 0, 8))
	binary.Write(buf, binary.BigEndian, msg)
	buf.Write(block)
	return p.sendMessage(pieceID, buf.Bytes())
}

func (p *Peer) sendMessage(id messageID, payload []byte) error {
	p.log.Debugf("Sending message of type: %q", id)
	buf := bufio.NewWriterSize(p.conn, 4+1+len(payload))
	var header = struct {
		Length uint32
		ID     messageID
	}{
		uint32(1 + len(payload)),
		id,
	}
	binary.Write(buf, binary.BigEndian, &header)
	buf.Write(payload)
	return buf.Flush()
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"io"
	"log"
	"net"
	"time"

	"github.com/zeebo/bencode"
)

const MAX_OUR_REQUESTS = 2
const MAX_PEER_REQUESTS = 10
const STANDARD_BLOCK_LENGTH = 16 * 1024

type peerMessage struct {
	peer    *peerState
	message []byte // nil means an error occurred
}

type peerState struct {
	address         string
	id              string
	writeChan       chan []byte
	writeChan2      chan []byte
	lastWriteTime   time.Time
	lastReadTime    time.Time
	have            *Bitset // What the peer has told us it has
	conn            net.Conn
	am_choking      bool // this client is choking the peer
	am_interested   bool // this client is interested in the peer
	peer_choking    bool // peer is choking this client
	peer_interested bool // peer is interested in this client
	peer_requests   map[uint64]bool
	our_requests    map[uint64]time.Time // What we requested, when we requested it

	// This field tells if the peer can send a bitfield or not
	can_receive_bitfield bool

	// Stores the bitfield they sent us but we can't verify yet (because
	// we don't have the torrent yet) and will commit when we can
	temporaryBitfield []byte

	theirExtensions map[string]int
}

func queueingWriter(in, out chan []byte) {
	queue := make(map[int][]byte)
	head, tail := 0, 0
L:
	for {
		if head == tail {
			select {
			case m, ok := <-in:
				if !ok {
					break L
				}
				queue[head] = m
				head++
			}
		} else {
			select {
			case m, ok := <-in:
				if !ok {
					break L
				}
				queue[head] = m
				head++
			case out <- queue[tail]:
				delete(queue, tail)
				tail++
			}
		}
	}
	// We throw away any messages waiting to be sent, including the
	// nil message that is automatically sent when the in channel is closed
	close(out)
}

func NewPeerState(conn net.Conn) *peerState {
	writeChan := make(chan []byte)
	writeChan2 := make(chan []byte)
	go queueingWriter(writeChan, writeChan2)
	return &peerState{writeChan: writeChan, writeChan2: writeChan2, conn: conn,
		am_choking: true, peer_choking: true,
		peer_requests:        make(map[uint64]bool, MAX_PEER_REQUESTS),
		our_requests:         make(map[uint64]time.Time, MAX_OUR_REQUESTS),
		can_receive_bitfield: true}
}

func (p *peerState) Close() {
	p.conn.Close()
	// No need to close p.writeChan. Further writes to p.conn will just fail.
}

func (p *peerState) AddRequest(index, begin, length uint32) {
	if !p.am_choking && len(p.peer_requests) < MAX_PEER_REQUESTS {
		offset := (uint64(index) << 32) | uint64(begin)
		p.peer_requests[offset] = true
	}
}

func (p *peerState) CancelRequest(index, begin, length uint32) {
	offset := (uint64(index) << 32) | uint64(begin)
	if _, ok := p.peer_requests[offset]; ok {
		delete(p.peer_requests, offset)
	}
}

func (p *peerState) RemoveRequest() (index, begin, length uint32, ok bool) {
	for k, _ := range p.peer_requests {
		index, begin = uint32(k>>32), uint32(k)
		length = STANDARD_BLOCK_LENGTH
		ok = true
		return
	}
	return
}

func (p *peerState) SetChoke(choke bool) {
	if choke != p.am_choking {
		p.am_choking = choke
		b := byte(UNCHOKE)
		if choke {
			b = CHOKE
			p.peer_requests = make(map[uint64]bool, MAX_PEER_REQUESTS)
		}
		p.sendOneCharMessage(b)
	}
}

func (p *peerState) SetInterested(interested bool) {
	if interested != p.am_interested {
		// log.Println("SetInterested", interested, p.address)
		p.am_interested = interested
		b := byte(NOT_INTERESTED)
		if interested {
			b = INTERESTED
		}
		p.sendOneCharMessage(b)
	}
}

func (p *peerState) SendBitfield(bs *Bitset) {
	msg := make([]byte, len(bs.Bytes())+1)
	msg[0] = BITFIELD
	copy(msg[1:], bs.Bytes())
	p.sendMessage(msg)
}

func (p *peerState) SendExtensions(supportedExtensions map[int]string,
	metadataSize int64) {

	handshake := ExtensionHandshake{
		M:            make(map[string]int, len(supportedExtensions)),
		V:            "Taipei-Torrent dev",
		MetadataSize: metadataSize,
	}

	for i, ext := range supportedExtensions {
		handshake.M[ext] = i
	}

	var buf bytes.Buffer
	err := bencode.NewEncoder(&buf).Encode(handshake)
	if err != nil {
		log.Println("Error when marshalling extension message")
		return
	}

	msg := make([]byte, 2+buf.Len())
	msg[0] = EXTENSION
	msg[1] = EXTENSION_HANDSHAKE
	copy(msg[2:], buf.Bytes())

	p.sendMessage(msg)
}

func (p *peerState) sendOneCharMessage(b byte) {
	// log.Println("ocm", b, p.address)
	p.sendMessage([]byte{b})
}

func (p *peerState) sendMessage(b []byte) {
	p.writeChan <- b
	p.lastWriteTime = time.Now()
}

func (p *peerState) keepAlive(now time.Time) {
	if now.Sub(p.lastWriteTime) >= 2*time.Minute {
		// log.Stderr("Sending keep alive", p)
		p.sendMessage([]byte{})
	}
}

// There's two goroutines per peer, one to read data from the peer, the other to
// send data to the peer.

func uint32ToBytes(buf []byte, n uint32) {
	buf[0] = byte(n >> 24)
	buf[1] = byte(n >> 16)
	buf[2] = byte(n >> 8)
	buf[3] = byte(n)
}

func writeNBOUint32(conn net.Conn, n uint32) (err error) {
	var buf []byte = make([]byte, 4)
	uint32ToBytes(buf, n)
	_, err = conn.Write(buf[0:])
	return
}

func bytesToUint32(buf []byte) uint32 {
	return (uint32(buf[0]) << 24) |
		(uint32(buf[1]) << 16) |
		(uint32(buf[2]) << 8) | uint32(buf[3])
}

func readNBOUint32(conn net.Conn) (n uint32, err error) {
	var buf [4]byte
	_, err = conn.Read(buf[0:])
	if err != nil {
		return
	}
	n = bytesToUint32(buf[0:])
	return
}

// This func is designed to be run as a goroutine. It
// listens for messages on a channel and sends them to a peer.

func (p *peerState) peerWriter(errorChan chan peerMessage) {
	// log.Println("Writing messages")
	for msg := range p.writeChan2 {
		// log.Println("Writing", uint32(len(msg)), p.conn.RemoteAddr())
		err := writeNBOUint32(p.conn, uint32(len(msg)))
		if err != nil {
			log.Println(err)
			break
		}
		_, err = p.conn.Write(msg)
		if err != nil {
			// log.Println("Failed to write a message", p.address, len(msg), msg, err)
			break
		}
	}
	// log.Println("peerWriter exiting")
	errorChan <- peerMessage{p, nil}
}

// This func is designed to be run as a goroutine. It
// listens for messages from the peer and forwards them to a channel.

func (p *peerState) peerReader(msgChan chan peerMessage) {
	// log.Println("Reading messages")
	for {
		var n uint32
		n, err := readNBOUint32(p.conn)
		if err != nil {
			break
		}
		if n > 130*1024 {
			// log.Println("Message size too large: ", n)
			break
		}

		var buf []byte
		if n == 0 {
			// keep-alive - we want an empty message
			buf = make([]byte, 1)
		} else {
			buf = make([]byte, n)
		}

		_, err = io.ReadFull(p.conn, buf)
		if err != nil {
			break
		}
		msgChan <- peerMessage{p, buf}
	}

	msgChan <- peerMessage{p, nil}
	// log.Println("peerReader exiting")
}

func (p *peerState) sendMetadataRequest(piece int) {
	metaMessage := MetadataMessage{
		MsgType: METADATA_REQUEST,
		Piece:   uint(piece),
	}

	p.sendExtensionMessage("ut_metadata", metaMessage)
}

func (p *peerState) sendExtensionMessage(typ string, data interface{}) {
	if _, ok := p.theirExtensions[typ]; !ok {
		// They don't understand this extension
		return
	}

	var payload bytes.Buffer
	err := bencode.NewEncoder(&payload).Encode(data)
	if err != nil {
		log.Printf("Couldn't marshal extension message: ", err)
	}

	msg := make([]byte, 2+payload.Len())
	msg[0] = EXTENSION
	msg[1] = byte(p.theirExtensions[typ])
	copy(msg[2:], payload.Bytes())

	p.sendMessage(msg)
}
<commit_msg>Really set keep-alive messages to length 0<commit_after>package main

import (
	"bytes"
	"io"
	"log"
	"net"
	"time"

	"github.com/zeebo/bencode"
)

const MAX_OUR_REQUESTS = 2
const MAX_PEER_REQUESTS = 10
const STANDARD_BLOCK_LENGTH = 16 * 1024

type peerMessage struct {
	peer    *peerState
	message []byte // nil means an error occurred
}

type peerState struct {
	address         string
	id              string
	writeChan       chan []byte
	writeChan2      chan []byte
	lastWriteTime   time.Time
	lastReadTime    time.Time
	have            *Bitset // What the peer has told us it has
	conn            net.Conn
	am_choking      bool // this client is choking the peer
	am_interested   bool // this client is interested in the peer
	peer_choking    bool // peer is choking this client
	peer_interested bool // peer is interested in this client
	peer_requests   map[uint64]bool
	our_requests    map[uint64]time.Time // What we requested, when we requested it

	// This field tells if the peer can send a bitfield or not
	can_receive_bitfield bool

	// Stores the bitfield they sent us but we can't verify yet (because
	// we don't have the torrent yet) and will commit when we can
	temporaryBitfield []byte

	theirExtensions map[string]int
}

func queueingWriter(in, out chan []byte) {
	queue := make(map[int][]byte)
	head, tail := 0, 0
L:
	for {
		if head == tail {
			select {
			case m, ok := <-in:
				if !ok {
					break L
				}
				queue[head] = m
				head++
			}
		} else {
			select {
			case m, ok := <-in:
				if !ok {
					break L
				}
				queue[head] = m
				head++
			case out <- queue[tail]:
				delete(queue, tail)
				tail++
			}
		}
	}
	// We throw away any messages waiting to be sent, including the
	// nil message that is automatically sent when the in channel is closed
	close(out)
}

func NewPeerState(conn net.Conn) *peerState {
	writeChan := make(chan []byte)
	writeChan2 := make(chan []byte)
	go queueingWriter(writeChan, writeChan2)
	return &peerState{writeChan: writeChan, writeChan2: writeChan2, conn: conn,
		am_choking: true, peer_choking: true,
		peer_requests:        make(map[uint64]bool, MAX_PEER_REQUESTS),
		our_requests:         make(map[uint64]time.Time, MAX_OUR_REQUESTS),
		can_receive_bitfield: true}
}

func (p *peerState) Close() {
	p.conn.Close()
	// No need to close p.writeChan. Further writes to p.conn will just fail.
}

func (p *peerState) AddRequest(index, begin, length uint32) {
	if !p.am_choking && len(p.peer_requests) < MAX_PEER_REQUESTS {
		offset := (uint64(index) << 32) | uint64(begin)
		p.peer_requests[offset] = true
	}
}

func (p *peerState) CancelRequest(index, begin, length uint32) {
	offset := (uint64(index) << 32) | uint64(begin)
	if _, ok := p.peer_requests[offset]; ok {
		delete(p.peer_requests, offset)
	}
}

func (p *peerState) RemoveRequest() (index, begin, length uint32, ok bool) {
	for k, _ := range p.peer_requests {
		index, begin = uint32(k>>32), uint32(k)
		length = STANDARD_BLOCK_LENGTH
		ok = true
		return
	}
	return
}

func (p *peerState) SetChoke(choke bool) {
	if choke != p.am_choking {
		p.am_choking = choke
		b := byte(UNCHOKE)
		if choke {
			b = CHOKE
			p.peer_requests = make(map[uint64]bool, MAX_PEER_REQUESTS)
		}
		p.sendOneCharMessage(b)
	}
}

func (p *peerState) SetInterested(interested bool) {
	if interested != p.am_interested {
		// log.Println("SetInterested", interested, p.address)
		p.am_interested = interested
		b := byte(NOT_INTERESTED)
		if interested {
			b = INTERESTED
		}
		p.sendOneCharMessage(b)
	}
}

func (p *peerState) SendBitfield(bs *Bitset) {
	msg := make([]byte, len(bs.Bytes())+1)
	msg[0] = BITFIELD
	copy(msg[1:], bs.Bytes())
	p.sendMessage(msg)
}

func (p *peerState) SendExtensions(supportedExtensions map[int]string,
	metadataSize int64) {

	handshake := ExtensionHandshake{
		M:            make(map[string]int, len(supportedExtensions)),
		V:            "Taipei-Torrent dev",
		MetadataSize: metadataSize,
	}

	for i, ext := range supportedExtensions {
		handshake.M[ext] = i
	}

	var buf bytes.Buffer
	err := bencode.NewEncoder(&buf).Encode(handshake)
	if err != nil {
		log.Println("Error when marshalling extension message")
		return
	}

	msg := make([]byte, 2+buf.Len())
	msg[0] = EXTENSION
	msg[1] = EXTENSION_HANDSHAKE
	copy(msg[2:], buf.Bytes())

	p.sendMessage(msg)
}

func (p *peerState) sendOneCharMessage(b byte) {
	// log.Println("ocm", b, p.address)
	p.sendMessage([]byte{b})
}

func (p *peerState) sendMessage(b []byte) {
	p.writeChan <- b
	p.lastWriteTime = time.Now()
}

func (p *peerState) keepAlive(now time.Time) {
	if now.Sub(p.lastWriteTime) >= 2*time.Minute {
		// log.Stderr("Sending keep alive", p)
		p.sendMessage([]byte{})
	}
}

// There's two goroutines per peer, one to read data from the peer, the other to
// send data to the peer.

func uint32ToBytes(buf []byte, n uint32) {
	buf[0] = byte(n >> 24)
	buf[1] = byte(n >> 16)
	buf[2] = byte(n >> 8)
	buf[3] = byte(n)
}

func writeNBOUint32(conn net.Conn, n uint32) (err error) {
	var buf []byte = make([]byte, 4)
	uint32ToBytes(buf, n)
	_, err = conn.Write(buf[0:])
	return
}

func bytesToUint32(buf []byte) uint32 {
	return (uint32(buf[0]) << 24) |
		(uint32(buf[1]) << 16) |
		(uint32(buf[2]) << 8) | uint32(buf[3])
}

func readNBOUint32(conn net.Conn) (n uint32, err error) {
	var buf [4]byte
	_, err = conn.Read(buf[0:])
	if err != nil {
		return
	}
	n = bytesToUint32(buf[0:])
	return
}

// This func is designed to be run as a goroutine. It
// listens for messages on a channel and sends them to a peer.

func (p *peerState) peerWriter(errorChan chan peerMessage) {
	// log.Println("Writing messages")
	for msg := range p.writeChan2 {
		// log.Println("Writing", uint32(len(msg)), p.conn.RemoteAddr())
		err := writeNBOUint32(p.conn, uint32(len(msg)))
		if err != nil {
			log.Println(err)
			break
		}
		_, err = p.conn.Write(msg)
		if err != nil {
			// log.Println("Failed to write a message", p.address, len(msg), msg, err)
			break
		}
	}
	// log.Println("peerWriter exiting")
	errorChan <- peerMessage{p, nil}
}

// This func is designed to be run as a goroutine. It
// listens for messages from the peer and forwards them to a channel.

func (p *peerState) peerReader(msgChan chan peerMessage) {
	// log.Println("Reading messages")
	for {
		var n uint32
		n, err := readNBOUint32(p.conn)
		if err != nil {
			break
		}
		if n > 130*1024 {
			// log.Println("Message size too large: ", n)
			break
		}

		buf := make([]byte, n)

		_, err = io.ReadFull(p.conn, buf)
		if err != nil {
			break
		}
		msgChan <- peerMessage{p, buf}
	}

	msgChan <- peerMessage{p, nil}
	// log.Println("peerReader exiting")
}

func (p *peerState) sendMetadataRequest(piece int) {
	metaMessage := MetadataMessage{
		MsgType: METADATA_REQUEST,
		Piece:   uint(piece),
	}

	p.sendExtensionMessage("ut_metadata", metaMessage)
}

func (p *peerState) sendExtensionMessage(typ string, data interface{}) {
	if _, ok := p.theirExtensions[typ]; !ok {
		// They don't understand this extension
		return
	}

	var payload bytes.Buffer
	err := bencode.NewEncoder(&payload).Encode(data)
	if err != nil {
		log.Printf("Couldn't marshal extension message: ", err)
	}

	msg := make([]byte, 2+payload.Len())
	msg[0] = EXTENSION
	msg[1] = byte(p.theirExtensions[typ])
	copy(msg[2:], payload.Bytes())

	p.sendMessage(msg)
}
<|endoftext|>