<commit_before>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"io"
	"sort"
	"sync"

	"github.com/keybase/client/go/logger"
	"github.com/pkg/errors"
	"golang.org/x/net/context"
)

const (
	defaultIndirectPointerPrefetchCount int = 20
	fileIndirectBlockPrefetchPriority   int = -100
	dirEntryPrefetchPriority            int = -200
	updatePointerPrefetchPriority       int = 0
	defaultPrefetchPriority             int = -1024
)

type prefetcherConfig interface {
	dataVersioner
	logMaker
	blockCacher
}

type prefetchRequest struct {
	priority int
	kmd      KeyMetadata
	ptr      BlockPointer
	block    Block
}

// blockRetriever specifies a method for retrieving blocks asynchronously.
type blockRetriever interface {
	Request(ctx context.Context, priority int, kmd KeyMetadata, ptr BlockPointer, block Block, lifetime BlockCacheLifetime) <-chan error
}

type blockPrefetcher struct {
	config prefetcherConfig
	log    logger.Logger
	// blockRetriever to retrieve blocks from the server
	retriever blockRetriever
	// channel to synchronize prefetch requests with the prefetcher shutdown
	progressCh chan prefetchRequest
	// channel that is idempotently closed when a shutdown occurs
	shutdownCh chan struct{}
	// channel that is closed when a shutdown completes and all pending
	// prefetch requests are complete
	doneCh chan struct{}
}

var _ Prefetcher = (*blockPrefetcher)(nil)

func newBlockPrefetcher(retriever blockRetriever, config prefetcherConfig) *blockPrefetcher {
	p := &blockPrefetcher{
		config:     config,
		retriever:  retriever,
		progressCh: make(chan prefetchRequest),
		shutdownCh: make(chan struct{}),
		doneCh:     make(chan struct{}),
	}
	if config != nil {
		p.log = config.MakeLogger("PRE")
	} else {
		p.log = logger.NewNull()
	}
	if retriever == nil {
		// If we pass in a nil retriever, this prefetcher shouldn't do
		// anything. Treat it as already shut down.
		p.Shutdown()
		close(p.doneCh)
	} else {
		go p.run()
	}
	return p
}

func (p *blockPrefetcher) run() {
	var wg sync.WaitGroup
	defer func() {
		wg.Wait()
		close(p.doneCh)
	}()
	for {
		select {
		case req := <-p.progressCh:
			ctx, cancel := context.WithCancel(context.TODO())
			errCh := p.retriever.Request(ctx, req.priority, req.kmd, req.ptr, req.block, TransientEntry)
			wg.Add(1)
			go func() {
				defer wg.Done()
				defer cancel()
				select {
				case err := <-errCh:
					if err != nil {
						p.log.CDebugf(ctx, "Done prefetch for block %s. Error: %+v", req.ptr.ID, err)
					}
				case <-p.shutdownCh:
					// Cancel but still wait so p.doneCh accurately represents
					// whether we still have requests pending.
					cancel()
					<-errCh
				}
			}()
		case <-p.shutdownCh:
			return
		}
	}
}

func (p *blockPrefetcher) request(priority int, kmd KeyMetadata, ptr BlockPointer, block Block, entryName string) error {
	if _, err := p.config.BlockCache().Get(ptr); err == nil {
		return nil
	}
	if err := checkDataVersion(p.config, path{}, ptr); err != nil {
		return err
	}
	select {
	case p.progressCh <- prefetchRequest{priority, kmd, ptr, block}:
		return nil
	case <-p.shutdownCh:
		return errors.Wrapf(io.EOF, "Skipping prefetch for block %v since the prefetcher is shutdown", ptr.ID)
	}
}

func (p *blockPrefetcher) prefetchIndirectFileBlock(b *FileBlock, kmd KeyMetadata) {
	// Prefetch the first <n> indirect block pointers.
	// TODO: do something smart with subsequent blocks.
	numIPtrs := len(b.IPtrs)
	if numIPtrs > defaultIndirectPointerPrefetchCount {
		numIPtrs = defaultIndirectPointerPrefetchCount
	}
	p.log.CDebugf(context.TODO(), "Prefetching pointers for indirect file block. Num pointers to prefetch: %d", numIPtrs)
	for _, ptr := range b.IPtrs[:numIPtrs] {
		p.request(fileIndirectBlockPrefetchPriority, kmd,
			ptr.BlockPointer, b.NewEmpty(), "")
	}
}

func (p *blockPrefetcher) prefetchIndirectDirBlock(b *DirBlock, kmd KeyMetadata) {
	// Prefetch the first <n> indirect block pointers.
	numIPtrs := len(b.IPtrs)
	if numIPtrs > defaultIndirectPointerPrefetchCount {
		numIPtrs = defaultIndirectPointerPrefetchCount
	}
	p.log.CDebugf(context.TODO(), "Prefetching pointers for indirect dir block. Num pointers to prefetch: %d", numIPtrs)
	for _, ptr := range b.IPtrs[:numIPtrs] {
		_ = p.request(fileIndirectBlockPrefetchPriority, kmd,
			ptr.BlockPointer, b.NewEmpty(), "")
	}
}

func (p *blockPrefetcher) prefetchDirectDirBlock(b *DirBlock, kmd KeyMetadata) {
	p.log.CDebugf(context.TODO(), "Prefetching entries for directory block. Num entries: %d", len(b.Children))
	// Prefetch all DirEntry root blocks.
	dirEntries := dirEntriesBySizeAsc{dirEntryMapToDirEntries(b.Children)}
	sort.Sort(dirEntries)
	for i, entry := range dirEntries.dirEntries {
		// Prioritize small files
		priority := dirEntryPrefetchPriority - i
		var block Block
		switch entry.Type {
		case Dir:
			block = &DirBlock{}
		case File:
			block = &FileBlock{}
		case Exec:
			block = &FileBlock{}
		default:
			p.log.CDebugf(context.TODO(), "Skipping prefetch for entry of unknown type %d", entry.Type)
			continue
		}
		p.request(priority, kmd, entry.BlockPointer, block, entry.entryName)
	}
}

// PrefetchBlock implements the Prefetcher interface for blockPrefetcher.
func (p *blockPrefetcher) PrefetchBlock(
	block Block, ptr BlockPointer, kmd KeyMetadata, priority int) error {
	p.log.CDebugf(context.TODO(), "Prefetching block by request from upstream component. Priority: %d", priority)
	return p.request(priority, kmd, ptr, block, "")
}

// PrefetchAfterBlockRetrieved implements the Prefetcher interface for
// blockPrefetcher.
func (p *blockPrefetcher) PrefetchAfterBlockRetrieved(
	b Block, ptr BlockPointer, kmd KeyMetadata, priority int,
	lifetime BlockCacheLifetime, hasPrefetched bool) {
	if hasPrefetched {
		// We discard the error because there's nothing we can do about it.
		_ = p.config.BlockCache().PutWithPrefetch(ptr, kmd.TlfID(), b,
			lifetime, true)
		return
	}
	if priority < defaultOnDemandRequestPriority {
		// Only on-demand or higher priority requests can trigger prefetches.
		// We discard the error because there's nothing we can do about it.
		_ = p.config.BlockCache().PutWithPrefetch(ptr, kmd.TlfID(), b,
			lifetime, false)
		return
	}
	// We must let the cache know at this point that we've prefetched.
	// 1) To prevent any other Gets from prefetching.
	// 2) To prevent prefetching if a cache Put fails, since prefetching if
	//	  only useful when combined with the cache.
	err := p.config.BlockCache().PutWithPrefetch(ptr, kmd.TlfID(), b,
		lifetime, true)
	if _, isCacheFullError := err.(cachePutCacheFullError); isCacheFullError {
		p.log.CDebugf(context.TODO(), "Skipping prefetch because the cache is full")
		return
	}
	switch b := b.(type) {
	case *FileBlock:
		if b.IsInd {
			p.prefetchIndirectFileBlock(b, kmd)
		}
	case *DirBlock:
		if b.IsInd {
			p.prefetchIndirectDirBlock(b, kmd)
		} else {
			p.prefetchDirectDirBlock(b, kmd)
		}
	default:
		p.log.CDebugf(context.TODO(), "Skipping prefetch for entry of unknown type %T", b)
	}
}

// Shutdown implements the Prefetcher interface for blockPrefetcher.
func (p *blockPrefetcher) Shutdown() <-chan struct{} {
	select {
	case <-p.shutdownCh:
	default:
		close(p.shutdownCh)
	}
	return p.doneCh
}
<commit_msg>prefetcher: Remove overly verbose log message<commit_after>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"io"
	"sort"
	"sync"

	"github.com/keybase/client/go/logger"
	"github.com/pkg/errors"
	"golang.org/x/net/context"
)

const (
	defaultIndirectPointerPrefetchCount int = 20
	fileIndirectBlockPrefetchPriority   int = -100
	dirEntryPrefetchPriority            int = -200
	updatePointerPrefetchPriority       int = 0
	defaultPrefetchPriority             int = -1024
)

type prefetcherConfig interface {
	dataVersioner
	logMaker
	blockCacher
}

type prefetchRequest struct {
	priority int
	kmd      KeyMetadata
	ptr      BlockPointer
	block    Block
}

// blockRetriever specifies a method for retrieving blocks asynchronously.
type blockRetriever interface {
	Request(ctx context.Context, priority int, kmd KeyMetadata, ptr BlockPointer, block Block, lifetime BlockCacheLifetime) <-chan error
}

type blockPrefetcher struct {
	config prefetcherConfig
	log    logger.Logger
	// blockRetriever to retrieve blocks from the server
	retriever blockRetriever
	// channel to synchronize prefetch requests with the prefetcher shutdown
	progressCh chan prefetchRequest
	// channel that is idempotently closed when a shutdown occurs
	shutdownCh chan struct{}
	// channel that is closed when a shutdown completes and all pending
	// prefetch requests are complete
	doneCh chan struct{}
}

var _ Prefetcher = (*blockPrefetcher)(nil)

func newBlockPrefetcher(retriever blockRetriever, config prefetcherConfig) *blockPrefetcher {
	p := &blockPrefetcher{
		config:     config,
		retriever:  retriever,
		progressCh: make(chan prefetchRequest),
		shutdownCh: make(chan struct{}),
		doneCh:     make(chan struct{}),
	}
	if config != nil {
		p.log = config.MakeLogger("PRE")
	} else {
		p.log = logger.NewNull()
	}
	if retriever == nil {
		// If we pass in a nil retriever, this prefetcher shouldn't do
		// anything. Treat it as already shut down.
		p.Shutdown()
		close(p.doneCh)
	} else {
		go p.run()
	}
	return p
}

func (p *blockPrefetcher) run() {
	var wg sync.WaitGroup
	defer func() {
		wg.Wait()
		close(p.doneCh)
	}()
	for {
		select {
		case req := <-p.progressCh:
			ctx, cancel := context.WithCancel(context.TODO())
			errCh := p.retriever.Request(ctx, req.priority, req.kmd, req.ptr, req.block, TransientEntry)
			wg.Add(1)
			go func() {
				defer wg.Done()
				defer cancel()
				select {
				case err := <-errCh:
					if err != nil {
						p.log.CDebugf(ctx, "Done prefetch for block %s. Error: %+v", req.ptr.ID, err)
					}
				case <-p.shutdownCh:
					// Cancel but still wait so p.doneCh accurately represents
					// whether we still have requests pending.
					cancel()
					<-errCh
				}
			}()
		case <-p.shutdownCh:
			return
		}
	}
}

func (p *blockPrefetcher) request(priority int, kmd KeyMetadata, ptr BlockPointer, block Block, entryName string) error {
	if _, err := p.config.BlockCache().Get(ptr); err == nil {
		return nil
	}
	if err := checkDataVersion(p.config, path{}, ptr); err != nil {
		return err
	}
	select {
	case p.progressCh <- prefetchRequest{priority, kmd, ptr, block}:
		return nil
	case <-p.shutdownCh:
		return errors.Wrapf(io.EOF, "Skipping prefetch for block %v since the prefetcher is shutdown", ptr.ID)
	}
}

func (p *blockPrefetcher) prefetchIndirectFileBlock(b *FileBlock, kmd KeyMetadata) {
	// Prefetch the first <n> indirect block pointers.
	// TODO: do something smart with subsequent blocks.
	numIPtrs := len(b.IPtrs)
	if numIPtrs > defaultIndirectPointerPrefetchCount {
		numIPtrs = defaultIndirectPointerPrefetchCount
	}
	p.log.CDebugf(context.TODO(), "Prefetching pointers for indirect file block. Num pointers to prefetch: %d", numIPtrs)
	for _, ptr := range b.IPtrs[:numIPtrs] {
		p.request(fileIndirectBlockPrefetchPriority, kmd,
			ptr.BlockPointer, b.NewEmpty(), "")
	}
}

func (p *blockPrefetcher) prefetchIndirectDirBlock(b *DirBlock, kmd KeyMetadata) {
	// Prefetch the first <n> indirect block pointers.
	numIPtrs := len(b.IPtrs)
	if numIPtrs > defaultIndirectPointerPrefetchCount {
		numIPtrs = defaultIndirectPointerPrefetchCount
	}
	p.log.CDebugf(context.TODO(), "Prefetching pointers for indirect dir block. Num pointers to prefetch: %d", numIPtrs)
	for _, ptr := range b.IPtrs[:numIPtrs] {
		_ = p.request(fileIndirectBlockPrefetchPriority, kmd,
			ptr.BlockPointer, b.NewEmpty(), "")
	}
}

func (p *blockPrefetcher) prefetchDirectDirBlock(b *DirBlock, kmd KeyMetadata) {
	p.log.CDebugf(context.TODO(), "Prefetching entries for directory block. Num entries: %d", len(b.Children))
	// Prefetch all DirEntry root blocks.
	dirEntries := dirEntriesBySizeAsc{dirEntryMapToDirEntries(b.Children)}
	sort.Sort(dirEntries)
	for i, entry := range dirEntries.dirEntries {
		// Prioritize small files
		priority := dirEntryPrefetchPriority - i
		var block Block
		switch entry.Type {
		case Dir:
			block = &DirBlock{}
		case File:
			block = &FileBlock{}
		case Exec:
			block = &FileBlock{}
		default:
			p.log.CDebugf(context.TODO(), "Skipping prefetch for entry of unknown type %d", entry.Type)
			continue
		}
		p.request(priority, kmd, entry.BlockPointer, block, entry.entryName)
	}
}

// PrefetchBlock implements the Prefetcher interface for blockPrefetcher.
func (p *blockPrefetcher) PrefetchBlock(
	block Block, ptr BlockPointer, kmd KeyMetadata, priority int) error {
	p.log.CDebugf(context.TODO(), "Prefetching block by request from upstream component. Priority: %d", priority)
	return p.request(priority, kmd, ptr, block, "")
}

// PrefetchAfterBlockRetrieved implements the Prefetcher interface for
// blockPrefetcher.
func (p *blockPrefetcher) PrefetchAfterBlockRetrieved(
	b Block, ptr BlockPointer, kmd KeyMetadata, priority int,
	lifetime BlockCacheLifetime, hasPrefetched bool) {
	if hasPrefetched {
		// We discard the error because there's nothing we can do about it.
		_ = p.config.BlockCache().PutWithPrefetch(ptr, kmd.TlfID(), b,
			lifetime, true)
		return
	}
	if priority < defaultOnDemandRequestPriority {
		// Only on-demand or higher priority requests can trigger prefetches.
		// We discard the error because there's nothing we can do about it.
		_ = p.config.BlockCache().PutWithPrefetch(ptr, kmd.TlfID(), b,
			lifetime, false)
		return
	}
	// We must let the cache know at this point that we've prefetched.
	// 1) To prevent any other Gets from prefetching.
	// 2) To prevent prefetching if a cache Put fails, since prefetching if
	//	  only useful when combined with the cache.
	err := p.config.BlockCache().PutWithPrefetch(ptr, kmd.TlfID(), b,
		lifetime, true)
	if _, isCacheFullError := err.(cachePutCacheFullError); isCacheFullError {
		p.log.CDebugf(context.TODO(), "Skipping prefetch because the cache is full")
		return
	}
	switch b := b.(type) {
	case *FileBlock:
		if b.IsInd {
			p.prefetchIndirectFileBlock(b, kmd)
		}
	case *DirBlock:
		if b.IsInd {
			p.prefetchIndirectDirBlock(b, kmd)
		} else {
			p.prefetchDirectDirBlock(b, kmd)
		}
	default:
		// Skipping prefetch for block of unknown type (likely CommonBlock)
	}
}

// Shutdown implements the Prefetcher interface for blockPrefetcher.
func (p *blockPrefetcher) Shutdown() <-chan struct{} {
	select {
	case <-p.shutdownCh:
	default:
		close(p.shutdownCh)
	}
	return p.doneCh
}
<|endoftext|>


<commit_before>package main

//go:generate esc -o templates.go ../../templates

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"text/tabwriter"

	"github.com/fsnotify/fsnotify"
	"github.com/subosito/snowboard"
	"github.com/subosito/snowboard/engines/drafter"
	"github.com/subosito/snowboard/engines/drafterc"
)

const versionStr = "v0.3.3"

var (
	version  = flag.Bool("v", false, "Display version information")
	input    = flag.String("i", "", "API Blueprint file")
	output   = flag.String("o", "index.html", "HTML output file")
	watch    = flag.Bool("w", false, "Watch input (and template, if any) file for changes")
	serve    = flag.Bool("s", false, "Serve HTML via 0.0.0.0:8088")
	tplFile  = flag.String("t", "alpha", "Custom template for documentation")
	engineF  = flag.String("e", "cgo", "Use different engine. Supported engines: cgo, cli")
	validate = flag.Bool("l", false, "Validate input only")
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage:\n  snowboard [OPTIONS]\n\nOptions:\n")
		flag.PrintDefaults()
		os.Exit(0)
	}

	flag.Parse()

	engine := parserEngine()

	if *version {
		fmt.Printf("Snowboard version: %s\n", versionStr)
		fmt.Println("Engine:")

		for name, v := range engine.Version() {
			fmt.Printf("  %s version: %s\n", name, v)
		}

		os.Exit(0)
	}

	if *input == "" {
		flag.Usage()
	}

	if *validate {
		cEngine := checkerEngine()

		b, err := readFile(*input)
		checkErr(err)

		bf := bytes.NewReader(b)

		out, err := snowboard.Validate(bf, cEngine)
		if err == nil && out == nil {
			fmt.Fprintf(os.Stdout, "OK\n")
			os.Exit(0)
		}

		if err != nil {
			checkErr(err)
		}

		s := "--------"
		w := tabwriter.NewWriter(os.Stdout, 8, 0, 0, ' ', tabwriter.Debug)
		fmt.Fprintln(w, "Row\tCol\tDescription")
		fmt.Fprintf(w, "%s\t%s\t%s\n", s, s, strings.Repeat(s, 12))

		for _, n := range out.Annotations {
			for _, m := range n.SourceMaps {
				fmt.Fprintf(w, "%d\t%d\t%s\n", m.Row, m.Col, n.Description)
			}
		}

		w.Flush()
		os.Exit(1)
	}

	if *watch {
		watcher, err := fsnotify.NewWatcher()
		checkErr(err)
		defer watcher.Close()

		done := make(chan bool)
		go func() {
			for {
				select {
				case event := <-watcher.Events:
					if event.Op&fsnotify.Write == fsnotify.Write {
						renderHTML(engine)
					}
				case err := <-watcher.Errors:
					log.Println("Error:", err)
				}
			}
		}()

		err = watcher.Add(*input)
		checkErr(err)

		_, err = os.Stat(*tplFile)
		if err == nil {
			err = watcher.Add(*tplFile)
			checkErr(err)
		}

		renderHTML(engine)
		serveHTML()

		<-done
	} else {
		renderHTML(engine)
		serveHTML()
	}
}

func readFile(fn string) ([]byte, error) {
	info, err := os.Stat(fn)
	if err != nil {
		return nil, errors.New("File is not exist")
	}

	if info.IsDir() {
		return nil, errors.New("File is a directory")
	}

	return ioutil.ReadFile(fn)
}

func readTemplate(fn string) ([]byte, error) {
	tf, err := readFile(fn)
	if err == nil {
		return tf, nil
	}

	fs := FS(false)
	ff, err := fs.Open("/templates/" + fn + ".html")
	if err != nil {
		return nil, err
	}

	defer ff.Close()
	return ioutil.ReadAll(ff)
}

func checkErr(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %s\n", err)
		os.Exit(1)
	}
}

func renderHTML(engine snowboard.Parser) {
	bp, err := snowboard.Load(*input, engine)
	logErr(err)

	of, err := os.Create(*output)
	logErr(err)
	defer of.Close()

	tf, err := readTemplate(*tplFile)
	logErr(err)

	err = snowboard.HTML(string(tf), of, bp)
	logErr(err)
	log.Println("HTML has been generated!")
}

func logErr(err error) {
	if err != nil {
		log.Fatalln("Error: ", err)
	}
}

func serveHTML() {
	if !*serve {
		return
	}

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, *output)
	})

	err := http.ListenAndServe(":8088", nil)
	logErr(err)
}

func parserEngine() snowboard.Parser {
	switch *engineF {
	case "cli":
		return drafterc.Engine{}
	}

	return drafter.Engine{}
}

func checkerEngine() snowboard.Parser {
	return drafter.Engine{}
}
<commit_msg>allow http listen addr and port customization; always watch changes when running as server<commit_after>package main

//go:generate esc -o templates.go ../../templates

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"text/tabwriter"

	"github.com/fsnotify/fsnotify"
	"github.com/subosito/snowboard"
	"github.com/subosito/snowboard/engines/drafter"
	"github.com/subosito/snowboard/engines/drafterc"
)

const versionStr = "v0.3.3"

var (
	version  = flag.Bool("v", false, "Display version information")
	input    = flag.String("i", "", "API Blueprint file")
	output   = flag.String("o", "index.html", "HTML output file")
	serve    = flag.String("s", "127.0.0.1:8088", "Serve HTML via HTTP server")
	tplFile  = flag.String("t", "alpha", "Custom template for documentation")
	engineF  = flag.String("e", "cgo", "Use different engine. Supported engines: cgo, cli")
	validate = flag.Bool("l", false, "Validate input only")
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage:\n  snowboard [OPTIONS]\n\nOptions:\n")
		flag.PrintDefaults()
		os.Exit(0)
	}

	flag.Parse()

	engine := parserEngine()

	if *version {
		fmt.Printf("Snowboard version: %s\n", versionStr)
		fmt.Println("Engine:")

		for name, v := range engine.Version() {
			fmt.Printf("  %s version: %s\n", name, v)
		}

		os.Exit(0)
	}

	if *input == "" {
		flag.Usage()
	}

	if *validate {
		cEngine := checkerEngine()

		b, err := readFile(*input)
		checkErr(err)

		bf := bytes.NewReader(b)

		out, err := snowboard.Validate(bf, cEngine)
		if err == nil && out == nil {
			fmt.Fprintf(os.Stdout, "OK\n")
			os.Exit(0)
		}

		if err != nil {
			checkErr(err)
		}

		s := "--------"
		w := tabwriter.NewWriter(os.Stdout, 8, 0, 0, ' ', tabwriter.Debug)
		fmt.Fprintln(w, "Row\tCol\tDescription")
		fmt.Fprintf(w, "%s\t%s\t%s\n", s, s, strings.Repeat(s, 12))

		for _, n := range out.Annotations {
			for _, m := range n.SourceMaps {
				fmt.Fprintf(w, "%d\t%d\t%s\n", m.Row, m.Col, n.Description)
			}
		}

		w.Flush()
		os.Exit(1)
	}

	if *serve != "" {
		watcher, err := fsnotify.NewWatcher()
		checkErr(err)
		defer watcher.Close()

		done := make(chan bool)
		go func() {
			for {
				select {
				case event := <-watcher.Events:
					if event.Op&fsnotify.Write == fsnotify.Write {
						renderHTML(engine)
					}
				case err := <-watcher.Errors:
					log.Println("Error:", err)
				}
			}
		}()

		err = watcher.Add(*input)
		checkErr(err)

		_, err = os.Stat(*tplFile)
		if err == nil {
			err = watcher.Add(*tplFile)
			checkErr(err)
		}

		renderHTML(engine)
		serveHTML()

		<-done
	} else {
		renderHTML(engine)
		serveHTML()
	}
}

func readFile(fn string) ([]byte, error) {
	info, err := os.Stat(fn)
	if err != nil {
		return nil, errors.New("File is not exist")
	}

	if info.IsDir() {
		return nil, errors.New("File is a directory")
	}

	return ioutil.ReadFile(fn)
}

func readTemplate(fn string) ([]byte, error) {
	tf, err := readFile(fn)
	if err == nil {
		return tf, nil
	}

	fs := FS(false)
	ff, err := fs.Open("/templates/" + fn + ".html")
	if err != nil {
		return nil, err
	}

	defer ff.Close()
	return ioutil.ReadAll(ff)
}

func checkErr(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %s\n", err)
		os.Exit(1)
	}
}

func renderHTML(engine snowboard.Parser) {
	bp, err := snowboard.Load(*input, engine)
	logErr(err)

	of, err := os.Create(*output)
	logErr(err)
	defer of.Close()

	tf, err := readTemplate(*tplFile)
	logErr(err)

	err = snowboard.HTML(string(tf), of, bp)
	logErr(err)
	log.Println("HTML has been generated!")
}

func logErr(err error) {
	if err != nil {
		log.Fatalln("Error: ", err)
	}
}

func serveHTML() {
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, *output)
	})

	err := http.ListenAndServe(*serve, nil)
	logErr(err)
}

func parserEngine() snowboard.Parser {
	switch *engineF {
	case "cli":
		return drafterc.Engine{}
	}

	return drafter.Engine{}
}

func checkerEngine() snowboard.Parser {
	return drafter.Engine{}
}
<|endoftext|>


<commit_before>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// See https://code.google.com/p/go/source/browse/CONTRIBUTORS
// Licensed under the same terms as Go itself:
// https://code.google.com/p/go/source/browse/LICENSE

// Package http2 implements the HTTP/2 protocol.
//
// This is a work in progress. This package is low-level and intended
// to be used directly by very few people. Most users will use it
// indirectly through integration with the net/http package. See
// ConfigureServer. That ConfigureServer call will likely be automatic
// or available via an empty import in the future.
//
// This package currently targets draft-14. See http://http2.github.io/
package http2

import (
	"bytes"
	"crypto/tls"
	"io"
	"log"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strconv"
	"strings"

	"github.com/bradfitz/http2/hpack"
)

const (
	// ClientPreface is the string that must be sent by new
	// connections from clients.
	ClientPreface = "PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"
)

var (
	clientPreface = []byte(ClientPreface)
)

const (
	npnProto = "h2-14"

	// http://http2.github.io/http2-spec/#SettingValues
	initialHeaderTableSize = 4096
)

// Server is an HTTP2 server.
type Server struct {
	// MaxStreams optionally ...
	MaxStreams int
}

func (srv *Server) handleConn(hs *http.Server, c *tls.Conn, h http.Handler) {
	sc := &serverConn{
		hs:             hs,
		conn:           c,
		handler:        h,
		framer:         NewFramer(c, c),
		streams:        make(map[uint32]*stream),
		canonHeader:    make(map[string]string),
		readFrameCh:    make(chan frameAndProcessed),
		readFrameErrCh: make(chan error, 1),
		doneServing:    make(chan struct{}),
	}
	sc.hpackDecoder = hpack.NewDecoder(initialHeaderTableSize, sc.onNewHeaderField)
	sc.serve()
}

// frameAndProcessed coordinates the readFrames and serve goroutines, since
// the Framer interface only permits the most recently-read Frame from being
// accessed. The serve goroutine sends on processed to signal to the readFrames
// goroutine that another frame may be read.
type frameAndProcessed struct {
	f         Frame
	processed chan struct{}
}

type serverConn struct {
	hs             *http.Server
	conn           *tls.Conn
	handler        http.Handler
	framer         *Framer
	doneServing    chan struct{}          // closed when serverConn.serve ends
	readFrameCh    chan frameAndProcessed // written by serverConn.readFrames
	readFrameErrCh chan error

	maxStreamID uint32 // max ever seen
	streams     map[uint32]*stream

	// State related to parsing current headers:
	hpackDecoder *hpack.Decoder
	header       http.Header
	canonHeader  map[string]string // http2-lower-case -> Go-Canonical-Case

	method, path, scheme, authority string

	// curHeaderStreamID is non-zero if we're in the middle
	// of parsing headers that span multiple frames.
	curHeaderStreamID uint32
	curStream         *stream
}

type streamState int

const (
	stateIdle streamState = iota
	stateOpen
	stateHalfClosedLocal
	stateHalfClosedRemote
	stateResvLocal
	stateResvRemote
	stateClosed
)

type stream struct {
	id    uint32
	state streamState // owned by serverConn's processing loop
}

func (sc *serverConn) state(streamID uint32) streamState {
	// http://http2.github.io/http2-spec/#rfc.section.5.1
	if st, ok := sc.streams[streamID]; ok {
		return st.state
	}
	// "The first use of a new stream identifier implicitly closes all
	// streams in the "idle" state that might have been initiated by
	// that peer with a lower-valued stream identifier. For example, if
	// a client sends a HEADERS frame on stream 7 without ever sending a
	// frame on stream 5, then stream 5 transitions to the "closed"
	// state when the first frame for stream 7 is sent or received."
	if streamID <= sc.maxStreamID {
		return stateClosed
	}
	return stateIdle
}

func (sc *serverConn) logf(format string, args ...interface{}) {
	if lg := sc.hs.ErrorLog; lg != nil {
		lg.Printf(format, args...)
	} else {
		log.Printf(format, args...)
	}
}

func (sc *serverConn) onNewHeaderField(f hpack.HeaderField) {
	log.Printf("Header field: +%v", f)
	if strings.HasPrefix(f.Name, ":") {
		switch f.Name {
		case ":method":
			sc.method = f.Value
		case ":path":
			sc.path = f.Value
		case ":scheme":
			sc.scheme = f.Value
		case ":authority":
			sc.authority = f.Value
		default:
			log.Printf("Ignoring unknown pseudo-header %q", f.Name)
		}
		return
	}
	sc.header.Add(sc.canonicalHeader(f.Name), f.Value)
}

func (sc *serverConn) canonicalHeader(v string) string {
	// TODO: use a sync.Pool instead of putting the cache on *serverConn?
	cv, ok := sc.canonHeader[v]
	if !ok {
		cv = http.CanonicalHeaderKey(v)
		sc.canonHeader[v] = cv
	}
	return cv
}

func (sc *serverConn) readFrames() {
	processed := make(chan struct{}, 1)
	for {
		f, err := sc.framer.ReadFrame()
		if err != nil {
			close(sc.readFrameCh)
			sc.readFrameErrCh <- err
			return
		}
		sc.readFrameCh <- frameAndProcessed{f, processed}
		<-processed
	}
}

func (sc *serverConn) serve() {
	defer sc.conn.Close()
	defer close(sc.doneServing)

	log.Printf("HTTP/2 connection from %v on %p", sc.conn.RemoteAddr(), sc.hs)

	// Read the client preface
	buf := make([]byte, len(ClientPreface))
	// TODO: timeout reading from the client
	if _, err := io.ReadFull(sc.conn, buf); err != nil {
		sc.logf("error reading client preface: %v", err)
		return
	}
	if !bytes.Equal(buf, clientPreface) {
		sc.logf("bogus greeting from client: %q", buf)
		return
	}
	log.Printf("client %v said hello", sc.conn.RemoteAddr())

	go sc.readFrames()

	for {
		select {
		case fp, ok := <-sc.readFrameCh:
			if !ok {
				err := <-sc.readFrameErrCh
				if err != io.EOF {
					sc.logf("client %s stopped sending frames: %v", sc.conn.RemoteAddr(), err)
				}
				return
			}
			f := fp.f
			log.Printf("got %v: %#v", f.Header(), f)
			err := sc.processFrame(f)
			fp.processed <- struct{}{} // let readFrames proceed
			if h2e, ok := err.(Error); ok {
				if h2e.IsConnectionError() {
					sc.logf("Disconnection; connection error: %v", err)
					return
				}
				// TODO: stream errors, etc
			}
			if err != nil {
				sc.logf("Disconnection due to other error: %v", err)
				return
			}
		}
	}
}

func (sc *serverConn) processFrame(f Frame) error {
	if s := sc.curHeaderStreamID; s != 0 {
		if cf, ok := f.(*ContinuationFrame); !ok {
			return ConnectionError(ErrCodeProtocol)
		} else if cf.Header().StreamID != s {
			return ConnectionError(ErrCodeProtocol)
		}
	}

	switch f := f.(type) {
	case *SettingsFrame:
		return sc.processSettings(f)
	case *HeadersFrame:
		return sc.processHeaders(f)
	case *ContinuationFrame:
		return sc.processContinuation(f)
	default:
		log.Printf("Ignoring unknown %v", f.Header)
		return nil
	}
}

func (sc *serverConn) processSettings(f *SettingsFrame) error {
	f.ForeachSetting(func(s Setting) {
		log.Printf("  setting %s = %v", s.ID, s.Val)
	})
	return nil
}

func (sc *serverConn) processHeaders(f *HeadersFrame) error {
	id := f.Header().StreamID

	// http://http2.github.io/http2-spec/#rfc.section.5.1.1
	if id%2 != 1 || id <= sc.maxStreamID {
		// Streams initiated by a client MUST use odd-numbered
		// stream identifiers. [...] The identifier of a newly
		// established stream MUST be numerically greater than all
		// streams that the initiating endpoint has opened or
		// reserved. [...]  An endpoint that receives an unexpected
		// stream identifier MUST respond with a connection error
		// (Section 5.4.1) of type PROTOCOL_ERROR.
		return ConnectionError(ErrCodeProtocol)
	}
	if id > sc.maxStreamID {
		sc.maxStreamID = id
	}

	st := &stream{
		id:    id,
		state: stateOpen,
	}
	sc.streams[id] = st
	sc.header = make(http.Header)
	sc.curHeaderStreamID = id
	sc.curStream = st
	return sc.processHeaderBlockFragment(f.HeaderBlockFragment(), f.HeadersEnded())
}

func (sc *serverConn) processContinuation(f *ContinuationFrame) error {
	return sc.processHeaderBlockFragment(f.HeaderBlockFragment(), f.HeadersEnded())
}

func (sc *serverConn) processHeaderBlockFragment(frag []byte, end bool) error {
	if _, err := sc.hpackDecoder.Write(frag); err != nil {
		// TODO: convert to stream error I assume?
		return err
	}
	if !end {
		return nil
	}
	if err := sc.hpackDecoder.Close(); err != nil {
		// TODO: convert to stream error I assume?
		return err
	}
	curStream := sc.curStream
	sc.curHeaderStreamID = 0
	sc.curStream = nil

	// TODO: transition streamID state
	go sc.startHandler(curStream, sc.method, sc.path, sc.scheme, sc.authority, sc.header)

	return nil
}

func (sc *serverConn) startHandler(st *stream, method, path, scheme, authority string, reqHeader http.Header) {
	var tlsState *tls.ConnectionState // make this non-nil if https
	if scheme == "https" {
		// TODO: get from sc's ConnectionState
		tlsState = &tls.ConnectionState{}
	}
	req := &http.Request{
		Method:     method,
		URL:        &url.URL{},
		RemoteAddr: sc.conn.RemoteAddr().String(),
		RequestURI: path,
		Proto:      "HTTP/2.0",
		ProtoMajor: 2,
		ProtoMinor: 0,
		TLS:        tlsState,
		Host:       authority,
	}
	if vv, ok := reqHeader["Content-Length"]; ok {
		req.ContentLength, _ = strconv.ParseInt(vv[0], 10, 64)
	} else {
		req.ContentLength = -1
	}
	var rw = httptest.NewRecorder()
	sc.handler.ServeHTTP(rw, req)
	log.Printf("Got code %d: body: %q\n", rw.Code, rw.Body.Bytes())
}

// ConfigureServer adds HTTP/2 support to a net/http Server.
//
// The configuration conf may be nil.
//
// ConfigureServer must be called before s begins serving.
func ConfigureServer(s *http.Server, conf *Server) {
	if conf == nil {
		conf = new(Server)
	}
	if s.TLSConfig == nil {
		s.TLSConfig = new(tls.Config)
	}
	haveNPN := false
	for _, p := range s.TLSConfig.NextProtos {
		if p == npnProto {
			haveNPN = true
			break
		}
	}
	if !haveNPN {
		s.TLSConfig.NextProtos = append(s.TLSConfig.NextProtos, npnProto)
	}

	if s.TLSNextProto == nil {
		s.TLSNextProto = map[string]func(*http.Server, *tls.Conn, http.Handler){}
	}
	s.TLSNextProto[npnProto] = func(hs *http.Server, c *tls.Conn, h http.Handler) {
		if testHookOnConn != nil {
			testHookOnConn()
		}
		conf.handleConn(hs, c, h)
	}
}

var testHookOnConn func() // for testing
<commit_msg>start of requestBody<commit_after>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// See https://code.google.com/p/go/source/browse/CONTRIBUTORS
// Licensed under the same terms as Go itself:
// https://code.google.com/p/go/source/browse/LICENSE

// Package http2 implements the HTTP/2 protocol.
//
// This is a work in progress. This package is low-level and intended
// to be used directly by very few people. Most users will use it
// indirectly through integration with the net/http package. See
// ConfigureServer. That ConfigureServer call will likely be automatic
// or available via an empty import in the future.
//
// This package currently targets draft-14. See http://http2.github.io/
package http2

import (
	"bytes"
	"crypto/tls"
	"errors"
	"io"
	"log"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strconv"
	"strings"

	"github.com/bradfitz/http2/hpack"
)

const (
	// ClientPreface is the string that must be sent by new
	// connections from clients.
	ClientPreface = "PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"
)

var (
	clientPreface = []byte(ClientPreface)
)

const (
	npnProto = "h2-14"

	// http://http2.github.io/http2-spec/#SettingValues
	initialHeaderTableSize = 4096
)

// Server is an HTTP2 server.
type Server struct {
	// MaxStreams optionally ...
	MaxStreams int
}

func (srv *Server) handleConn(hs *http.Server, c *tls.Conn, h http.Handler) {
	sc := &serverConn{
		hs:             hs,
		conn:           c,
		handler:        h,
		framer:         NewFramer(c, c),
		streams:        make(map[uint32]*stream),
		canonHeader:    make(map[string]string),
		readFrameCh:    make(chan frameAndProcessed),
		readFrameErrCh: make(chan error, 1),
		doneServing:    make(chan struct{}),
	}
	sc.hpackDecoder = hpack.NewDecoder(initialHeaderTableSize, sc.onNewHeaderField)
	sc.serve()
}

// frameAndProcessed coordinates the readFrames and serve goroutines, since
// the Framer interface only permits the most recently-read Frame from being
// accessed. The serve goroutine sends on processed to signal to the readFrames
// goroutine that another frame may be read.
type frameAndProcessed struct {
	f         Frame
	processed chan struct{}
}

type serverConn struct {
	hs             *http.Server
	conn           *tls.Conn
	handler        http.Handler
	framer         *Framer
	doneServing    chan struct{}          // closed when serverConn.serve ends
	readFrameCh    chan frameAndProcessed // written by serverConn.readFrames
	readFrameErrCh chan error

	maxStreamID uint32 // max ever seen
	streams     map[uint32]*stream

	// State related to parsing current headers:
	hpackDecoder *hpack.Decoder
	header       http.Header
	canonHeader  map[string]string // http2-lower-case -> Go-Canonical-Case

	method, path, scheme, authority string

	// curHeaderStreamID is non-zero if we're in the middle
	// of parsing headers that span multiple frames.
	curHeaderStreamID uint32
	curStream         *stream
}

type streamState int

const (
	stateIdle streamState = iota
	stateOpen
	stateHalfClosedLocal
	stateHalfClosedRemote
	stateResvLocal
	stateResvRemote
	stateClosed
)

type stream struct {
	id    uint32
	state streamState // owned by serverConn's processing loop
}

func (sc *serverConn) state(streamID uint32) streamState {
	// http://http2.github.io/http2-spec/#rfc.section.5.1
	if st, ok := sc.streams[streamID]; ok {
		return st.state
	}
	// "The first use of a new stream identifier implicitly closes all
	// streams in the "idle" state that might have been initiated by
	// that peer with a lower-valued stream identifier. For example, if
	// a client sends a HEADERS frame on stream 7 without ever sending a
	// frame on stream 5, then stream 5 transitions to the "closed"
	// state when the first frame for stream 7 is sent or received."
	if streamID <= sc.maxStreamID {
		return stateClosed
	}
	return stateIdle
}

func (sc *serverConn) logf(format string, args ...interface{}) {
	if lg := sc.hs.ErrorLog; lg != nil {
		lg.Printf(format, args...)
	} else {
		log.Printf(format, args...)
	}
}

func (sc *serverConn) onNewHeaderField(f hpack.HeaderField) {
	log.Printf("Header field: +%v", f)
	if strings.HasPrefix(f.Name, ":") {
		switch f.Name {
		case ":method":
			sc.method = f.Value
		case ":path":
			sc.path = f.Value
		case ":scheme":
			sc.scheme = f.Value
		case ":authority":
			sc.authority = f.Value
		default:
			log.Printf("Ignoring unknown pseudo-header %q", f.Name)
		}
		return
	}
	sc.header.Add(sc.canonicalHeader(f.Name), f.Value)
}

func (sc *serverConn) canonicalHeader(v string) string {
	// TODO: use a sync.Pool instead of putting the cache on *serverConn?
	cv, ok := sc.canonHeader[v]
	if !ok {
		cv = http.CanonicalHeaderKey(v)
		sc.canonHeader[v] = cv
	}
	return cv
}

func (sc *serverConn) readFrames() {
	processed := make(chan struct{}, 1)
	for {
		f, err := sc.framer.ReadFrame()
		if err != nil {
			close(sc.readFrameCh)
			sc.readFrameErrCh <- err
			return
		}
		sc.readFrameCh <- frameAndProcessed{f, processed}
		<-processed
	}
}

func (sc *serverConn) serve() {
	defer sc.conn.Close()
	defer close(sc.doneServing)

	log.Printf("HTTP/2 connection from %v on %p", sc.conn.RemoteAddr(), sc.hs)

	// Read the client preface
	buf := make([]byte, len(ClientPreface))
	// TODO: timeout reading from the client
	if _, err := io.ReadFull(sc.conn, buf); err != nil {
		sc.logf("error reading client preface: %v", err)
		return
	}
	if !bytes.Equal(buf, clientPreface) {
		sc.logf("bogus greeting from client: %q", buf)
		return
	}
	log.Printf("client %v said hello", sc.conn.RemoteAddr())

	go sc.readFrames()

	for {
		select {
		case fp, ok := <-sc.readFrameCh:
			if !ok {
				err := <-sc.readFrameErrCh
				if err != io.EOF {
					sc.logf("client %s stopped sending frames: %v", sc.conn.RemoteAddr(), err)
				}
				return
			}
			f := fp.f
			log.Printf("got %v: %#v", f.Header(), f)
			err := sc.processFrame(f)
			fp.processed <- struct{}{} // let readFrames proceed
			if h2e, ok := err.(Error); ok {
				if h2e.IsConnectionError() {
					sc.logf("Disconnection; connection error: %v", err)
					return
				}
				// TODO: stream errors, etc
			}
			if err != nil {
				sc.logf("Disconnection due to other error: %v", err)
				return
			}
		}
	}
}

func (sc *serverConn) processFrame(f Frame) error {
	if s := sc.curHeaderStreamID; s != 0 {
		if cf, ok := f.(*ContinuationFrame); !ok {
			return ConnectionError(ErrCodeProtocol)
		} else if cf.Header().StreamID != s {
			return ConnectionError(ErrCodeProtocol)
		}
	}

	switch f := f.(type) {
	case *SettingsFrame:
		return sc.processSettings(f)
	case *HeadersFrame:
		return sc.processHeaders(f)
	case *ContinuationFrame:
		return sc.processContinuation(f)
	default:
		log.Printf("Ignoring unknown %v", f.Header)
		return nil
	}
}

func (sc *serverConn) processSettings(f *SettingsFrame) error {
	f.ForeachSetting(func(s Setting) {
		log.Printf("  setting %s = %v", s.ID, s.Val)
	})
	return nil
}

func (sc *serverConn) processHeaders(f *HeadersFrame) error {
	id := f.Header().StreamID

	// http://http2.github.io/http2-spec/#rfc.section.5.1.1
	if id%2 != 1 || id <= sc.maxStreamID {
		// Streams initiated by a client MUST use odd-numbered
		// stream identifiers. [...] The identifier of a newly
		// established stream MUST be numerically greater than all
		// streams that the initiating endpoint has opened or
		// reserved. [...]  An endpoint that receives an unexpected
		// stream identifier MUST respond with a connection error
		// (Section 5.4.1) of type PROTOCOL_ERROR.
		return ConnectionError(ErrCodeProtocol)
	}
	if id > sc.maxStreamID {
		sc.maxStreamID = id
	}

	st := &stream{
		id:    id,
		state: stateOpen,
	}
	if f.Header().Flags.Has(FlagHeadersEndStream) {
		st.state = stateHalfClosedRemote
	}
	sc.streams[id] = st
	sc.header = make(http.Header)
	sc.curHeaderStreamID = id
	sc.curStream = st
	return sc.processHeaderBlockFragment(f.HeaderBlockFragment(), f.HeadersEnded())
}

func (sc *serverConn) processContinuation(f *ContinuationFrame) error {
	return sc.processHeaderBlockFragment(f.HeaderBlockFragment(), f.HeadersEnded())
}

func (sc *serverConn) processHeaderBlockFragment(frag []byte, end bool) error {
	if _, err := sc.hpackDecoder.Write(frag); err != nil {
		// TODO: convert to stream error I assume?
		return err
	}
	if !end {
		return nil
	}
	if err := sc.hpackDecoder.Close(); err != nil {
		// TODO: convert to stream error I assume?
		return err
	}
	curStream := sc.curStream
	sc.curHeaderStreamID = 0
	sc.curStream = nil

	// TODO: transition streamID state
	go sc.startHandler(curStream.id, curStream.state == stateOpen, sc.method, sc.path, sc.scheme, sc.authority, sc.header)

	return nil
}

func (sc *serverConn) startHandler(streamID uint32, bodyOpen bool, method, path, scheme, authority string, reqHeader http.Header) {
	var tlsState *tls.ConnectionState // make this non-nil if https
	if scheme == "https" {
		// TODO: get from sc's ConnectionState
		tlsState = &tls.ConnectionState{}
	}
	req := &http.Request{
		Method:     method,
		URL:        &url.URL{},
		RemoteAddr: sc.conn.RemoteAddr().String(),
		RequestURI: path,
		Proto:      "HTTP/2.0",
		ProtoMajor: 2,
		ProtoMinor: 0,
		TLS:        tlsState,
		Host:       authority,
		Body: &requestBody{
			sc:       sc,
			streamID: streamID,
			hasBody:  bodyOpen,
		},
	}
	if vv, ok := reqHeader["Content-Length"]; ok {
		req.ContentLength, _ = strconv.ParseInt(vv[0], 10, 64)
	} else {
		req.ContentLength = -1
	}
	var rw = httptest.NewRecorder()
	sc.handler.ServeHTTP(rw, req)
	log.Printf("Got code %d: body: %q\n", rw.Code, rw.Body.Bytes())
}

// ConfigureServer adds HTTP/2 support to a net/http Server.
//
// The configuration conf may be nil.
//
// ConfigureServer must be called before s begins serving.
func ConfigureServer(s *http.Server, conf *Server) {
	if conf == nil {
		conf = new(Server)
	}
	if s.TLSConfig == nil {
		s.TLSConfig = new(tls.Config)
	}
	haveNPN := false
	for _, p := range s.TLSConfig.NextProtos {
		if p == npnProto {
			haveNPN = true
			break
		}
	}
	if !haveNPN {
		s.TLSConfig.NextProtos = append(s.TLSConfig.NextProtos, npnProto)
	}

	if s.TLSNextProto == nil {
		s.TLSNextProto = map[string]func(*http.Server, *tls.Conn, http.Handler){}
	}
	s.TLSNextProto[npnProto] = func(hs *http.Server, c *tls.Conn, h http.Handler) {
		if testHookOnConn != nil {
			testHookOnConn()
		}
		conf.handleConn(hs, c, h)
	}
}

type requestBody struct {
	sc       *serverConn
	streamID uint32
	hasBody  bool
	closed   bool
}

func (b *requestBody) Close() error {
	b.closed = true
	return nil
}

func (b *requestBody) Read(p []byte) (n int, err error) {
	if !b.hasBody {
		return 0, io.EOF
	}
	// TODO: implement
	return 0, errors.New("TODO: we don't handle request bodies yet")
}

var testHookOnConn func() // for testing
<|endoftext|>


<commit_before>package config

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/user"
	"path"
	"strings"
	"testing"

	"github.com/opentable/sous/ext/docker"
)

func TestDefaultStateLocation(t *testing.T) {
	testCases := []struct {
		XDGDataHome, StateLocation string
	}{
		{"", "~/.local/share/sous/state"}, // Note: ~ is handled in code below.
		{"some/dir", "some/dir/sous/state"},
	}

	c := &Config{}
	for i, tc := range testCases {
		os.Setenv("XDG_DATA_HOME", tc.XDGDataHome)
		expected := tc.StateLocation
		u, err := user.Current()
		if err != nil {
			t.Error(err)
			continue
		}
		if expected[0] == '~' {
			expected = strings.TrimPrefix(expected, "~")
			expected = path.Join(u.HomeDir, expected)
		}
		actual, err := c.defaultStateLocation()
		if err != nil {
			t.Error(err)
			continue
		}
		if actual != expected {
			t.Errorf("%d: %q got %q; want %q", i, tc.XDGDataHome, actual, expected)
		}
	}
}

func TestConfig_Validate(t *testing.T) {
	cfg := &Config{
		Server:      "not_a_url",
		SiblingURLs: map[string]string{"left": "tcp://is.a.url"},
	}

	checkNotValid := func() {
		if err := cfg.Validate(); err == nil {
			t.Errorf("%v returns true from Validate()", cfg)
		}
	}

	checkValid := func() {
		if err := cfg.Validate(); err != nil {
			t.Errorf("%v returns false from Validate()", cfg)
		}
	}

	checkNotValid()

	cfg.Server = "https://a.sous.server"
	checkNotValid()

	cfg.SiblingURLs["left"] = "http://sibling.sous"
	checkValid()

	cfg.Server = "zxcv"
	checkNotValid()

	cfg.Server = ""
	checkValid()
}

func TestConfig_Equals(t *testing.T) {
	expected := &Config{
		StateLocation: "statelocation",
		Server:        "server",
		SiblingURLs:   map[string]string{"x": "sibling", "y": "urls"},
		BuildStateDir: "buildstatedir",
		Docker: docker.Config{
			RegistryHost:       "registryhost",
			DatabaseDriver:     "databasedriver",
			DatabaseConnection: "databaseconnection",
		},
	}
	var actual *Config

	checkNotEqual := func() {
		if expected.Equal(actual) {
			t.Errorf("expected.Equal(actual) was true:\n%v\n%v", expected, actual)
		}
		if actual.Equal(expected) {
			t.Errorf("actual.Equal(expected) was true:\n%v\n%v", actual, expected)
		}
	}

	actual = &Config{}
	checkNotEqual()

	actual.StateLocation = "statelocation"
	checkNotEqual()

	actual.Server = "server"
	checkNotEqual()

	actual.BuildStateDir = "buildstatedir"
	checkNotEqual()

	actual.Docker = expected.Docker
	checkNotEqual()

	actual.SiblingURLs = map[string]string{"x": "sibling", "y": "urls"}
	if !expected.Equal(actual) {
		t.Errorf("expected.Equal(actual) was false:\n%v\n%v", expected, actual)
	}
	if !actual.Equal(expected) {
		t.Errorf("actual.Equal(expected) was false:\n%v\n%v", actual, expected)
	}

	// Confirming that the two Docker structs are separate memory
	actual.Docker.DatabaseConnection = ""
	checkNotEqual()
}

func TestConfigParses(t *testing.T) {

}

func TestEnsureDirExists(t *testing.T) {
	testDataDir := "testdata/gen"
	if err := os.RemoveAll(testDataDir); err != nil {
		t.Fatal(err)
	}
	extantDir := path.Join(testDataDir, "test-dir")
	if err := os.MkdirAll(extantDir, 0777); err != nil {
		t.Fatal(err)
	}
	if err := EnsureDirExists(extantDir); err != nil {
		t.Error(err)
	}
	nonExistentDir := path.Join(testDataDir, "other-dir")
	if err := EnsureDirExists(nonExistentDir); err != nil {
		t.Error(err)
	}
	s, err := os.Stat(nonExistentDir)
	if err != nil {
		t.Fatal(err)
	}
	if !s.IsDir() {
		t.Errorf("created path %q is not a directory", nonExistentDir)
	}
	filePath := path.Join(testDataDir, "a-file")
	if err := ioutil.WriteFile(filePath, []byte("x"), 0777); err != nil {
		t.Fatal(err)
	}
	expectedErr := fmt.Sprintf("%q exists and is not a directory", filePath)
	err = EnsureDirExists(filePath)
	if err == nil {
		t.Errorf("got nil error; want %q", expectedErr)
	}
	actualErr := err.Error()
	if actualErr != expectedErr {
		t.Errorf("got error %q; want %q", actualErr, expectedErr)
	}
}
<commit_msg>REmoving dead test<commit_after>package config

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/user"
	"path"
	"strings"
	"testing"

	"github.com/opentable/sous/ext/docker"
)

func TestDefaultStateLocation(t *testing.T) {
	testCases := []struct {
		XDGDataHome, StateLocation string
	}{
		{"", "~/.local/share/sous/state"}, // Note: ~ is handled in code below.
		{"some/dir", "some/dir/sous/state"},
	}

	c := &Config{}
	for i, tc := range testCases {
		os.Setenv("XDG_DATA_HOME", tc.XDGDataHome)
		expected := tc.StateLocation
		u, err := user.Current()
		if err != nil {
			t.Error(err)
			continue
		}
		if expected[0] == '~' {
			expected = strings.TrimPrefix(expected, "~")
			expected = path.Join(u.HomeDir, expected)
		}
		actual, err := c.defaultStateLocation()
		if err != nil {
			t.Error(err)
			continue
		}
		if actual != expected {
			t.Errorf("%d: %q got %q; want %q", i, tc.XDGDataHome, actual, expected)
		}
	}
}

func TestConfig_Validate(t *testing.T) {
	cfg := &Config{
		Server:      "not_a_url",
		SiblingURLs: map[string]string{"left": "tcp://is.a.url"},
	}

	checkNotValid := func() {
		if err := cfg.Validate(); err == nil {
			t.Errorf("%v returns true from Validate()", cfg)
		}
	}

	checkValid := func() {
		if err := cfg.Validate(); err != nil {
			t.Errorf("%v returns false from Validate()", cfg)
		}
	}

	checkNotValid()

	cfg.Server = "https://a.sous.server"
	checkNotValid()

	cfg.SiblingURLs["left"] = "http://sibling.sous"
	checkValid()

	cfg.Server = "zxcv"
	checkNotValid()

	cfg.Server = ""
	checkValid()
}

func TestConfig_Equals(t *testing.T) {
	expected := &Config{
		StateLocation: "statelocation",
		Server:        "server",
		SiblingURLs:   map[string]string{"x": "sibling", "y": "urls"},
		BuildStateDir: "buildstatedir",
		Docker: docker.Config{
			RegistryHost:       "registryhost",
			DatabaseDriver:     "databasedriver",
			DatabaseConnection: "databaseconnection",
		},
	}
	var actual *Config

	checkNotEqual := func() {
		if expected.Equal(actual) {
			t.Errorf("expected.Equal(actual) was true:\n%v\n%v", expected, actual)
		}
		if actual.Equal(expected) {
			t.Errorf("actual.Equal(expected) was true:\n%v\n%v", actual, expected)
		}
	}

	actual = &Config{}
	checkNotEqual()

	actual.StateLocation = "statelocation"
	checkNotEqual()

	actual.Server = "server"
	checkNotEqual()

	actual.BuildStateDir = "buildstatedir"
	checkNotEqual()

	actual.Docker = expected.Docker
	checkNotEqual()

	actual.SiblingURLs = map[string]string{"x": "sibling", "y": "urls"}
	if !expected.Equal(actual) {
		t.Errorf("expected.Equal(actual) was false:\n%v\n%v", expected, actual)
	}
	if !actual.Equal(expected) {
		t.Errorf("actual.Equal(expected) was false:\n%v\n%v", actual, expected)
	}

	// Confirming that the two Docker structs are separate memory
	actual.Docker.DatabaseConnection = ""
	checkNotEqual()
}

func TestEnsureDirExists(t *testing.T) {
	testDataDir := "testdata/gen"
	if err := os.RemoveAll(testDataDir); err != nil {
		t.Fatal(err)
	}
	extantDir := path.Join(testDataDir, "test-dir")
	if err := os.MkdirAll(extantDir, 0777); err != nil {
		t.Fatal(err)
	}
	if err := EnsureDirExists(extantDir); err != nil {
		t.Error(err)
	}
	nonExistentDir := path.Join(testDataDir, "other-dir")
	if err := EnsureDirExists(nonExistentDir); err != nil {
		t.Error(err)
	}
	s, err := os.Stat(nonExistentDir)
	if err != nil {
		t.Fatal(err)
	}
	if !s.IsDir() {
		t.Errorf("created path %q is not a directory", nonExistentDir)
	}
	filePath := path.Join(testDataDir, "a-file")
	if err := ioutil.WriteFile(filePath, []byte("x"), 0777); err != nil {
		t.Fatal(err)
	}
	expectedErr := fmt.Sprintf("%q exists and is not a directory", filePath)
	err = EnsureDirExists(filePath)
	if err == nil {
		t.Errorf("got nil error; want %q", expectedErr)
	}
	actualErr := err.Error()
	if actualErr != expectedErr {
		t.Errorf("got error %q; want %q", actualErr, expectedErr)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"image"
	"image/jpeg"
	"log"
	"os"
	"time"

	"github.com/mdlayher/tapcam/camera"
	"github.com/mdlayher/tapcam/tapcamclient"

	"github.com/disintegration/imaging"
)

const (
	envHost    = "TAPCAMCTL_HOST"
	envUser    = "TAPCAMCTL_USER"
	envKeyFile = "TAPCAMCTL_KEY"
	envTarget  = "TAPCAMCTL_TARGET"
	envInvert  = "TAPCAMCTL_INVERT"
)

var (
	device = flag.String("d", camera.DefaultDevice, "webcam device location")
	format = flag.String("f", string(camera.FormatJPEG), "webcam image capture format")
	size   = flag.String("s", camera.Resolution1080p.String(), "webcam image size")
	delay  = flag.Int("delay", 0, "delay in seconds before taking a picture")

	host    = flag.String("host", "", "tapcamd host")
	user    = flag.String("user", "", "tapcamd SSH user")
	keyFile = flag.String("key", "", "tapcamd SSH private key file")
	target  = flag.String("target", "", "tapcamd image upload target location")
	invert  = flag.Bool("i", false, "invert image after capture")
)

func main() {
	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)
	flag.Parse()

	config, err := checkFlags(
		*host,
		*user,
		*keyFile,
		*target,
		*invert,
	)
	if err != nil {
		log.Fatal(err)
	}

	tcc, err := tapcamclient.New(
		config.host,
		tapcamclient.SSHUserKeyFile(config.user, config.keyFile),
	)
	if err != nil {
		log.Fatal(err)
	}

	resolution, err := camera.NewResolution(*size)
	if err != nil {
		log.Fatal(err)
	}

	cam, err := camera.New(
		camera.SetDevice(*device),
		camera.SetFormat(camera.Format(*format)),
		camera.SetResolution(resolution),
	)
	if err != nil {
		log.Fatal(err)
	}

	if d := *delay; d > 0 {
		log.Printf("taking picture in %d seconds", d)
		for i := 0; i < d; i++ {
			fmt.Print(". ")
			time.Sleep(1 * time.Second)
		}
		fmt.Println()
	}

	rc, err := cam.Capture()
	if err != nil {
		log.Fatal(err)
	}

	if !config.invert {
		if err := tcc.Upload(config.target, rc); err != nil {
			log.Fatal(err)
		}

		if err := tcc.Close(); err != nil {
			log.Fatal(err)
		}

		return
	}

	img, _, err := image.Decode(rc)
	if err != nil {
		log.Fatal(err)
	}

	if err := rc.Close(); err != nil {
		log.Fatal(err)
	}

	img = imaging.FlipV(img)
	img = imaging.FlipH(img)

	var buf bytes.Buffer
	if err := jpeg.Encode(&buf, img, nil); err != nil {
		log.Fatal(err)
	}

	if err := tcc.Upload(config.target, &buf); err != nil {
		log.Fatal(err)
	}

	if err := tcc.Close(); err != nil {
		log.Fatal(err)
	}
}

type config struct {
	host    string
	user    string
	keyFile string
	target  string
	invert  bool
}

func checkFlags(
	host string,
	user string,
	keyFile string,
	target string,
	invert bool,
) (*config, error) {
	c := new(config)

	if host == "" {
		host = os.Getenv(envHost)
	}
	if host == "" {
		return nil, fmt.Errorf("must specify SFTP host using -host flag or $%s", envHost)
	}
	c.host = host

	if user == "" {
		user = os.Getenv(envUser)
	}
	if user == "" {
		return nil, fmt.Errorf("must specify SFTP user using -user flag or $%s", envUser)
	}
	c.user = user

	if keyFile == "" {
		keyFile = os.Getenv(envKeyFile)
	}
	if keyFile == "" {
		return nil, fmt.Errorf("must specify SFTP private key file using -key flag or $%s", envKeyFile)
	}
	c.keyFile = keyFile

	if target == "" {
		target = os.Getenv(envTarget)
	}
	if target == "" {
		return nil, fmt.Errorf("must specify SFTP file target using -target flag or $%s", envTarget)
	}
	c.target = target

	if !invert {
		invert = os.Getenv(envInvert) == "true"
	}
	c.invert = invert

	return c, nil
}
<commit_msg>tapcamctl: add logging for each processing step<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"image"
	"image/jpeg"
	"log"
	"os"
	"time"

	"github.com/mdlayher/tapcam/camera"
	"github.com/mdlayher/tapcam/tapcamclient"

	"github.com/disintegration/imaging"
)

const (
	envHost    = "TAPCAMCTL_HOST"
	envUser    = "TAPCAMCTL_USER"
	envKeyFile = "TAPCAMCTL_KEY"
	envTarget  = "TAPCAMCTL_TARGET"
	envInvert  = "TAPCAMCTL_INVERT"
)

var (
	device = flag.String("d", camera.DefaultDevice, "webcam device location")
	format = flag.String("f", string(camera.FormatJPEG), "webcam image capture format")
	size   = flag.String("s", camera.Resolution1080p.String(), "webcam image size")
	delay  = flag.Int("delay", 0, "delay in seconds before taking a picture")

	host    = flag.String("host", "", "tapcamd host")
	user    = flag.String("user", "", "tapcamd SSH user")
	keyFile = flag.String("key", "", "tapcamd SSH private key file")
	target  = flag.String("target", "", "tapcamd image upload target location")
	invert  = flag.Bool("i", false, "invert image after capture")
)

func main() {
	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)
	flag.Parse()

	config, err := checkFlags(
		*host,
		*user,
		*keyFile,
		*target,
		*invert,
	)
	if err != nil {
		log.Fatal(err)
	}

	tcc, err := tapcamclient.New(
		config.host,
		tapcamclient.SSHUserKeyFile(config.user, config.keyFile),
	)
	if err != nil {
		log.Fatal(err)
	}

	resolution, err := camera.NewResolution(*size)
	if err != nil {
		log.Fatal(err)
	}

	cam, err := camera.New(
		camera.SetDevice(*device),
		camera.SetFormat(camera.Format(*format)),
		camera.SetResolution(resolution),
	)
	if err != nil {
		log.Fatal(err)
	}

	if d := *delay; d > 0 {
		log.Printf("capturing image in %d seconds", d)
		for i := 0; i < d; i++ {
			fmt.Print(". ")
			time.Sleep(1 * time.Second)
		}
		fmt.Println()
	}

	log.Println("capturing image with camera")

	rc, err := cam.Capture()
	if err != nil {
		log.Fatal(err)
	}

	if !config.invert {
		log.Println("uploading image to server")

		if err := tcc.Upload(config.target, rc); err != nil {
			log.Fatal(err)
		}

		if err := tcc.Close(); err != nil {
			log.Fatal(err)
		}

		log.Println("done!")
		return
	}

	log.Println("inverting image before upload")

	img, _, err := image.Decode(rc)
	if err != nil {
		log.Fatal(err)
	}

	if err := rc.Close(); err != nil {
		log.Fatal(err)
	}

	img = imaging.FlipV(img)
	img = imaging.FlipH(img)

	var buf bytes.Buffer
	if err := jpeg.Encode(&buf, img, nil); err != nil {
		log.Fatal(err)
	}

	log.Println("uploading image to server")

	if err := tcc.Upload(config.target, &buf); err != nil {
		log.Fatal(err)
	}

	if err := tcc.Close(); err != nil {
		log.Fatal(err)
	}

	log.Println("done!")
}

type config struct {
	host    string
	user    string
	keyFile string
	target  string
	invert  bool
}

func checkFlags(
	host string,
	user string,
	keyFile string,
	target string,
	invert bool,
) (*config, error) {
	c := new(config)

	if host == "" {
		host = os.Getenv(envHost)
	}
	if host == "" {
		return nil, fmt.Errorf("must specify SFTP host using -host flag or $%s", envHost)
	}
	c.host = host

	if user == "" {
		user = os.Getenv(envUser)
	}
	if user == "" {
		return nil, fmt.Errorf("must specify SFTP user using -user flag or $%s", envUser)
	}
	c.user = user

	if keyFile == "" {
		keyFile = os.Getenv(envKeyFile)
	}
	if keyFile == "" {
		return nil, fmt.Errorf("must specify SFTP private key file using -key flag or $%s", envKeyFile)
	}
	c.keyFile = keyFile

	if target == "" {
		target = os.Getenv(envTarget)
	}
	if target == "" {
		return nil, fmt.Errorf("must specify SFTP file target using -target flag or $%s", envTarget)
	}
	c.target = target

	if !invert {
		invert = os.Getenv(envInvert) == "true"
	}
	c.invert = invert

	return c, nil
}
<|endoftext|>


<commit_before>package config

import (
	"github.com/ghthor/gospec"
	. "github.com/ghthor/gospec"
	"testing"
)

func TestUnitSpecs(t *testing.T) {
	r := gospec.NewRunner()

	r.AddSpec(DescribeConfigLoading)

	gospec.MainGoTest(r, t)
}

func DescribeConfigLoading(c gospec.Context) {
	c.Specify("Config can be parsed from json file", func() {
		expectedConfig := Config{
			"a/path/to/a/git/directory",
		}

		config, err := ReadFromFile("config.example.json")
		c.Assume(err, IsNil)
		c.Expect(config, Equals, expectedConfig)
	})
}
<commit_msg>Change the name the test's are closer to integration | But they are both unit+integration cause the package is so simple<commit_after>package config

import (
	"github.com/ghthor/gospec"
	. "github.com/ghthor/gospec"
	"testing"
)

func TestSpecs(t *testing.T) {
	r := gospec.NewRunner()

	r.AddSpec(DescribeConfigLoading)

	gospec.MainGoTest(r, t)
}

func DescribeConfigLoading(c gospec.Context) {
	c.Specify("Config can be parsed from json file", func() {
		expectedConfig := Config{
			"a/path/to/a/git/directory",
		}

		config, err := ReadFromFile("config.example.json")
		c.Assume(err, IsNil)
		c.Expect(config, Equals, expectedConfig)
	})
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sort"
	"sync"
	"time"

	"github.com/alexbakker/tox4go/bootstrap"
	"github.com/alexbakker/tox4go/crypto"
	"github.com/alexbakker/tox4go/dht"
	"github.com/alexbakker/tox4go/dht/ping"
	"github.com/alexbakker/tox4go/relay"
	"github.com/alexbakker/tox4go/transport"
	"github.com/didip/tollbooth"
)

const (
	enableIpv6  = true
	probeRate   = 1 * time.Minute
	refreshRate = 5 * time.Minute
)

var (
	lastScan    int64
	lastRefresh int64
	nodes       = []*toxNode{}
	nodesMutex  = sync.Mutex{}
	tcpPorts    = []int{443, 3389, 33445}
)

func main() {
	if parseFlags() {
		return
	}

	// load state if available
	state, err := loadState()
	if err != nil {
		log.Fatalf("error loading state: %s", err.Error())
	}
	lastScan = state.LastScan
	lastRefresh = state.LastRefresh
	nodes = state.Nodes

	if err := loadCountries(); err != nil {
		log.Fatalf("error loading countries.json: %s", err.Error())
	}

	inst, err := NewInstance(":33450")
	if err != nil {
		log.Fatalf("fatal: %s", err.Error())
	}
	inst.UDPTransport.Handle(dht.PacketIDSendNodes, inst.handleSendNodesPacket)
	inst.UDPTransport.Handle(bootstrap.PacketIDBootstrapInfo, handleBootstrapInfoPacket)

	//handle stop signal
	interruptChan := make(chan os.Signal)
	signal.Notify(interruptChan, os.Interrupt)

	//setup http server
	listener, err := net.Listen("tcp4", fmt.Sprintf("localhost:%d", httpListenPort))
	if err != nil {
		log.Fatalf("error in net.Listen: %s", err.Error())
	}
	limiter := tollbooth.NewLimiter(1, 2*time.Second)
	limiter.Methods = []string{"POST"}
	limiter.IPLookups = []string{"X-Forwarded-For", "RemoteAddr", "X-Real-IP"}
	serveMux := http.NewServeMux()
	serveMux.HandleFunc("/", handleHTTPRequest)
	serveMux.Handle("/test", tollbooth.LimitFuncHandler(limiter, handleHTTPRequest))
	serveMux.HandleFunc("/json", handleJSONRequest)
	go func() {
		err := http.Serve(listener, serveMux)
		if err != nil {
			log.Printf("http server error: %s\n", err.Error())
			interruptChan <- os.Interrupt
		}
	}()

	//listen for tox packets
	go func() {
		err := inst.UDPTransport.Listen()
		if err != nil {
			log.Printf("udp transport error: %s\n", err.Error())
			interruptChan <- os.Interrupt
		}
	}()
	//go tcpTransport.Listen()

	err = refreshNodes()
	if err != nil {
		log.Fatal(err.Error())
	}
	inst.probeNodes()

	probeTicker := time.NewTicker(probeRate)
	refreshTicker := time.NewTicker(refreshRate)
	updateTicker := time.NewTicker(30 * time.Second)
	run := true

	for run {
		select {
		case <-interruptChan:
			fmt.Printf("killing routines\n")
			probeTicker.Stop()
			refreshTicker.Stop()
			updateTicker.Stop()
			inst.UDPTransport.Stop()
			//tcpTransport.Stop()
			listener.Close()
			run = false
		case <-probeTicker.C:
			// we want an empty ping list at the start of every probe
			inst.PingsMutex.Lock()
			inst.Pings.Clear(false)
			inst.PingsMutex.Unlock()

			nodesMutex.Lock()
			err := inst.probeNodes()
			nodesMutex.Unlock()
			if err != nil {
				log.Printf("error while trying to probe nodes: %s", err.Error())
			}
		case <-refreshTicker.C:
			err := refreshNodes()
			if err != nil {
				log.Printf("error while trying to refresh nodes: %s", err.Error())
			}
		case <-updateTicker.C:
			inst.PingsMutex.Lock()
			inst.Pings.Clear(true)
			inst.PingsMutex.Unlock()

			nodesMutex.Lock()
			for _, node := range nodes {
				if time.Now().Sub(time.Unix(node.LastPing, 0)) > time.Minute*2 {
					node.UDPStatus = false
				}
			}
			sort.Stable(nodeSlice(nodes))

			state := getState()
			saveState(state)
			nodesMutex.Unlock()
		}
	}
}

func refreshNodes() error {
	parsedNodes, err := parseNodes()
	if err != nil {
		return err
	}

	nodesMutex.Lock()
	for _, freshNode := range parsedNodes {
		for _, node := range nodes {
			if freshNode.PublicKey == node.PublicKey {
				freshNode.LastPing = node.LastPing
				freshNode.UDPStatus = node.UDPStatus
				freshNode.TCPStatus = node.TCPStatus
				freshNode.TCPPorts = node.TCPPorts
				freshNode.MOTD = node.MOTD
				freshNode.Version = node.Version
				break
			}
		}
	}
	nodes = parsedNodes
	sort.Stable(nodeSlice(nodes))
	nodesMutex.Unlock()

	lastRefresh = time.Now().Unix()
	return nil
}

func (i *instance) probeNodes() error {
	for _, node := range nodes {
		err := i.getBootstrapInfo(node)
		if err != nil {
			fmt.Println(err.Error())
		}

		p, err := i.getNodes(node)
		if err != nil {
			fmt.Println(err.Error())
		} else {
			i.PingsMutex.Lock()
			err = i.Pings.Add(p)
			if err != nil {
				fmt.Println(err.Error())
			}
			i.PingsMutex.Unlock()
		}

		ports := tcpPorts
		exists := false
		for _, i := range ports {
			if i == node.Port {
				exists = true
			}
		}
		if !exists {
			ports = append(ports, node.Port)
		}

		go i.probeNodeTCPPorts(node, ports)
	}

	lastScan = time.Now().Unix()
	return nil
}

func (i *instance) probeNodeTCPPorts(node *toxNode, ports []int) {
	c := make(chan int)
	for _, port := range ports {
		go func(p int) {
			conn, err := connectTCP(node, p)
			if err != nil {
				fmt.Printf("%s\n", err.Error())
				c <- -1
				return
			}

			err = i.tcpHandshake(node, conn)
			if err != nil {
				fmt.Printf("%s\n", err.Error())
				c <- -1
			} else {
				c <- p
			}
			conn.Close()
		}(port)
	}

	newPorts := []int{}
	for i := 0; i < len(ports); i++ {
		port := <-c
		if port != -1 {
			fmt.Printf("tcp port for %s: %d\n", node.Maintainer, port)
			newPorts = append(newPorts, port)
		}
	}

	nodesMutex.Lock()
	node.TCPPorts = newPorts
	if len(node.TCPPorts) > 0 {
		node.LastPing = time.Now().Unix()
	}
	node.TCPStatus = len(node.TCPPorts) > 0
	sort.Stable(nodeSlice(nodes))
	nodesMutex.Unlock()
}

func (i *instance) tcpHandshake(node *toxNode, conn *net.TCPConn) error {
	nodePublicKey := new([crypto.PublicKeySize]byte)
	decPublicKey, err := hex.DecodeString(node.PublicKey)
	if err != nil {
		return err
	}
	copy(nodePublicKey[:], decPublicKey)

	relayConn, err := relay.NewConnection()
	if err != nil {
		return err
	}

	req, err := relayConn.StartHandshake()
	if err != nil {
		return err
	}

	reqBytes, err := req.MarshalBinary()
	if err != nil {
		return err
	}

	encryptedReqBytes, nonce, err := i.Ident.EncryptBlob(reqBytes, nodePublicKey)
	if err != nil {
		return err
	}

	reqPacket := &relay.HandshakeRequestPacket{
		PublicKey: i.Ident.PublicKey,
		Nonce:     nonce,
		Payload:   encryptedReqBytes,
	}

	reqPacketBytes, err := reqPacket.MarshalBinary()
	if err != nil {
		return err
	}

	conn.SetDeadline(time.Now().Add(2 * time.Second))
	_, err = conn.Write(reqPacketBytes)
	if err != nil {
		return err
	}

	buffer := make([]byte, 96)
	left := len(buffer)
	for left > 0 {
		read, readErr := conn.Read(buffer[len(buffer)-left:])
		if readErr != nil {
			return readErr
		}
		left -= read
	}

	res := relay.HandshakeResponsePacket{}
	err = res.UnmarshalBinary(buffer)
	if err != nil {
		return err
	}

	decryptedBytes, err := i.Ident.DecryptBlob(res.Payload, nodePublicKey, res.Nonce)
	if err != nil {
		return err
	}

	resPacket := &relay.HandshakePayload{}
	err = resPacket.UnmarshalBinary(decryptedBytes)
	if err != nil {
		return err
	}

	return relayConn.EndHandshake(resPacket)
}

func (i *instance) getNodes(node *toxNode) (*ping.Ping, error) {
	nodePublicKey := new([crypto.PublicKeySize]byte)
	decPublicKey, err := hex.DecodeString(node.PublicKey)
	if err != nil {
		return nil, err
	}
	copy(nodePublicKey[:], decPublicKey)

	p, err := ping.NewPing(nodePublicKey)
	if err != nil {
		return nil, err
	}

	packet := &dht.GetNodesPacket{
		PublicKey: i.Ident.PublicKey,
		PingID:    p.ID,
	}

	dhtPacket, err := i.Ident.EncryptPacket(transport.Packet(packet), nodePublicKey)
	if err != nil {
		return nil, err
	}

	payload, err := dhtPacket.MarshalBinary()
	if err != nil {
		return nil, err
	}

	if err := i.sendToUDP(payload, node); err != nil {
		return nil, err
	}

	return p, nil
}

func (i *instance) getBootstrapInfo(node *toxNode) error {
	packet, err := bootstrap.ConstructPacket(&bootstrap.InfoRequestPacket{})
	if err != nil {
		return err
	}

	payload, err := packet.MarshalBinary()
	if err != nil {
		return err
	}

	return i.sendToUDP(payload, node)
}

func (i *instance) sendToUDP(data []byte, node *toxNode) error {
	ip, err := getNodeIP(node)
	if err != nil {
		return err
	}

	return i.UDPTransport.Send(
		&transport.Message{
			Data: data,
			Addr: &net.UDPAddr{
				IP:   ip,
				Port: node.Port,
			},
		},
	)
}

func getNodeIP(node *toxNode) (net.IP, error) {
	if node.ip4 != nil {
		return node.ip4, nil
	} else if enableIpv6 && node.ip6 != nil {
		return node.ip6, nil
	}

	return nil, fmt.Errorf("no valid ip found for %s", node.Maintainer)
}

func connectTCP(node *toxNode, port int) (*net.TCPConn, error) {
	ip, err := getNodeIP(node)
	if err != nil {
		return nil, err
	}

	dialer := net.Dialer{}
	dialer.Deadline = time.Now().Add(2 * time.Second)

	tempConn, err := dialer.Dial("tcp", fmt.Sprintf("%s:%d", ip, port))
	if err != nil {
		return nil, err
	}

	conn, ok := tempConn.(*net.TCPConn)
	if !ok {
		return nil, errors.New("not a tcp conn")
	}

	return conn, nil
}

func (i *instance) handleSendNodesPacket(msg *transport.Message) error {
	dhtPacket := &dht.Packet{}
	err := dhtPacket.UnmarshalBinary(msg.Data)
	if err != nil {
		return err
	}

	decryptedPacket, err := i.Ident.DecryptPacket(dhtPacket)
	if err != nil {
		return err
	}

	packet, ok := decryptedPacket.(*dht.SendNodesPacket)
	if !ok {
		return nil
	}

	i.PingsMutex.Lock()
	ping := i.Pings.Find(dhtPacket.SenderPublicKey, packet.PingID, true)
	i.PingsMutex.Unlock()

	if ping == nil {
		err := fmt.Errorf("sendnodes packet from unknown node: %s:%d", msg.Addr.IP, msg.Addr.Port)
		fmt.Printf("error: %s", err.Error())
		return err
	}

	nodesMutex.Lock()
	for _, node := range nodes {
		publicKey, err := hex.DecodeString(node.PublicKey)
		if err != nil {
			continue
		}

		if bytes.Equal(publicKey, dhtPacket.SenderPublicKey[:]) {
			node.UDPStatus = true
			node.LastPing = time.Now().Unix()
			break
		}
	}
	sort.Stable(nodeSlice(nodes))
	nodesMutex.Unlock()

	return nil
}

func handleBootstrapInfoPacket(msg *transport.Message) error {
	bootstrapPacket := &bootstrap.Packet{}
	err := bootstrapPacket.UnmarshalBinary(msg.Data)
	if err != nil {
		return err
	}

	transPacket, err := bootstrap.DestructPacket(bootstrapPacket)
	if err != nil {
		return err
	}

	packet, ok := transPacket.(*bootstrap.InfoResponsePacket)
	if !ok {
		return errors.New("wtf")
	}

	nodesMutex.Lock()
	for _, node := range nodes {
		if (node.ip4 != nil && node.ip4.Equal(msg.Addr.IP)) ||
			(node.ip6 != nil && node.ip6.Equal(msg.Addr.IP)) {
			node.MOTD = packet.MOTD
			node.Version = fmt.Sprintf("%d", packet.Version)
			break
		}
	}
	nodesMutex.Unlock()

	return nil
}
<commit_msg>Fix the TCP check not being able to do IPv6<commit_after>package main

import (
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sort"
	"sync"
	"time"

	"github.com/alexbakker/tox4go/bootstrap"
	"github.com/alexbakker/tox4go/crypto"
	"github.com/alexbakker/tox4go/dht"
	"github.com/alexbakker/tox4go/dht/ping"
	"github.com/alexbakker/tox4go/relay"
	"github.com/alexbakker/tox4go/transport"
	"github.com/didip/tollbooth"
)

const (
	enableIpv6  = true
	probeRate   = 1 * time.Minute
	refreshRate = 5 * time.Minute
)

var (
	lastScan    int64
	lastRefresh int64
	nodes       = []*toxNode{}
	nodesMutex  = sync.Mutex{}
	tcpPorts    = []int{443, 3389, 33445}
)

func main() {
	if parseFlags() {
		return
	}

	// load state if available
	state, err := loadState()
	if err != nil {
		log.Fatalf("error loading state: %s", err.Error())
	}
	lastScan = state.LastScan
	lastRefresh = state.LastRefresh
	nodes = state.Nodes

	if err := loadCountries(); err != nil {
		log.Fatalf("error loading countries.json: %s", err.Error())
	}

	inst, err := NewInstance(":33450")
	if err != nil {
		log.Fatalf("fatal: %s", err.Error())
	}
	inst.UDPTransport.Handle(dht.PacketIDSendNodes, inst.handleSendNodesPacket)
	inst.UDPTransport.Handle(bootstrap.PacketIDBootstrapInfo, handleBootstrapInfoPacket)

	//handle stop signal
	interruptChan := make(chan os.Signal)
	signal.Notify(interruptChan, os.Interrupt)

	//setup http server
	listener, err := net.Listen("tcp4", fmt.Sprintf("localhost:%d", httpListenPort))
	if err != nil {
		log.Fatalf("error in net.Listen: %s", err.Error())
	}
	limiter := tollbooth.NewLimiter(1, 2*time.Second)
	limiter.Methods = []string{"POST"}
	limiter.IPLookups = []string{"X-Forwarded-For", "RemoteAddr", "X-Real-IP"}
	serveMux := http.NewServeMux()
	serveMux.HandleFunc("/", handleHTTPRequest)
	serveMux.Handle("/test", tollbooth.LimitFuncHandler(limiter, handleHTTPRequest))
	serveMux.HandleFunc("/json", handleJSONRequest)
	go func() {
		err := http.Serve(listener, serveMux)
		if err != nil {
			log.Printf("http server error: %s\n", err.Error())
			interruptChan <- os.Interrupt
		}
	}()

	//listen for tox packets
	go func() {
		err := inst.UDPTransport.Listen()
		if err != nil {
			log.Printf("udp transport error: %s\n", err.Error())
			interruptChan <- os.Interrupt
		}
	}()
	//go tcpTransport.Listen()

	err = refreshNodes()
	if err != nil {
		log.Fatal(err.Error())
	}
	inst.probeNodes()

	probeTicker := time.NewTicker(probeRate)
	refreshTicker := time.NewTicker(refreshRate)
	updateTicker := time.NewTicker(30 * time.Second)
	run := true

	for run {
		select {
		case <-interruptChan:
			fmt.Printf("killing routines\n")
			probeTicker.Stop()
			refreshTicker.Stop()
			updateTicker.Stop()
			inst.UDPTransport.Stop()
			//tcpTransport.Stop()
			listener.Close()
			run = false
		case <-probeTicker.C:
			// we want an empty ping list at the start of every probe
			inst.PingsMutex.Lock()
			inst.Pings.Clear(false)
			inst.PingsMutex.Unlock()

			nodesMutex.Lock()
			err := inst.probeNodes()
			nodesMutex.Unlock()
			if err != nil {
				log.Printf("error while trying to probe nodes: %s", err.Error())
			}
		case <-refreshTicker.C:
			err := refreshNodes()
			if err != nil {
				log.Printf("error while trying to refresh nodes: %s", err.Error())
			}
		case <-updateTicker.C:
			inst.PingsMutex.Lock()
			inst.Pings.Clear(true)
			inst.PingsMutex.Unlock()

			nodesMutex.Lock()
			for _, node := range nodes {
				if time.Now().Sub(time.Unix(node.LastPing, 0)) > time.Minute*2 {
					node.UDPStatus = false
				}
			}
			sort.Stable(nodeSlice(nodes))

			state := getState()
			saveState(state)
			nodesMutex.Unlock()
		}
	}
}

func refreshNodes() error {
	parsedNodes, err := parseNodes()
	if err != nil {
		return err
	}

	nodesMutex.Lock()
	for _, freshNode := range parsedNodes {
		for _, node := range nodes {
			if freshNode.PublicKey == node.PublicKey {
				freshNode.LastPing = node.LastPing
				freshNode.UDPStatus = node.UDPStatus
				freshNode.TCPStatus = node.TCPStatus
				freshNode.TCPPorts = node.TCPPorts
				freshNode.MOTD = node.MOTD
				freshNode.Version = node.Version
				break
			}
		}
	}
	nodes = parsedNodes
	sort.Stable(nodeSlice(nodes))
	nodesMutex.Unlock()

	lastRefresh = time.Now().Unix()
	return nil
}

func (i *instance) probeNodes() error {
	for _, node := range nodes {
		err := i.getBootstrapInfo(node)
		if err != nil {
			fmt.Println(err.Error())
		}

		p, err := i.getNodes(node)
		if err != nil {
			fmt.Println(err.Error())
		} else {
			i.PingsMutex.Lock()
			err = i.Pings.Add(p)
			if err != nil {
				fmt.Println(err.Error())
			}
			i.PingsMutex.Unlock()
		}

		ports := tcpPorts
		exists := false
		for _, i := range ports {
			if i == node.Port {
				exists = true
			}
		}
		if !exists {
			ports = append(ports, node.Port)
		}

		go i.probeNodeTCPPorts(node, ports)
	}

	lastScan = time.Now().Unix()
	return nil
}

func (i *instance) probeNodeTCPPorts(node *toxNode, ports []int) {
	c := make(chan int)
	for _, port := range ports {
		go func(p int) {
			conn, err := connectTCP(node, p)
			if err != nil {
				fmt.Printf("%s\n", err.Error())
				c <- -1
				return
			}

			err = i.tcpHandshake(node, conn)
			if err != nil {
				fmt.Printf("%s\n", err.Error())
				c <- -1
			} else {
				c <- p
			}
			conn.Close()
		}(port)
	}

	newPorts := []int{}
	for i := 0; i < len(ports); i++ {
		port := <-c
		if port != -1 {
			fmt.Printf("tcp port for %s: %d\n", node.Maintainer, port)
			newPorts = append(newPorts, port)
		}
	}

	nodesMutex.Lock()
	node.TCPPorts = newPorts
	if len(node.TCPPorts) > 0 {
		node.LastPing = time.Now().Unix()
	}
	node.TCPStatus = len(node.TCPPorts) > 0
	sort.Stable(nodeSlice(nodes))
	nodesMutex.Unlock()
}

func (i *instance) tcpHandshake(node *toxNode, conn *net.TCPConn) error {
	nodePublicKey := new([crypto.PublicKeySize]byte)
	decPublicKey, err := hex.DecodeString(node.PublicKey)
	if err != nil {
		return err
	}
	copy(nodePublicKey[:], decPublicKey)

	relayConn, err := relay.NewConnection()
	if err != nil {
		return err
	}

	req, err := relayConn.StartHandshake()
	if err != nil {
		return err
	}

	reqBytes, err := req.MarshalBinary()
	if err != nil {
		return err
	}

	encryptedReqBytes, nonce, err := i.Ident.EncryptBlob(reqBytes, nodePublicKey)
	if err != nil {
		return err
	}

	reqPacket := &relay.HandshakeRequestPacket{
		PublicKey: i.Ident.PublicKey,
		Nonce:     nonce,
		Payload:   encryptedReqBytes,
	}

	reqPacketBytes, err := reqPacket.MarshalBinary()
	if err != nil {
		return err
	}

	conn.SetDeadline(time.Now().Add(2 * time.Second))
	_, err = conn.Write(reqPacketBytes)
	if err != nil {
		return err
	}

	buffer := make([]byte, 96)
	left := len(buffer)
	for left > 0 {
		read, readErr := conn.Read(buffer[len(buffer)-left:])
		if readErr != nil {
			return readErr
		}
		left -= read
	}

	res := relay.HandshakeResponsePacket{}
	err = res.UnmarshalBinary(buffer)
	if err != nil {
		return err
	}

	decryptedBytes, err := i.Ident.DecryptBlob(res.Payload, nodePublicKey, res.Nonce)
	if err != nil {
		return err
	}

	resPacket := &relay.HandshakePayload{}
	err = resPacket.UnmarshalBinary(decryptedBytes)
	if err != nil {
		return err
	}

	return relayConn.EndHandshake(resPacket)
}

func (i *instance) getNodes(node *toxNode) (*ping.Ping, error) {
	nodePublicKey := new([crypto.PublicKeySize]byte)
	decPublicKey, err := hex.DecodeString(node.PublicKey)
	if err != nil {
		return nil, err
	}
	copy(nodePublicKey[:], decPublicKey)

	p, err := ping.NewPing(nodePublicKey)
	if err != nil {
		return nil, err
	}

	packet := &dht.GetNodesPacket{
		PublicKey: i.Ident.PublicKey,
		PingID:    p.ID,
	}

	dhtPacket, err := i.Ident.EncryptPacket(transport.Packet(packet), nodePublicKey)
	if err != nil {
		return nil, err
	}

	payload, err := dhtPacket.MarshalBinary()
	if err != nil {
		return nil, err
	}

	if err := i.sendToUDP(payload, node); err != nil {
		return nil, err
	}

	return p, nil
}

func (i *instance) getBootstrapInfo(node *toxNode) error {
	packet, err := bootstrap.ConstructPacket(&bootstrap.InfoRequestPacket{})
	if err != nil {
		return err
	}

	payload, err := packet.MarshalBinary()
	if err != nil {
		return err
	}

	return i.sendToUDP(payload, node)
}

func (i *instance) sendToUDP(data []byte, node *toxNode) error {
	ip, err := getNodeIP(node)
	if err != nil {
		return err
	}

	return i.UDPTransport.Send(
		&transport.Message{
			Data: data,
			Addr: &net.UDPAddr{
				IP:   ip,
				Port: node.Port,
			},
		},
	)
}

func getNodeIP(node *toxNode) (net.IP, error) {
	if node.ip4 != nil {
		return node.ip4, nil
	} else if enableIpv6 && node.ip6 != nil {
		return node.ip6, nil
	}

	return nil, fmt.Errorf("no valid ip found for %s", node.Maintainer)
}

func connectTCP(node *toxNode, port int) (*net.TCPConn, error) {
	ip, err := getNodeIP(node)
	if err != nil {
		return nil, err
	}

	dialer := net.Dialer{}
	dialer.Deadline = time.Now().Add(2 * time.Second)
	fmtIp := ""
	if (ip.To4() == nil) {
		// Wrap IPv6 in brackets
		fmtIp = fmt.Sprintf("[%s]:%d", ip, port)
	} else {
		fmtIp = fmt.Sprintf("%s:%d", ip, port)
	}

	tempConn, err := dialer.Dial("tcp", fmtIp)
	if err != nil {
		return nil, err
	}

	conn, ok := tempConn.(*net.TCPConn)
	if !ok {
		return nil, errors.New("not a tcp conn")
	}

	return conn, nil
}

func (i *instance) handleSendNodesPacket(msg *transport.Message) error {
	dhtPacket := &dht.Packet{}
	err := dhtPacket.UnmarshalBinary(msg.Data)
	if err != nil {
		return err
	}

	decryptedPacket, err := i.Ident.DecryptPacket(dhtPacket)
	if err != nil {
		return err
	}

	packet, ok := decryptedPacket.(*dht.SendNodesPacket)
	if !ok {
		return nil
	}

	i.PingsMutex.Lock()
	ping := i.Pings.Find(dhtPacket.SenderPublicKey, packet.PingID, true)
	i.PingsMutex.Unlock()

	if ping == nil {
		err := fmt.Errorf("sendnodes packet from unknown node: %s:%d", msg.Addr.IP, msg.Addr.Port)
		fmt.Printf("error: %s", err.Error())
		return err
	}

	nodesMutex.Lock()
	for _, node := range nodes {
		publicKey, err := hex.DecodeString(node.PublicKey)
		if err != nil {
			continue
		}

		if bytes.Equal(publicKey, dhtPacket.SenderPublicKey[:]) {
			node.UDPStatus = true
			node.LastPing = time.Now().Unix()
			break
		}
	}
	sort.Stable(nodeSlice(nodes))
	nodesMutex.Unlock()

	return nil
}

func handleBootstrapInfoPacket(msg *transport.Message) error {
	bootstrapPacket := &bootstrap.Packet{}
	err := bootstrapPacket.UnmarshalBinary(msg.Data)
	if err != nil {
		return err
	}

	transPacket, err := bootstrap.DestructPacket(bootstrapPacket)
	if err != nil {
		return err
	}

	packet, ok := transPacket.(*bootstrap.InfoResponsePacket)
	if !ok {
		return errors.New("wtf")
	}

	nodesMutex.Lock()
	for _, node := range nodes {
		if (node.ip4 != nil && node.ip4.Equal(msg.Addr.IP)) ||
			(node.ip6 != nil && node.ip6.Equal(msg.Addr.IP)) {
			node.MOTD = packet.MOTD
			node.Version = fmt.Sprintf("%d", packet.Version)
			break
		}
	}
	nodesMutex.Unlock()

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	sscore "github.com/shadowsocks/go-shadowsocks2/core"
	vcore "v2ray.com/core"
	vproxyman "v2ray.com/core/app/proxyman"
	vrouting "v2ray.com/core/features/routing"

	"github.com/eycorsican/go-tun2socks/core"
	"github.com/eycorsican/go-tun2socks/proxy/echo"
	"github.com/eycorsican/go-tun2socks/proxy/shadowsocks"
	"github.com/eycorsican/go-tun2socks/proxy/socks"
	"github.com/eycorsican/go-tun2socks/proxy/v2ray"
	"github.com/eycorsican/go-tun2socks/tun"

	_ "github.com/eycorsican/go-tun2socks/cmd/tun2socks/features"
)

const (
	MTU = 1500
)

func main() {
	tunName := flag.String("tunName", "tun1", "TUN interface name")
	tunAddr := flag.String("tunAddr", "240.0.0.2", "TUN interface address")
	tunGw := flag.String("tunGw", "240.0.0.1", "TUN interface gateway")
	gateway := flag.String("gateway", "", "The gateway adrress of your default network, set this to enable dynamic routing, and root/admin privileges may also required for using dynamic routing (V2Ray only)")
	tunMask := flag.String("tunMask", "255.255.255.0", "TUN interface netmask")
	dnsServer := flag.String("dnsServer", "114.114.114.114,223.5.5.5", "DNS resolvers for TUN interface. (Only take effect on Windows)")
	proxyType := flag.String("proxyType", "socks", "Proxy handler type: socks, shadowsocks, v2ray")
	vconfig := flag.String("vconfig", "config.json", "Config file for v2ray, in JSON format, and note that routing in v2ray could not violate routes in the routing table")
	sniffingType := flag.String("sniffingType", "http,tls", "Enable domain sniffing for specific kind of traffic in v2ray")
	proxyServer := flag.String("proxyServer", "1.2.3.4:1087", "Proxy server address (host:port) for socks and Shadowsocks proxies")
	proxyCipher := flag.String("proxyCipher", "AEAD_CHACHA20_POLY1305", "Cipher used for Shadowsocks proxy, available ciphers: "+strings.Join(sscore.ListCipher(), " "))
	proxyPassword := flag.String("proxyPassword", "", "Password used for Shadowsocks proxy")
	delayICMP := flag.Int("delayICMP", 10, "Delay ICMP packets for a short period of time, in milliseconds")
	udpTimeout := flag.Duration("udpTimeout", 1*time.Minute, "Set timeout for UDP proxy connections in socks and Shadowsocks")

	flag.Parse()

	parts := strings.Split(*proxyServer, ":")
	if len(parts) != 2 {
		log.Fatal("invalid server address")
	}
	proxyAddr := parts[0]
	port, err := strconv.Atoi(parts[1])
	if err != nil {
		log.Fatal("invalid server port")
	}
	proxyPort := uint16(port)

	// Open the tun device.
	dnsServers := strings.Split(*dnsServer, ",")
	tunDev, err := tun.OpenTunDevice(*tunName, *tunAddr, *tunGw, *tunMask, dnsServers)
	if err != nil {
		log.Fatalf("failed to open tun device: %v", err)
	}

	// Setup TCP/IP stack.
	lwipWriter := core.NewLWIPStack().(io.Writer)

	// Wrap a writer to delay ICMP packets if delay time is not zero.
	if *delayICMP > 0 {
		lwipWriter = &icmpDelayedWriter{writer: lwipWriter, delay: *delayICMP}
	}

	// Register TCP and UDP handlers to handle accepted connections.
	switch *proxyType {
	case "echo":
		core.RegisterTCPConnectionHandler(echo.NewTCPHandler())
		core.RegisterUDPConnectionHandler(echo.NewUDPHandler())
		break
	case "socks":
		core.RegisterTCPConnectionHandler(socks.NewTCPHandler(proxyAddr, proxyPort))
		core.RegisterUDPConnectionHandler(socks.NewUDPHandler(proxyAddr, proxyPort, *udpTimeout))
		break
	case "shadowsocks":
		if *proxyCipher == "" || *proxyPassword == "" {
			log.Fatal("invalid cipher or password")
		}
		log.Printf("creat Shadowsocks handler: %v:%v@%v:%v", *proxyCipher, *proxyPassword, proxyAddr, proxyPort)
		core.RegisterTCPConnectionHandler(shadowsocks.NewTCPHandler(net.JoinHostPort(proxyAddr, strconv.Itoa(int(proxyPort))), *proxyCipher, *proxyPassword))
		core.RegisterUDPConnectionHandler(shadowsocks.NewUDPHandler(net.JoinHostPort(proxyAddr, strconv.Itoa(int(proxyPort))), *proxyCipher, *proxyPassword, *udpTimeout))
		break
	case "v2ray":
		configBytes, err := ioutil.ReadFile(*vconfig)
		if err != nil {
			log.Fatal("invalid vconfig file")
		}
		var validSniffings []string
		sniffings := strings.Split(*sniffingType, ",")
		for _, s := range sniffings {
			if s == "http" || s == "tls" {
				validSniffings = append(validSniffings, s)
			}
		}

		v, err := vcore.StartInstance("json", configBytes)
		if err != nil {
			log.Fatal("start V instance failed: %v", err)
		}

		// Wrap a writer for adding routes according to V2Ray's routing results if dynamic routing is enabled.
		if *gateway != "" {
			router := v.GetFeature(vrouting.RouterType()).(vrouting.Router)
			lwipWriter = &routingAwareWriter{writer: lwipWriter, router: router, gateway: *gateway}
		}

		sniffingConfig := &vproxyman.SniffingConfig{
			Enabled:             true,
			DestinationOverride: validSniffings,
		}
		if len(validSniffings) == 0 {
			sniffingConfig.Enabled = false
		}

		ctx := vproxyman.ContextWithSniffingConfig(context.Background(), sniffingConfig)

		vhandler := v2ray.NewHandler(ctx, v)
		core.RegisterTCPConnectionHandler(vhandler)
		core.RegisterUDPConnectionHandler(vhandler)
		break
	default:
		log.Fatal("unsupported proxy type")
	}

	// Register an output function to write packets output from lwip stack to tun
	// device, output function should be set before input any packets.
	core.RegisterOutputFn(func(data []byte) (int, error) {
		return tunDev.Write(data)
	})

	// Copy packets from tun device to lwip stack.
	go func() {
		_, err := io.CopyBuffer(lwipWriter, tunDev, make([]byte, MTU))
		if err != nil {
			log.Fatal("copying data failed: %v", err)
		}
	}()

	log.Printf("running tun2socks")

	osSignals := make(chan os.Signal, 1)
	signal.Notify(osSignals, os.Interrupt, os.Kill, syscall.SIGTERM, syscall.SIGHUP)
	<-osSignals
}
<commit_msg>do not log proxy info<commit_after>package main

import (
	"context"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	sscore "github.com/shadowsocks/go-shadowsocks2/core"
	vcore "v2ray.com/core"
	vproxyman "v2ray.com/core/app/proxyman"
	vrouting "v2ray.com/core/features/routing"

	"github.com/eycorsican/go-tun2socks/core"
	"github.com/eycorsican/go-tun2socks/proxy/echo"
	"github.com/eycorsican/go-tun2socks/proxy/shadowsocks"
	"github.com/eycorsican/go-tun2socks/proxy/socks"
	"github.com/eycorsican/go-tun2socks/proxy/v2ray"
	"github.com/eycorsican/go-tun2socks/tun"

	_ "github.com/eycorsican/go-tun2socks/cmd/tun2socks/features"
)

const (
	MTU = 1500
)

func main() {
	tunName := flag.String("tunName", "tun1", "TUN interface name")
	tunAddr := flag.String("tunAddr", "240.0.0.2", "TUN interface address")
	tunGw := flag.String("tunGw", "240.0.0.1", "TUN interface gateway")
	gateway := flag.String("gateway", "", "The gateway adrress of your default network, set this to enable dynamic routing, and root/admin privileges may also required for using dynamic routing (V2Ray only)")
	tunMask := flag.String("tunMask", "255.255.255.0", "TUN interface netmask")
	dnsServer := flag.String("dnsServer", "114.114.114.114,223.5.5.5", "DNS resolvers for TUN interface. (Only take effect on Windows)")
	proxyType := flag.String("proxyType", "socks", "Proxy handler type: socks, shadowsocks, v2ray")
	vconfig := flag.String("vconfig", "config.json", "Config file for v2ray, in JSON format, and note that routing in v2ray could not violate routes in the routing table")
	sniffingType := flag.String("sniffingType", "http,tls", "Enable domain sniffing for specific kind of traffic in v2ray")
	proxyServer := flag.String("proxyServer", "1.2.3.4:1087", "Proxy server address (host:port) for socks and Shadowsocks proxies")
	proxyCipher := flag.String("proxyCipher", "AEAD_CHACHA20_POLY1305", "Cipher used for Shadowsocks proxy, available ciphers: "+strings.Join(sscore.ListCipher(), " "))
	proxyPassword := flag.String("proxyPassword", "", "Password used for Shadowsocks proxy")
	delayICMP := flag.Int("delayICMP", 10, "Delay ICMP packets for a short period of time, in milliseconds")
	udpTimeout := flag.Duration("udpTimeout", 1*time.Minute, "Set timeout for UDP proxy connections in socks and Shadowsocks")

	flag.Parse()

	parts := strings.Split(*proxyServer, ":")
	if len(parts) != 2 {
		log.Fatal("invalid server address")
	}
	proxyAddr := parts[0]
	port, err := strconv.Atoi(parts[1])
	if err != nil {
		log.Fatal("invalid server port")
	}
	proxyPort := uint16(port)

	// Open the tun device.
	dnsServers := strings.Split(*dnsServer, ",")
	tunDev, err := tun.OpenTunDevice(*tunName, *tunAddr, *tunGw, *tunMask, dnsServers)
	if err != nil {
		log.Fatalf("failed to open tun device: %v", err)
	}

	// Setup TCP/IP stack.
	lwipWriter := core.NewLWIPStack().(io.Writer)

	// Wrap a writer to delay ICMP packets if delay time is not zero.
	if *delayICMP > 0 {
		lwipWriter = &icmpDelayedWriter{writer: lwipWriter, delay: *delayICMP}
	}

	// Register TCP and UDP handlers to handle accepted connections.
	switch *proxyType {
	case "echo":
		core.RegisterTCPConnectionHandler(echo.NewTCPHandler())
		core.RegisterUDPConnectionHandler(echo.NewUDPHandler())
		break
	case "socks":
		core.RegisterTCPConnectionHandler(socks.NewTCPHandler(proxyAddr, proxyPort))
		core.RegisterUDPConnectionHandler(socks.NewUDPHandler(proxyAddr, proxyPort, *udpTimeout))
		break
	case "shadowsocks":
		if *proxyCipher == "" || *proxyPassword == "" {
			log.Fatal("invalid cipher or password")
		}
		core.RegisterTCPConnectionHandler(shadowsocks.NewTCPHandler(net.JoinHostPort(proxyAddr, strconv.Itoa(int(proxyPort))), *proxyCipher, *proxyPassword))
		core.RegisterUDPConnectionHandler(shadowsocks.NewUDPHandler(net.JoinHostPort(proxyAddr, strconv.Itoa(int(proxyPort))), *proxyCipher, *proxyPassword, *udpTimeout))
		break
	case "v2ray":
		configBytes, err := ioutil.ReadFile(*vconfig)
		if err != nil {
			log.Fatal("invalid vconfig file")
		}
		var validSniffings []string
		sniffings := strings.Split(*sniffingType, ",")
		for _, s := range sniffings {
			if s == "http" || s == "tls" {
				validSniffings = append(validSniffings, s)
			}
		}

		v, err := vcore.StartInstance("json", configBytes)
		if err != nil {
			log.Fatal("start V instance failed: %v", err)
		}

		// Wrap a writer for adding routes according to V2Ray's routing results if dynamic routing is enabled.
		if *gateway != "" {
			router := v.GetFeature(vrouting.RouterType()).(vrouting.Router)
			lwipWriter = &routingAwareWriter{writer: lwipWriter, router: router, gateway: *gateway}
		}

		sniffingConfig := &vproxyman.SniffingConfig{
			Enabled:             true,
			DestinationOverride: validSniffings,
		}
		if len(validSniffings) == 0 {
			sniffingConfig.Enabled = false
		}

		ctx := vproxyman.ContextWithSniffingConfig(context.Background(), sniffingConfig)

		vhandler := v2ray.NewHandler(ctx, v)
		core.RegisterTCPConnectionHandler(vhandler)
		core.RegisterUDPConnectionHandler(vhandler)
		break
	default:
		log.Fatal("unsupported proxy type")
	}

	// Register an output function to write packets output from lwip stack to tun
	// device, output function should be set before input any packets.
	core.RegisterOutputFn(func(data []byte) (int, error) {
		return tunDev.Write(data)
	})

	// Copy packets from tun device to lwip stack.
	go func() {
		_, err := io.CopyBuffer(lwipWriter, tunDev, make([]byte, MTU))
		if err != nil {
			log.Fatal("copying data failed: %v", err)
		}
	}()

	log.Printf("running tun2socks")

	osSignals := make(chan os.Signal, 1)
	signal.Notify(osSignals, os.Interrupt, os.Kill, syscall.SIGTERM, syscall.SIGHUP)
	<-osSignals
}
<|endoftext|>