<commit_before>package server

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"sync"
	"time"

	"go.pedge.io/lion/proto"
	"go.pedge.io/pb/go/google/protobuf"
	"go.pedge.io/proto/rpclog"
	"go.pedge.io/proto/stream"
	"golang.org/x/net/context"

	"github.com/cenkalti/backoff"
	"github.com/gogo/protobuf/proto"
	pfsclient "github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"
)

type objBlockAPIServer struct {
	protorpclog.Logger
	dir         string
	localServer *localBlockAPIServer
	objClient   obj.Client
}

func newObjBlockAPIServer(dir string, objClient obj.Client) (*objBlockAPIServer, error) {
	localServer, err := newLocalBlockAPIServer(dir)
	if err != nil {
		return nil, err
	}
	return &objBlockAPIServer{
		Logger:      protorpclog.NewLogger("pachyderm.pfsclient.objBlockAPIServer"),
		dir:         dir,
		localServer: localServer,
		objClient:   objClient,
	}, nil
}

func newAmazonBlockAPIServer(dir string) (*objBlockAPIServer, error) {
	bucket, err := ioutil.ReadFile("/amazon-secret/bucket")
	if err != nil {
		return nil, err
	}
	id, err := ioutil.ReadFile("/amazon-secret/id")
	if err != nil {
		return nil, err
	}
	secret, err := ioutil.ReadFile("/amazon-secret/secret")
	if err != nil {
		return nil, err
	}
	token, err := ioutil.ReadFile("/amazon-secret/token")
	if err != nil {
		return nil, err
	}
	region, err := ioutil.ReadFile("/amazon-secret/region")
	if err != nil {
		return nil, err
	}
	objClient, err := obj.NewAmazonClient(string(bucket), string(id), string(secret), string(token), string(region))
	if err != nil {
		return nil, err
	}
	return newObjBlockAPIServer(dir, objClient)
}

func newGoogleBlockAPIServer(dir string) (*objBlockAPIServer, error) {
	bucket, err := ioutil.ReadFile("/google-secret/bucket")
	if err != nil {
		return nil, err
	}
	objClient, err := obj.NewGoogleClient(context.Background(), string(bucket))
	if err != nil {
		return nil, err
	}
	return newObjBlockAPIServer(dir, objClient)
}

func (s *objBlockAPIServer) PutBlock(putBlockServer pfsclient.BlockAPI_PutBlockServer) (retErr error) {
	result := &pfsclient.BlockRefs{}
	defer func(start time.Time) { s.Log(nil, result, retErr, time.Since(start)) }(time.Now())
	defer drainBlockServer(putBlockServer)
	putBlockRequest, err := putBlockServer.Recv()
	if err != nil {
		if err != io.EOF {
			return err
		}
		return putBlockServer.SendAndClose(result)
	}
	reader := bufio.NewReader(&putBlockReader{
		server: putBlockServer,
		buffer: bytes.NewBuffer(putBlockRequest.Value),
	})
	var wg sync.WaitGroup
	errCh := make(chan error, 1)
	decoder := json.NewDecoder(reader)
	for {
		blockRef, data, err := readBlock(putBlockRequest.Delimiter, reader, decoder)
		if err != nil {
			return err
		}
		result.BlockRef = append(result.BlockRef, blockRef)
		wg.Add(1)
		go func() {
			defer wg.Done()
			path := s.localServer.blockPath(blockRef.Block)
			// We don't want to overwrite blocks that already exist, since:
			// 1) blocks are content-addressable, so it will be the same block
			// 2) we risk exceeding the object store's rate limit
			if s.objClient.Exists(path) {
				return
			}
			writer, err := s.objClient.Writer(path)
			if err != nil {
				select {
				case errCh <- err:
				default:
				}
				return
			}
			defer func() {
				if err := writer.Close(); err != nil {
					select {
					case errCh <- err:
					default:
					}
					return
				}
			}()
			if _, err := writer.Write(data); err != nil {
				select {
				case errCh <- err:
				default:
				}
			}
		}()
		if (blockRef.Range.Upper - blockRef.Range.Lower) < uint64(blockSize) {
			break
		}
	}
	wg.Wait()
	select {
	case err := <-errCh:
		return err
	default:
	}
	return putBlockServer.SendAndClose(result)
}

func (s *objBlockAPIServer) GetBlock(request *pfsclient.GetBlockRequest, getBlockServer pfsclient.BlockAPI_GetBlockServer) (retErr error) {
	defer func(start time.Time) { s.Log(request, nil, retErr, time.Since(start)) }(time.Now())
	var reader io.ReadCloser
	var err error
	backoff.RetryNotify(func() error {
		reader, err = s.objClient.Reader(s.localServer.blockPath(request.Block), request.OffsetBytes, request.SizeBytes)
		if err != nil && s.objClient.IsRetryable(err) {
			return err
		}
		return nil
	}, obj.NewExponentialBackOffConfig(), func(err error, d time.Duration) {
		protolion.Infof("Error creating reader; retrying in %s: %#v", d, obj.RetryError{
			Err:               err.Error(),
			TimeTillNextRetry: d.String(),
		})
	})
	if err != nil {
		return err
	}
	defer func() {
		if err := reader.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	return protostream.WriteToStreamingBytesServer(reader, getBlockServer)
}

func (s *objBlockAPIServer) DeleteBlock(ctx context.Context, request *pfsclient.DeleteBlockRequest) (response *google_protobuf.Empty, retErr error) {
	defer func(start time.Time) { s.Log(request, response, retErr, time.Since(start)) }(time.Now())
	if err := s.objClient.Delete(s.localServer.blockPath(request.Block)); err != nil && !s.objClient.IsNotExist(err) {
		return nil, err
	}
	return google_protobuf.EmptyInstance, nil
}

func (s *objBlockAPIServer) InspectBlock(ctx context.Context, request *pfsclient.InspectBlockRequest) (response *pfsclient.BlockInfo, retErr error) {
	return nil, fmt.Errorf("not implemented")
}

func (s *objBlockAPIServer) ListBlock(ctx context.Context, request *pfsclient.ListBlockRequest) (response *pfsclient.BlockInfos, retErr error) {
	defer func(start time.Time) { s.Log(request, response, retErr, time.Since(start)) }(time.Now())
	return nil, fmt.Errorf("not implemented")
}

func (s *objBlockAPIServer) CreateDiff(ctx context.Context, request *pfsclient.DiffInfo) (response *google_protobuf.Empty, retErr error) {
	defer func(start time.Time) { s.Log(request, response, retErr, time.Since(start)) }(time.Now())
	data, err := proto.Marshal(request)
	if err != nil {
		return nil, err
	}
	writer, err := s.objClient.Writer(s.localServer.diffPath(request.Diff))
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := writer.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	if _, err := writer.Write(data); err != nil {
		return nil, err
	}
	return google_protobuf.EmptyInstance, nil
}

func (s *objBlockAPIServer) InspectDiff(ctx context.Context, request *pfsclient.InspectDiffRequest) (response *pfsclient.DiffInfo, retErr error) {
	defer func(start time.Time) { s.Log(request, response, retErr, time.Since(start)) }(time.Now())
	return s.readDiff(request.Diff)
}

func (s *objBlockAPIServer) ListDiff(request *pfsclient.ListDiffRequest, listDiffServer pfsclient.BlockAPI_ListDiffServer) (retErr error) {
	defer func(start time.Time) { s.Log(request, nil, retErr, time.Since(start)) }(time.Now())
	if err := s.objClient.Walk(s.localServer.diffDir(), func(path string) error {
		diff := s.localServer.pathToDiff(path)
		if diff == nil {
			return fmt.Errorf("couldn't parse %s", path)
		}
		if diff.Shard == request.Shard {
			diffInfo, err := s.readDiff(diff)
			if err != nil {
				return err
			}
			if err := listDiffServer.Send(diffInfo); err != nil {
				return err
			}
		}
		return nil
	}); err != nil {
		return err
	}
	return nil
}

func (s *objBlockAPIServer) DeleteDiff(ctx context.Context, request *pfsclient.DeleteDiffRequest) (response *google_protobuf.Empty, retErr error) {
	defer func(start time.Time) { s.Log(request, response, retErr, time.Since(start)) }(time.Now())
	if err := s.objClient.Delete(s.localServer.diffPath(request.Diff)); err != nil && !s.objClient.IsNotExist(err) {
		return nil, err
	}
	return google_protobuf.EmptyInstance, nil
}

func (s *objBlockAPIServer) readDiff(diff *pfsclient.Diff) (*pfsclient.DiffInfo, error) {
	reader, err := s.objClient.Reader(s.localServer.diffPath(diff), 0, 0)
	if err != nil {
		return nil, err
	}
	data, err := ioutil.ReadAll(reader)
	result := &pfsclient.DiffInfo{}
	if err := proto.Unmarshal(data, result); err != nil {
		return nil, err
	}
	return result, nil
}
<commit_msg>Add backoff to DeleteBlock<commit_after>package server

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"sync"
	"time"

	"go.pedge.io/lion/proto"
	"go.pedge.io/pb/go/google/protobuf"
	"go.pedge.io/proto/rpclog"
	"go.pedge.io/proto/stream"
	"golang.org/x/net/context"

	"github.com/cenkalti/backoff"
	"github.com/gogo/protobuf/proto"
	pfsclient "github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"
)

type objBlockAPIServer struct {
	protorpclog.Logger
	dir         string
	localServer *localBlockAPIServer
	objClient   obj.Client
}

func newObjBlockAPIServer(dir string, objClient obj.Client) (*objBlockAPIServer, error) {
	localServer, err := newLocalBlockAPIServer(dir)
	if err != nil {
		return nil, err
	}
	return &objBlockAPIServer{
		Logger:      protorpclog.NewLogger("pachyderm.pfsclient.objBlockAPIServer"),
		dir:         dir,
		localServer: localServer,
		objClient:   objClient,
	}, nil
}

func newAmazonBlockAPIServer(dir string) (*objBlockAPIServer, error) {
	bucket, err := ioutil.ReadFile("/amazon-secret/bucket")
	if err != nil {
		return nil, err
	}
	id, err := ioutil.ReadFile("/amazon-secret/id")
	if err != nil {
		return nil, err
	}
	secret, err := ioutil.ReadFile("/amazon-secret/secret")
	if err != nil {
		return nil, err
	}
	token, err := ioutil.ReadFile("/amazon-secret/token")
	if err != nil {
		return nil, err
	}
	region, err := ioutil.ReadFile("/amazon-secret/region")
	if err != nil {
		return nil, err
	}
	objClient, err := obj.NewAmazonClient(string(bucket), string(id), string(secret), string(token), string(region))
	if err != nil {
		return nil, err
	}
	return newObjBlockAPIServer(dir, objClient)
}

func newGoogleBlockAPIServer(dir string) (*objBlockAPIServer, error) {
	bucket, err := ioutil.ReadFile("/google-secret/bucket")
	if err != nil {
		return nil, err
	}
	objClient, err := obj.NewGoogleClient(context.Background(), string(bucket))
	if err != nil {
		return nil, err
	}
	return newObjBlockAPIServer(dir, objClient)
}

func (s *objBlockAPIServer) PutBlock(putBlockServer pfsclient.BlockAPI_PutBlockServer) (retErr error) {
	result := &pfsclient.BlockRefs{}
	defer func(start time.Time) { s.Log(nil, result, retErr, time.Since(start)) }(time.Now())
	defer drainBlockServer(putBlockServer)
	putBlockRequest, err := putBlockServer.Recv()
	if err != nil {
		if err != io.EOF {
			return err
		}
		return putBlockServer.SendAndClose(result)
	}
	reader := bufio.NewReader(&putBlockReader{
		server: putBlockServer,
		buffer: bytes.NewBuffer(putBlockRequest.Value),
	})
	var wg sync.WaitGroup
	errCh := make(chan error, 1)
	decoder := json.NewDecoder(reader)
	for {
		blockRef, data, err := readBlock(putBlockRequest.Delimiter, reader, decoder)
		if err != nil {
			return err
		}
		result.BlockRef = append(result.BlockRef, blockRef)
		wg.Add(1)
		go func() {
			defer wg.Done()
			path := s.localServer.blockPath(blockRef.Block)
			// We don't want to overwrite blocks that already exist, since:
			// 1) blocks are content-addressable, so it will be the same block
			// 2) we risk exceeding the object store's rate limit
			if s.objClient.Exists(path) {
				return
			}
			writer, err := s.objClient.Writer(path)
			if err != nil {
				select {
				case errCh <- err:
				default:
				}
				return
			}
			defer func() {
				if err := writer.Close(); err != nil {
					select {
					case errCh <- err:
					default:
					}
					return
				}
			}()
			if _, err := writer.Write(data); err != nil {
				select {
				case errCh <- err:
				default:
				}
			}
		}()
		if (blockRef.Range.Upper - blockRef.Range.Lower) < uint64(blockSize) {
			break
		}
	}
	wg.Wait()
	select {
	case err := <-errCh:
		return err
	default:
	}
	return putBlockServer.SendAndClose(result)
}

func (s *objBlockAPIServer) GetBlock(request *pfsclient.GetBlockRequest, getBlockServer pfsclient.BlockAPI_GetBlockServer) (retErr error) {
	defer func(start time.Time) { s.Log(request, nil, retErr, time.Since(start)) }(time.Now())
	var reader io.ReadCloser
	var err error
	backoff.RetryNotify(func() error {
		reader, err = s.objClient.Reader(s.localServer.blockPath(request.Block), request.OffsetBytes, request.SizeBytes)
		if err != nil && s.objClient.IsRetryable(err) {
			return err
		}
		return nil
	}, obj.NewExponentialBackOffConfig(), func(err error, d time.Duration) {
		protolion.Infof("Error creating reader; retrying in %s: %#v", d, obj.RetryError{
			Err:               err.Error(),
			TimeTillNextRetry: d.String(),
		})
	})
	if err != nil {
		return err
	}
	defer func() {
		if err := reader.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	return protostream.WriteToStreamingBytesServer(reader, getBlockServer)
}

func (s *objBlockAPIServer) DeleteBlock(ctx context.Context, request *pfsclient.DeleteBlockRequest) (response *google_protobuf.Empty, retErr error) {
	defer func(start time.Time) { s.Log(request, response, retErr, time.Since(start)) }(time.Now())
	backoff.RetryNotify(func() error {
		if err := s.objClient.Delete(s.localServer.blockPath(request.Block)); err != nil && !s.objClient.IsNotExist(err) {
			return err
		}
		return nil
	}, obj.NewExponentialBackOffConfig(), func(err error, d time.Duration) {
		protolion.Infof("Error deleting block; retrying in %s: %#v", d, obj.RetryError{
			Err:               err.Error(),
			TimeTillNextRetry: d.String(),
		})
	})
	return google_protobuf.EmptyInstance, nil
}

func (s *objBlockAPIServer) InspectBlock(ctx context.Context, request *pfsclient.InspectBlockRequest) (response *pfsclient.BlockInfo, retErr error) {
	return nil, fmt.Errorf("not implemented")
}

func (s *objBlockAPIServer) ListBlock(ctx context.Context, request *pfsclient.ListBlockRequest) (response *pfsclient.BlockInfos, retErr error) {
	defer func(start time.Time) { s.Log(request, response, retErr, time.Since(start)) }(time.Now())
	return nil, fmt.Errorf("not implemented")
}

func (s *objBlockAPIServer) CreateDiff(ctx context.Context, request *pfsclient.DiffInfo) (response *google_protobuf.Empty, retErr error) {
	defer func(start time.Time) { s.Log(request, response, retErr, time.Since(start)) }(time.Now())
	data, err := proto.Marshal(request)
	if err != nil {
		return nil, err
	}
	writer, err := s.objClient.Writer(s.localServer.diffPath(request.Diff))
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := writer.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	if _, err := writer.Write(data); err != nil {
		return nil, err
	}
	return google_protobuf.EmptyInstance, nil
}

func (s *objBlockAPIServer) InspectDiff(ctx context.Context, request *pfsclient.InspectDiffRequest) (response *pfsclient.DiffInfo, retErr error) {
	defer func(start time.Time) { s.Log(request, response, retErr, time.Since(start)) }(time.Now())
	return s.readDiff(request.Diff)
}

func (s *objBlockAPIServer) ListDiff(request *pfsclient.ListDiffRequest, listDiffServer pfsclient.BlockAPI_ListDiffServer) (retErr error) {
	defer func(start time.Time) { s.Log(request, nil, retErr, time.Since(start)) }(time.Now())
	if err := s.objClient.Walk(s.localServer.diffDir(), func(path string) error {
		diff := s.localServer.pathToDiff(path)
		if diff == nil {
			return fmt.Errorf("couldn't parse %s", path)
		}
		if diff.Shard == request.Shard {
			diffInfo, err := s.readDiff(diff)
			if err != nil {
				return err
			}
			if err := listDiffServer.Send(diffInfo); err != nil {
				return err
			}
		}
		return nil
	}); err != nil {
		return err
	}
	return nil
}

func (s *objBlockAPIServer) DeleteDiff(ctx context.Context, request *pfsclient.DeleteDiffRequest) (response *google_protobuf.Empty, retErr error) {
	defer func(start time.Time) { s.Log(request, response, retErr, time.Since(start)) }(time.Now())
	if err := s.objClient.Delete(s.localServer.diffPath(request.Diff)); err != nil && !s.objClient.IsNotExist(err) {
		return nil, err
	}
	return google_protobuf.EmptyInstance, nil
}

func (s *objBlockAPIServer) readDiff(diff *pfsclient.Diff) (*pfsclient.DiffInfo, error) {
	reader, err := s.objClient.Reader(s.localServer.diffPath(diff), 0, 0)
	if err != nil {
		return nil, err
	}
	data, err := ioutil.ReadAll(reader)
	result := &pfsclient.DiffInfo{}
	if err := proto.Unmarshal(data, result); err != nil {
		return nil, err
	}
	return result, nil
}
<|endoftext|>


<commit_before>package httpfwd

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	// "io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"code.google.com/p/gopacket"
	"code.google.com/p/gopacket/layers"
	"code.google.com/p/gopacket/pcap"
	"code.google.com/p/gopacket/tcpassembly"
	"code.google.com/p/gopacket/tcpassembly/tcpreader"
)

type Forwarder struct {
	Sources      AddrList
	Destinations AddrList
	Interfaces   []string
	Headers      map[string]string
	Log          *log.Logger
	Verbose      bool
	Bufsize      int32
	Timeout      time.Duration
	Transport    http.Transport
}

type Options struct {
	Sources      []string `short:"s" long:"src"     description:"Source(s) to wiretap HTTP traffic from" value-name:"HOST[:PORT]" default:"*:80" default-mask:"*:80 by default"`
	Destinations []string `short:"d" long:"dst"     description:"Destination(s) to forward copy of HTTP traffic to" value-name:"HOST[:PORT]" required:"true"`
	Headers      []string `short:"H" long:"header"  description:"Set or replace request header in duplicated traffic" value-name:"LINE"`
	Verbose      bool     `short:"v" long:"verbose" description:"Show extra information, including all request headers"`
}

func NewForwarder(opts Options) *Forwarder {
	sources, err := ResolveAddrPatterns(opts.Sources)
	if err != nil {
		panic(err)
	}

	destinations, err := ResolveAddrPatterns(opts.Destinations)
	if err != nil {
		panic(err)
	}

	headers := make(map[string]string)
	for _, header := range opts.Headers {
		parts := strings.SplitN(header, ":", 2)
		headers[strings.ToLower(parts[0])] = strings.TrimLeft(parts[1], " ")
	}

	return &Forwarder{
		Sources:      sources,
		Destinations: destinations,
		Interfaces:   FindInterfaces(),
		Headers:      headers,
		Log:          log.New(os.Stdout, "", log.LstdFlags),
		Verbose:      opts.Verbose,
		Bufsize:      65535,
		Timeout:      10 * time.Millisecond,
		Transport:    http.Transport{MaxIdleConnsPerHost: 6},
	}
}

func PcapVersion() string {
	return pcap.Version()
}

func (fwd *Forwarder) Start() {
	pool := tcpassembly.NewStreamPool(fwd)
	assembler := tcpassembly.NewAssembler(pool)

	packets := fwd.wiretap()
	ticker := time.Tick(time.Minute)

	if fwd.Verbose {
		fmt.Fprintf(os.Stderr, "Listening on interfaces %s\n", strings.Join(fwd.Interfaces, ", "))
	}

	fmt.Fprintf(os.Stderr, "Wiretapping HTTP traffic to %s and forwarding to %s...\n", fwd.Sources, fwd.Destinations)

	for {
		select {
		case packet := <-packets:
			assembler.Assemble(
				packet.NetworkLayer().NetworkFlow(),
				packet.TransportLayer().(*layers.TCP))
		case <-ticker:
			assembler.FlushOlderThan(time.Now().Add(-2 * time.Minute))
		}
	}
}

func (fwd *Forwarder) New(netFlow, tcpFlow gopacket.Flow) tcpassembly.Stream {
	stream := tcpreader.NewReaderStream()

	go func() {
		buf := bufio.NewReader(&stream)

		for {
			req, err := http.ReadRequest(buf)
			if err == io.EOF {
				return
			} else if err != nil {
				fwd.Log.Println("Error:", err)
			} else {
				body := new(bytes.Buffer)
				// if _, err := io.Copy(body, req.Body); err != nil {
				// 	fwd.Log.Println("Error:", err)
				// }
				go fwd.forwardRequest(&netFlow, req, body)
			}
		}
	}()

	return &stream
}

func (fwd *Forwarder) wiretap() chan gopacket.Packet {
	channel := make(chan gopacket.Packet, 100)
	filter := fwd.Sources.Filter()

	for _, intf := range fwd.Interfaces {
		handle, err := pcap.OpenLive(intf, fwd.Bufsize, fwd.Sources.RequiresPromisc(), fwd.Timeout)
		if err != nil {
			panic(err)
		}

		if err := handle.SetBPFFilter(filter); err != nil {
			panic(err)
		}

		go func(handle *pcap.Handle) {
			defer handle.Close()

			src := gopacket.NewPacketSource(handle, handle.LinkType())
			src.DecodeOptions = gopacket.NoCopy

			for {
				packet, err := src.NextPacket()
				if err == io.EOF {
					return
				} else if err == nil {
					channel <- packet
				}
			}
		}(handle)
	}

	return channel
}

func (fwd *Forwarder) forwardRequest(netFlow *gopacket.Flow, req *http.Request, body *bytes.Buffer) {
	req.URL.Scheme = "http"
	req.URL.Host = req.Host

	originalPeer := netFlow.Src().String()
	originalURL := req.URL.String()

	for key, value := range fwd.Headers {
		if value == "" {
			req.Header.Del(key)
			if key == "user-agent" {
				/* "Use the defaultUserAgent unless the Header contains one,
				   which may be blank to not send the header." */
				req.Header.Set(key, "")
			}
		} else {
			req.Header.Set(key, value)
			if key == "host" {
				req.Host = value
			}
		}
	}

	for _, dst := range fwd.Destinations {
		copy := copyRequest(req)
		copy.URL.Host = dst.String()
		// copy.Body = ioutil.NopCloser(bytes.NewReader(body.Bytes()))
		go fwd.sendRequest(copy, originalPeer, originalURL)
	}
}

func (fwd *Forwarder) sendRequest(req *http.Request, originalPeer, originalURL string) {
	res, err := fwd.Transport.RoundTrip(req)
	if err != nil {
		fwd.Log.Println("Error:", err)
	} else {
		/* "The client must close the response body when finished with it." */
		defer res.Body.Close()
		fwd.Log.Printf("%s %s %s (%s) %d\n", originalPeer, req.Method, originalURL, req.URL.Host, res.StatusCode)
		if fwd.Verbose {
			req.Body = nil
			req.Write(os.Stdout)
		}
	}
}

func copyRequest(req *http.Request) *http.Request {
	url := *req.URL
	copy := *req
	copy.URL = &url
	return &copy
}
<commit_msg>Re-enable copying of request bodies.<commit_after>package httpfwd

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"code.google.com/p/gopacket"
	"code.google.com/p/gopacket/layers"
	"code.google.com/p/gopacket/pcap"
	"code.google.com/p/gopacket/tcpassembly"
	"code.google.com/p/gopacket/tcpassembly/tcpreader"
)

type Forwarder struct {
	Sources      AddrList
	Destinations AddrList
	Interfaces   []string
	Headers      map[string]string
	Log          *log.Logger
	Verbose      bool
	Bufsize      int32
	Timeout      time.Duration
	Transport    http.Transport
}

type Options struct {
	Sources      []string `short:"s" long:"src"     description:"Source(s) to wiretap HTTP traffic from" value-name:"HOST[:PORT]" default:"*:80" default-mask:"*:80 by default"`
	Destinations []string `short:"d" long:"dst"     description:"Destination(s) to forward copy of HTTP traffic to" value-name:"HOST[:PORT]" required:"true"`
	Headers      []string `short:"H" long:"header"  description:"Set or replace request header in duplicated traffic" value-name:"LINE"`
	Verbose      bool     `short:"v" long:"verbose" description:"Show extra information, including all request headers"`
}

func NewForwarder(opts Options) *Forwarder {
	sources, err := ResolveAddrPatterns(opts.Sources)
	if err != nil {
		panic(err)
	}

	destinations, err := ResolveAddrPatterns(opts.Destinations)
	if err != nil {
		panic(err)
	}

	headers := make(map[string]string)
	for _, header := range opts.Headers {
		parts := strings.SplitN(header, ":", 2)
		headers[strings.ToLower(parts[0])] = strings.TrimLeft(parts[1], " ")
	}

	return &Forwarder{
		Sources:      sources,
		Destinations: destinations,
		Interfaces:   FindInterfaces(),
		Headers:      headers,
		Log:          log.New(os.Stdout, "", log.LstdFlags),
		Verbose:      opts.Verbose,
		Bufsize:      65535,
		Timeout:      10 * time.Millisecond,
		Transport:    http.Transport{MaxIdleConnsPerHost: 6},
	}
}

func PcapVersion() string {
	return pcap.Version()
}

func (fwd *Forwarder) Start() {
	pool := tcpassembly.NewStreamPool(fwd)
	assembler := tcpassembly.NewAssembler(pool)

	packets := fwd.wiretap()
	ticker := time.Tick(time.Minute)

	if fwd.Verbose {
		fmt.Fprintf(os.Stderr, "Listening on interfaces %s\n", strings.Join(fwd.Interfaces, ", "))
	}

	fmt.Fprintf(os.Stderr, "Wiretapping HTTP traffic to %s and forwarding to %s...\n", fwd.Sources, fwd.Destinations)

	for {
		select {
		case packet := <-packets:
			assembler.Assemble(
				packet.NetworkLayer().NetworkFlow(),
				packet.TransportLayer().(*layers.TCP))
		case <-ticker:
			assembler.FlushOlderThan(time.Now().Add(-2 * time.Minute))
		}
	}
}

func (fwd *Forwarder) New(netFlow, tcpFlow gopacket.Flow) tcpassembly.Stream {
	stream := tcpreader.NewReaderStream()

	go func() {
		buf := bufio.NewReader(&stream)

		for {
			req, err := http.ReadRequest(buf)
			if err == io.EOF {
				return
			} else if err != nil {
				fwd.Log.Println("Error:", err)
			} else {
				body := new(bytes.Buffer)
				if _, err := io.Copy(body, req.Body); err != nil {
					fwd.Log.Println("Error:", err)
				}
				go fwd.forwardRequest(&netFlow, req, body)
			}
		}
	}()

	return &stream
}

func (fwd *Forwarder) wiretap() chan gopacket.Packet {
	channel := make(chan gopacket.Packet, 100)
	filter := fwd.Sources.Filter()

	for _, intf := range fwd.Interfaces {
		handle, err := pcap.OpenLive(intf, fwd.Bufsize, fwd.Sources.RequiresPromisc(), fwd.Timeout)
		if err != nil {
			panic(err)
		}

		if err := handle.SetBPFFilter(filter); err != nil {
			panic(err)
		}

		go func(handle *pcap.Handle) {
			defer handle.Close()

			src := gopacket.NewPacketSource(handle, handle.LinkType())
			src.DecodeOptions = gopacket.NoCopy

			for {
				packet, err := src.NextPacket()
				if err == io.EOF {
					return
				} else if err == nil {
					channel <- packet
				}
			}
		}(handle)
	}

	return channel
}

func (fwd *Forwarder) forwardRequest(netFlow *gopacket.Flow, req *http.Request, body *bytes.Buffer) {
	req.URL.Scheme = "http"
	req.URL.Host = req.Host

	originalPeer := netFlow.Src().String()
	originalURL := req.URL.String()

	for key, value := range fwd.Headers {
		if value == "" {
			req.Header.Del(key)
			if key == "user-agent" {
				/* "Use the defaultUserAgent unless the Header contains one,
				   which may be blank to not send the header." */
				req.Header.Set(key, "")
			}
		} else {
			req.Header.Set(key, value)
			if key == "host" {
				req.Host = value
			}
		}
	}

	for _, dst := range fwd.Destinations {
		copy := copyRequest(req)
		copy.URL.Host = dst.String()
		copy.Body = ioutil.NopCloser(bytes.NewReader(body.Bytes()))
		go fwd.sendRequest(copy, originalPeer, originalURL)
	}
}

func (fwd *Forwarder) sendRequest(req *http.Request, originalPeer, originalURL string) {
	res, err := fwd.Transport.RoundTrip(req)
	if err != nil {
		fwd.Log.Println("Error:", err)
	} else {
		/* "The client must close the response body when finished with it." */
		defer res.Body.Close()
		fwd.Log.Printf("%s %s %s (%s) %d\n", originalPeer, req.Method, originalURL, req.URL.Host, res.StatusCode)
		if fwd.Verbose {
			req.Body = nil
			req.Write(os.Stdout)
		}
	}
}

func copyRequest(req *http.Request) *http.Request {
	url := *req.URL
	copy := *req
	copy.URL = &url
	return &copy
}
<|endoftext|>


<commit_before>// Package flagconfig provides a flag to specifiy a config file which
// will in turn be used to read unspecified flag values.
package flagconfig

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
)

var configFile = flag.String(
	"c", defaultConfig(), "Config file to read flags from.")

func Usage() {
	Parse()
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	flag.VisitAll(func(f *flag.Flag) {
		fmt.Fprintf(os.Stderr, "  -%s=%s: %s\n", f.Name, f.Value.String(), f.Usage)
	})
}

func defaultConfig() string {
	home := os.Getenv("HOME")
	path := filepath.Join(home, ".config", os.Args[0], "config")
	_, err := os.Open(path)
	if err == nil {
		return path
	}
	path = filepath.Join("/", "etc", "conf.d", os.Args[0])
	_, err = os.Open(path)
	if err == nil {
		return path
	}
	return ""
}

func contains(list []*flag.Flag, f *flag.Flag) bool {
	for _, i := range list {
		if i == f {
			return true
		}
	}
	return false
}

func readConfig() map[string]string {
	bytes, err := ioutil.ReadFile(*configFile)
	if err != nil {
		log.Fatalf("Failed to read config file %s: %s", *configFile, err)
	}
	lines := strings.Split(string(bytes), "\n")
	result := make(map[string]string, len(lines))
	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if trimmed == "" || trimmed[0] == '#' {
			continue
		}
		parts := strings.Split(line, "=")
		if len(parts) != 2 {
			log.Fatalf("Invalid config line: %s", line)
		}
		result[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1])
	}
	return result
}

func Parse() {
	if *configFile == "" {
		return
	}
	config := readConfig()
	explicit := make([]*flag.Flag, 0)
	all := make([]*flag.Flag, 0)
	flag.Visit(func(f *flag.Flag) {
		explicit = append(explicit, f)
	})
	flag.VisitAll(func(f *flag.Flag) {
		all = append(all, f)
		if !contains(explicit, f) {
			val := config[f.Name]
			if val != "" {
				err := f.Value.Set(val)
				if err != nil {
					log.Fatalf("Failed to set flag %s with value %s", f.Name, val)
				}
			}
		}
	})
Outer:
	for name, val := range config {
		for _, f := range all {
			if f.Name == name {
				continue Outer
			}
		}
		log.Fatalf("Unknown flag %s=%s in config file.", name, val)
	}
}
<commit_msg>only use base name of os.Args[0] to look for config<commit_after>// Package flagconfig provides a flag to specifiy a config file which
// will in turn be used to read unspecified flag values.
package flagconfig

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
)

var configFile = flag.String(
	"c", defaultConfig(), "Config file to read flags from.")

func Usage() {
	Parse()
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	flag.VisitAll(func(f *flag.Flag) {
		fmt.Fprintf(os.Stderr, "  -%s=%s: %s\n", f.Name, f.Value.String(), f.Usage)
	})
}

func defaultConfig() string {
	home := os.Getenv("HOME")
	basename := filepath.Base(os.Args[0])
	path := filepath.Join(home, ".config", basename, "config")
	_, err := os.Open(path)
	if err == nil {
		return path
	}
	path = filepath.Join("/", "etc", "conf.d", basename)
	_, err = os.Open(path)
	if err == nil {
		return path
	}
	return ""
}

func contains(list []*flag.Flag, f *flag.Flag) bool {
	for _, i := range list {
		if i == f {
			return true
		}
	}
	return false
}

func readConfig() map[string]string {
	bytes, err := ioutil.ReadFile(*configFile)
	if err != nil {
		log.Fatalf("Failed to read config file %s: %s", *configFile, err)
	}
	lines := strings.Split(string(bytes), "\n")
	result := make(map[string]string, len(lines))
	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if trimmed == "" || trimmed[0] == '#' {
			continue
		}
		parts := strings.Split(line, "=")
		if len(parts) != 2 {
			log.Fatalf("Invalid config line: %s", line)
		}
		result[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1])
	}
	return result
}

func Parse() {
	if *configFile == "" {
		return
	}
	config := readConfig()
	explicit := make([]*flag.Flag, 0)
	all := make([]*flag.Flag, 0)
	flag.Visit(func(f *flag.Flag) {
		explicit = append(explicit, f)
	})
	flag.VisitAll(func(f *flag.Flag) {
		all = append(all, f)
		if !contains(explicit, f) {
			val := config[f.Name]
			if val != "" {
				err := f.Value.Set(val)
				if err != nil {
					log.Fatalf("Failed to set flag %s with value %s", f.Name, val)
				}
			}
		}
	})
Outer:
	for name, val := range config {
		for _, f := range all {
			if f.Name == name {
				continue Outer
			}
		}
		log.Fatalf("Unknown flag %s=%s in config file.", name, val)
	}
}
<|endoftext|>


<commit_before>package dht

import (
	"context"
	"crypto/rand"
	"fmt"
	"strings"
	"sync"
	"time"

	u "github.com/ipfs/go-ipfs-util"
	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/routing"
	"github.com/multiformats/go-multiaddr"
	_ "github.com/multiformats/go-multiaddr-dns"
	"github.com/pkg/errors"
)

var DefaultBootstrapPeers []multiaddr.Multiaddr

func init() {
	for _, s := range []string{
		"/dnsaddr/bootstrap.libp2p.io/ipfs/QmNnooDu7bfjPFoTZYxMNLWUQJyrVwtbZg5gBMjTezGAJN",
		"/dnsaddr/bootstrap.libp2p.io/ipfs/QmQCU2EcMqAqQPR2i9bChDtGNJchTbq5TbXJJ16u19uLTa",
		"/dnsaddr/bootstrap.libp2p.io/ipfs/QmbLHAnMoJPWSCR5Zhtx6BHJX9KiKNN6tpvbUcqanj75Nb",
		"/dnsaddr/bootstrap.libp2p.io/ipfs/QmcZf59bWwK5XFi76CZX8cbJ4BhTzzA3gU1ZjYZcYW3dwt",
		"/ip4/104.131.131.82/tcp/4001/ipfs/QmaCpDMGvV2BGHeYERUEnRQAwe3N8SzbUtfsmvsqQLuvuJ",            // mars.i.ipfs.io
		"/ip4/104.236.179.241/tcp/4001/ipfs/QmSoLPppuBtQSGwKDZT2M73ULpjvfd3aZ6ha4oFGL1KrGM",           // pluto.i.ipfs.io
		"/ip4/128.199.219.111/tcp/4001/ipfs/QmSoLSafTMBsPKadTEgaXctDQVcqN88CNLHXMkTNwMKPnu",           // saturn.i.ipfs.io
		"/ip4/104.236.76.40/tcp/4001/ipfs/QmSoLV4Bbm51jM9C4gDYZQ9Cy3U6aXMJDAbzgu2fzaDs64",             // venus.i.ipfs.io
		"/ip4/178.62.158.247/tcp/4001/ipfs/QmSoLer265NRgSp2LA3dPaeykiS1J6DifTC88f5uVQKNAd",            // earth.i.ipfs.io
		"/ip6/2604:a880:1:20::203:d001/tcp/4001/ipfs/QmSoLPppuBtQSGwKDZT2M73ULpjvfd3aZ6ha4oFGL1KrGM",  // pluto.i.ipfs.io
		"/ip6/2400:6180:0:d0::151:6001/tcp/4001/ipfs/QmSoLSafTMBsPKadTEgaXctDQVcqN88CNLHXMkTNwMKPnu",  // saturn.i.ipfs.io
		"/ip6/2604:a880:800:10::4a:5001/tcp/4001/ipfs/QmSoLV4Bbm51jM9C4gDYZQ9Cy3U6aXMJDAbzgu2fzaDs64", // venus.i.ipfs.io
		"/ip6/2a03:b0c0:0:1010::23:1001/tcp/4001/ipfs/QmSoLer265NRgSp2LA3dPaeykiS1J6DifTC88f5uVQKNAd", // earth.i.ipfs.io
	} {
		ma, err := multiaddr.NewMultiaddr(s)
		if err != nil {
			panic(err)
		}
		DefaultBootstrapPeers = append(DefaultBootstrapPeers, ma)
	}
}

// BootstrapConfig specifies parameters used for bootstrapping the DHT.
type BootstrapConfig struct {
	BucketPeriod             time.Duration // how long to wait for a k-bucket to be queried before doing a random walk on it
	Timeout                  time.Duration // how long to wait for a bootstrap query to run
	RoutingTableScanInterval time.Duration // how often to scan the RT for k-buckets that haven't been queried since the given period
	SelfQueryInterval        time.Duration // how often to query for self
}

var DefaultBootstrapConfig = BootstrapConfig{
	// same as that mentioned in the kad dht paper
	BucketPeriod: 1 * time.Hour,

	// since the default bucket period is 1 hour, a scan interval of 30 minutes sounds reasonable
	RoutingTableScanInterval: 30 * time.Minute,

	Timeout: 10 * time.Second,

	SelfQueryInterval: 1 * time.Hour,
}

// A method in the IpfsRouting interface. It calls BootstrapWithConfig with
// the default bootstrap config.
func (dht *IpfsDHT) Bootstrap(ctx context.Context) error {
	return dht.BootstrapWithConfig(ctx, DefaultBootstrapConfig)
}

// Runs cfg.Queries bootstrap queries every cfg.BucketPeriod.
func (dht *IpfsDHT) BootstrapWithConfig(ctx context.Context, cfg BootstrapConfig) error {
	// we should query for self periodically so we can discover closer peers
	go func() {
		for {
			err := dht.BootstrapSelf(ctx)
			if err != nil {
				logger.Warningf("error bootstrapping while searching for my self (I'm Too Shallow ?): %s", err)
			}
			select {
			case <-time.After(cfg.SelfQueryInterval):
			case <-ctx.Done():
				return
			}
		}
	}()

	// scan the RT table periodically & do a random walk on k-buckets that haven't been queried since the given bucket period
	go func() {
		for {
			err := dht.runBootstrap(ctx, cfg)
			if err != nil {
				logger.Warningf("error bootstrapping: %s", err)
			}
			select {
			case <-time.After(cfg.RoutingTableScanInterval):
			case <-ctx.Done():
				return
			}
		}
	}()
	return nil
}

func newRandomPeerId() peer.ID {
	id := make([]byte, 32) // SHA256 is the default. TODO: Use a more canonical way to generate random IDs.
	rand.Read(id)
	id = u.Hash(id) // TODO: Feed this directly into the multihash instead of hashing it.
	return peer.ID(id)
}

// Traverse the DHT toward the given ID.
func (dht *IpfsDHT) walk(ctx context.Context, target peer.ID) (peer.AddrInfo, error) {
	// TODO: Extract the query action (traversal logic?) inside FindPeer,
	// don't actually call through the FindPeer machinery, which can return
	// things out of the peer store etc.
	return dht.FindPeer(ctx, target)
}

// Traverse the DHT toward a random ID.
func (dht *IpfsDHT) randomWalk(ctx context.Context) error {
	id := newRandomPeerId()
	p, err := dht.walk(ctx, id)
	switch err {
	case routing.ErrNotFound:
		return nil
	case nil:
		// We found a peer from a randomly generated ID. This should be very
		// unlikely.
		logger.Warningf("random walk toward %s actually found peer: %s", id, p)
		return nil
	default:
		return err
	}
}

// Traverse the DHT toward the self ID
func (dht *IpfsDHT) selfWalk(ctx context.Context) error {
	_, err := dht.walk(ctx, dht.self)
	if err == routing.ErrNotFound {
		return nil
	}
	return err
}

//scan the RT,& do a random walk on k-buckets that haven't been queried since the given bucket period
func (dht *IpfsDHT) runBootstrap(ctx context.Context, cfg BootstrapConfig) error {
	doQuery := func(n int, target string, f func(context.Context) error) error {
		logger.Infof("starting bootstrap query for bucket %d to %s (routing table size was %d)",
			n, target, dht.routingTable.Size())
		defer func() {
			logger.Infof("finished bootstrap query for bucket %d to %s (routing table size is now %d)",
				n, target, dht.routingTable.Size())
		}()
		queryCtx, cancel := context.WithTimeout(ctx, cfg.Timeout)
		defer cancel()
		err := f(queryCtx)
		if err == context.DeadlineExceeded && queryCtx.Err() == context.DeadlineExceeded && ctx.Err() == nil {
			return nil
		}
		return err
	}

	buckets := dht.routingTable.GetAllBuckets()
	var wg sync.WaitGroup
	errChan := make(chan error)

	for bucketID, bucket := range buckets {
		if time.Since(bucket.LastQueriedAt()) > cfg.BucketPeriod {
			wg.Add(1)
			go func(bucketID int, errChan chan<- error) {
				defer wg.Done()
				// gen rand peer in the bucket
				randPeerInBucket, err := dht.routingTable.GenRandPeerID(bucketID)
				if err != nil {
					errChan <- errors.Wrapf(err, "failed to generate random peer ID in bucket %d", bucketID)
					return
				}

				// walk to the generated peer
				walkFnc := func(c context.Context) error {
					_, err := dht.walk(ctx, randPeerInBucket)
					if err == routing.ErrNotFound {
						return nil
					}
					return err
				}

				if err := doQuery(bucketID, randPeerInBucket.String(), walkFnc); err != nil {
					errChan <- errors.Wrapf(err, "failed to do a random walk on bucket %d", bucketID)
				}
			}(bucketID, errChan)
		}
	}

	// wait for all walks to finish & close the error channel
	go func() {
		wg.Wait()
		close(errChan)
	}()

	// accumulate errors from all go-routines
	var errStrings []string
	for err := range errChan {
		errStrings = append(errStrings, err.Error())
	}
	if len(errStrings) == 0 {
		return nil
	} else {
		return fmt.Errorf("errors encountered while running bootstrap on RT: %s", strings.Join(errStrings, "\n"))
	}
}

// This is a synchronous bootstrap.
func (dht *IpfsDHT) BootstrapOnce(ctx context.Context, cfg BootstrapConfig) error {
	if err := dht.BootstrapSelf(ctx); err != nil {
		return errors.Wrap(err, "failed bootstrap while searching for self")
	} else {
		return dht.runBootstrap(ctx, cfg)
	}
}

func (dht *IpfsDHT) BootstrapRandom(ctx context.Context) error {
	return dht.randomWalk(ctx)
}

func (dht *IpfsDHT) BootstrapSelf(ctx context.Context) error {
	return dht.selfWalk(ctx)
}
<commit_msg>dep: update go-libp2p-kbucket<commit_after>package dht

import (
	"context"
	"crypto/rand"
	"fmt"
	"strings"
	"sync"
	"time"

	u "github.com/ipfs/go-ipfs-util"
	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/routing"
	"github.com/multiformats/go-multiaddr"
	_ "github.com/multiformats/go-multiaddr-dns"
	"github.com/pkg/errors"
)

var DefaultBootstrapPeers []multiaddr.Multiaddr

func init() {
	for _, s := range []string{
		"/dnsaddr/bootstrap.libp2p.io/ipfs/QmNnooDu7bfjPFoTZYxMNLWUQJyrVwtbZg5gBMjTezGAJN",
		"/dnsaddr/bootstrap.libp2p.io/ipfs/QmQCU2EcMqAqQPR2i9bChDtGNJchTbq5TbXJJ16u19uLTa",
		"/dnsaddr/bootstrap.libp2p.io/ipfs/QmbLHAnMoJPWSCR5Zhtx6BHJX9KiKNN6tpvbUcqanj75Nb",
		"/dnsaddr/bootstrap.libp2p.io/ipfs/QmcZf59bWwK5XFi76CZX8cbJ4BhTzzA3gU1ZjYZcYW3dwt",
		"/ip4/104.131.131.82/tcp/4001/ipfs/QmaCpDMGvV2BGHeYERUEnRQAwe3N8SzbUtfsmvsqQLuvuJ",            // mars.i.ipfs.io
		"/ip4/104.236.179.241/tcp/4001/ipfs/QmSoLPppuBtQSGwKDZT2M73ULpjvfd3aZ6ha4oFGL1KrGM",           // pluto.i.ipfs.io
		"/ip4/128.199.219.111/tcp/4001/ipfs/QmSoLSafTMBsPKadTEgaXctDQVcqN88CNLHXMkTNwMKPnu",           // saturn.i.ipfs.io
		"/ip4/104.236.76.40/tcp/4001/ipfs/QmSoLV4Bbm51jM9C4gDYZQ9Cy3U6aXMJDAbzgu2fzaDs64",             // venus.i.ipfs.io
		"/ip4/178.62.158.247/tcp/4001/ipfs/QmSoLer265NRgSp2LA3dPaeykiS1J6DifTC88f5uVQKNAd",            // earth.i.ipfs.io
		"/ip6/2604:a880:1:20::203:d001/tcp/4001/ipfs/QmSoLPppuBtQSGwKDZT2M73ULpjvfd3aZ6ha4oFGL1KrGM",  // pluto.i.ipfs.io
		"/ip6/2400:6180:0:d0::151:6001/tcp/4001/ipfs/QmSoLSafTMBsPKadTEgaXctDQVcqN88CNLHXMkTNwMKPnu",  // saturn.i.ipfs.io
		"/ip6/2604:a880:800:10::4a:5001/tcp/4001/ipfs/QmSoLV4Bbm51jM9C4gDYZQ9Cy3U6aXMJDAbzgu2fzaDs64", // venus.i.ipfs.io
		"/ip6/2a03:b0c0:0:1010::23:1001/tcp/4001/ipfs/QmSoLer265NRgSp2LA3dPaeykiS1J6DifTC88f5uVQKNAd", // earth.i.ipfs.io
	} {
		ma, err := multiaddr.NewMultiaddr(s)
		if err != nil {
			panic(err)
		}
		DefaultBootstrapPeers = append(DefaultBootstrapPeers, ma)
	}
}

// BootstrapConfig specifies parameters used for bootstrapping the DHT.
type BootstrapConfig struct {
	BucketPeriod             time.Duration // how long to wait for a k-bucket to be queried before doing a random walk on it
	Timeout                  time.Duration // how long to wait for a bootstrap query to run
	RoutingTableScanInterval time.Duration // how often to scan the RT for k-buckets that haven't been queried since the given period
	SelfQueryInterval        time.Duration // how often to query for self
}

var DefaultBootstrapConfig = BootstrapConfig{
	// same as that mentioned in the kad dht paper
	BucketPeriod: 1 * time.Hour,

	// since the default bucket period is 1 hour, a scan interval of 30 minutes sounds reasonable
	RoutingTableScanInterval: 30 * time.Minute,

	Timeout: 10 * time.Second,

	SelfQueryInterval: 1 * time.Hour,
}

// A method in the IpfsRouting interface. It calls BootstrapWithConfig with
// the default bootstrap config.
func (dht *IpfsDHT) Bootstrap(ctx context.Context) error {
	return dht.BootstrapWithConfig(ctx, DefaultBootstrapConfig)
}

// Runs cfg.Queries bootstrap queries every cfg.BucketPeriod.
func (dht *IpfsDHT) BootstrapWithConfig(ctx context.Context, cfg BootstrapConfig) error {
	// we should query for self periodically so we can discover closer peers
	go func() {
		for {
			err := dht.BootstrapSelf(ctx)
			if err != nil {
				logger.Warningf("error bootstrapping while searching for my self (I'm Too Shallow ?): %s", err)
			}
			select {
			case <-time.After(cfg.SelfQueryInterval):
			case <-ctx.Done():
				return
			}
		}
	}()

	// scan the RT table periodically & do a random walk on k-buckets that haven't been queried since the given bucket period
	go func() {
		for {
			err := dht.runBootstrap(ctx, cfg)
			if err != nil {
				logger.Warningf("error bootstrapping: %s", err)
			}
			select {
			case <-time.After(cfg.RoutingTableScanInterval):
			case <-ctx.Done():
				return
			}
		}
	}()
	return nil
}

func newRandomPeerId() peer.ID {
	id := make([]byte, 32) // SHA256 is the default. TODO: Use a more canonical way to generate random IDs.
	rand.Read(id)
	id = u.Hash(id) // TODO: Feed this directly into the multihash instead of hashing it.
	return peer.ID(id)
}

// Traverse the DHT toward the given ID.
func (dht *IpfsDHT) walk(ctx context.Context, target peer.ID) (peer.AddrInfo, error) {
	// TODO: Extract the query action (traversal logic?) inside FindPeer,
	// don't actually call through the FindPeer machinery, which can return
	// things out of the peer store etc.
	return dht.FindPeer(ctx, target)
}

// Traverse the DHT toward a random ID.
func (dht *IpfsDHT) randomWalk(ctx context.Context) error {
	id := newRandomPeerId()
	p, err := dht.walk(ctx, id)
	switch err {
	case routing.ErrNotFound:
		return nil
	case nil:
		// We found a peer from a randomly generated ID. This should be very
		// unlikely.
		logger.Warningf("random walk toward %s actually found peer: %s", id, p)
		return nil
	default:
		return err
	}
}

// Traverse the DHT toward the self ID
func (dht *IpfsDHT) selfWalk(ctx context.Context) error {
	_, err := dht.walk(ctx, dht.self)
	if err == routing.ErrNotFound {
		return nil
	}
	return err
}

//scan the RT,& do a random walk on k-buckets that haven't been queried since the given bucket period
func (dht *IpfsDHT) runBootstrap(ctx context.Context, cfg BootstrapConfig) error {
	doQuery := func(n int, target string, f func(context.Context) error) error {
		logger.Infof("starting bootstrap query for bucket %d to %s (routing table size was %d)",
			n, target, dht.routingTable.Size())
		defer func() {
			logger.Infof("finished bootstrap query for bucket %d to %s (routing table size is now %d)",
				n, target, dht.routingTable.Size())
		}()
		queryCtx, cancel := context.WithTimeout(ctx, cfg.Timeout)
		defer cancel()
		err := f(queryCtx)
		if err == context.DeadlineExceeded && queryCtx.Err() == context.DeadlineExceeded && ctx.Err() == nil {
			return nil
		}
		return err
	}

	buckets := dht.routingTable.GetAllBuckets()
	var wg sync.WaitGroup
	errChan := make(chan error)

	for bucketID, bucket := range buckets {
		if time.Since(bucket.RefreshedAt()) > cfg.BucketPeriod {
			wg.Add(1)
			go func(bucketID int, errChan chan<- error) {
				defer wg.Done()
				// gen rand peer in the bucket
				randPeerInBucket := dht.routingTable.GenRandPeerID(bucketID)

				// walk to the generated peer
				walkFnc := func(c context.Context) error {
					_, err := dht.walk(ctx, randPeerInBucket)
					if err == routing.ErrNotFound {
						return nil
					}
					return err
				}

				if err := doQuery(bucketID, randPeerInBucket.String(), walkFnc); err != nil {
					errChan <- errors.Wrapf(err, "failed to do a random walk on bucket %d", bucketID)
				}
			}(bucketID, errChan)
		}
	}

	// wait for all walks to finish & close the error channel
	go func() {
		wg.Wait()
		close(errChan)
	}()

	// accumulate errors from all go-routines
	var errStrings []string
	for err := range errChan {
		errStrings = append(errStrings, err.Error())
	}
	if len(errStrings) == 0 {
		return nil
	} else {
		return fmt.Errorf("errors encountered while running bootstrap on RT: %s", strings.Join(errStrings, "\n"))
	}
}

// This is a synchronous bootstrap.
func (dht *IpfsDHT) BootstrapOnce(ctx context.Context, cfg BootstrapConfig) error {
	if err := dht.BootstrapSelf(ctx); err != nil {
		return errors.Wrap(err, "failed bootstrap while searching for self")
	} else {
		return dht.runBootstrap(ctx, cfg)
	}
}

func (dht *IpfsDHT) BootstrapRandom(ctx context.Context) error {
	return dht.randomWalk(ctx)
}

func (dht *IpfsDHT) BootstrapSelf(ctx context.Context) error {
	return dht.selfWalk(ctx)
}
<|endoftext|>


<commit_before>// Copyright 2017 The go-interpreter Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package disasm provides functions for disassembling WebAssembly bytecode.
package disasm

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
	"math"

	"github.com/go-interpreter/wagon/internal/stack"
	"github.com/go-interpreter/wagon/wasm"
	"github.com/go-interpreter/wagon/wasm/leb128"
	ops "github.com/go-interpreter/wagon/wasm/operators"
)

// Instr describes an instruction, consisting of an operator, with its
// appropriate immediate value(s).
type Instr struct {
	Op ops.Op

	// Immediates are arguments to an operator in the bytecode stream itself.
	// Valid value types are:
	// - (u)(int/float)(32/64)
	// - wasm.BlockType
	Immediates []interface{}
	NewStack   *StackInfo // non-nil if the instruction requires the current stack to be unwound.
	Block      *BlockInfo // non-nil if the instruction starts a new block.
}

// StackInfo stores details about a new stack ended by an instruction.
type StackInfo struct {
	StackTopDiff int64 // The difference between the stack depths at the end of the block
	PreserveTop  bool  // Whether the value on the top of the stack should be preserved while unwinding
}

// BlockInfo stores details about a block created or ended by an instruction.
type BlockInfo struct {
	Start     bool           // If true, this instruction starts a block. Else this instruction ends it.
	Signature wasm.BlockType // The block signature
	// The index to the accompanying control operator.
	// For 'if', this is an index to the 'else' operator
	// For else/loop/block, the index is to the 'end' operator
	PairIndex int
}

// Disassembly is the result of disassembling a WebAssembly function.
type Disassembly struct {
	Code     []Instr
	MaxDepth int // The maximum stack depth that can be reached while executing this function
}

func (d *Disassembly) checkMaxDepth(depth int) {
	if depth > d.MaxDepth {
		d.MaxDepth = depth
	}
}

var ErrStackUnderflow = errors.New("disasm: stack underflow")

// Disassemble disassembles the given function. It also takes the function's
// parent module as an argument for locating any other functions referenced by
// fn.
func Disassemble(fn wasm.Function, module *wasm.Module) (*Disassembly, error) {
	code := fn.Body.Code
	reader := bytes.NewReader(code)
	disas := &Disassembly{}

	// a stack of current execution stack depth values, so that the depth for each
	// stack is maintained indepepdently for calculating discard values
	stackDepths := &stack.Stack{}
	stackDepths.Push(0)
	blockIndices := &stack.Stack{} // a stack of indices to operators which start new blocks
	curIndex := 0
	var lastOpReturn bool

	for {
		op, err := reader.ReadByte()
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}
		logger.Printf("stack top is %d", stackDepths.Top())

		opStr, err := ops.New(op)
		if err != nil {
			return nil, err
		}
		instr := Instr{
			Op:         opStr,
			Immediates: [](interface{}){},
		}

		logger.Printf("Name is %s", opStr.Name)
		if !opStr.Polymorphic {
			top := int(stackDepths.Top())
			top -= len(opStr.Args)
			stackDepths.SetTop(uint64(top))
			if top < -1 {
				return nil, ErrStackUnderflow
			}
			if opStr.Returns != wasm.ValueType(wasm.BlockTypeEmpty) {
				top++
				stackDepths.SetTop(uint64(top))
			}
			disas.checkMaxDepth(top)
		}

		switch op {
		case ops.Drop:
			stackDepths.SetTop(stackDepths.Top() - 1)
		case ops.Select:
			stackDepths.SetTop(stackDepths.Top() - 3)
		case ops.Return:
			stackDepths.SetTop(stackDepths.Top() - uint64(len(fn.Sig.ReturnTypes)))
			lastOpReturn = true
		case ops.End, ops.Else:
			// The max depth reached while execing the current block
			curDepth := stackDepths.Top()
			blockStartIndex := blockIndices.Pop()
			blockSig := disas.Code[blockStartIndex].Block.Signature
			instr.Block = &BlockInfo{
				Start:     false,
				Signature: blockSig,
				PairIndex: int(blockStartIndex),
			}

			// The max depth reached while execing the last block
			// If the signature of the current block is not empty,
			// this will be incremented
			prevDepthIndex := stackDepths.Len() - 2
			prevDepth := stackDepths.Get(prevDepthIndex)

			if blockSig != wasm.BlockTypeEmpty {
				stackDepths.Set(prevDepthIndex, prevDepth+1)
				disas.checkMaxDepth(int(stackDepths.Get(prevDepthIndex)))
			}

			if !lastOpReturn {
				elemsDiscard := int(curDepth) - int(prevDepth)
				if elemsDiscard < -1 {
					return nil, ErrStackUnderflow
				}
				instr.NewStack = &StackInfo{
					StackTopDiff: int64(elemsDiscard),
					PreserveTop:  blockSig != wasm.BlockTypeEmpty,
				}
			} else {
				instr.NewStack = &StackInfo{}
			}

			stackDepths.Pop()
			if op == ops.Else {
				stackDepths.Push(0)
				blockIndices.Push(uint64(curIndex))
			}
		case ops.Block, ops.Loop, ops.If:
			sig, err := leb128.ReadVarint32(reader)
			if err != nil {
				return nil, err
			}
			logger.Printf("if, depth is %d", stackDepths.Top())
			stackDepths.Push(stackDepths.Top())
			instr.Block = &BlockInfo{
				Start:     true,
				Signature: wasm.BlockType(sig),
			}

			blockIndices.Push(uint64(curIndex))
			instr.Immediates = append(instr.Immediates, wasm.BlockType(sig))
		case ops.Br, ops.BrIf:
			depth, err := leb128.ReadVarUint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, depth)

			curDepth := stackDepths.Top()
			elemsDiscard := int64(int(curDepth) - int(stackDepths.Get(int(depth))))
			if elemsDiscard < 0 {
				return nil, ErrStackUnderflow
			}

			// we use getBlockDepth because a getBlockIndex would
			// have the same code
			index := blockIndices.Get(int(depth))
			instr.NewStack = &StackInfo{
				StackTopDiff: elemsDiscard,
				PreserveTop:  disas.Code[index].Block.Signature != wasm.BlockTypeEmpty,
			}

		case ops.BrTable:
			targetCount, err := leb128.ReadVarUint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, targetCount)
			for i := uint32(0); i < targetCount; i++ {
				entry, err := leb128.ReadVarUint32(reader)
				if err != nil {
					return nil, err
				}
				instr.Immediates = append(instr.Immediates, entry)
			}

			defaultTarget, err := leb128.ReadVarUint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, defaultTarget)
			stackDepths.SetTop(stackDepths.Top() - 1)
		case ops.Call, ops.CallIndirect:
			index, err := leb128.ReadVarUint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, index)
			if op == ops.CallIndirect {
				reserved, err := leb128.ReadVarUint32(reader)
				if err != nil {
					return nil, err
				}
				instr.Immediates = append(instr.Immediates, reserved)
			}
			fn := module.GetFunction(int(index))
			top := int(stackDepths.Top())
			top -= len(fn.Sig.ParamTypes)
			top += len(fn.Sig.ReturnTypes)
			stackDepths.SetTop(uint64(top))
			disas.checkMaxDepth(top)
		case ops.GetLocal, ops.SetLocal, ops.TeeLocal, ops.GetGlobal, ops.SetGlobal:
			index, err := leb128.ReadVarUint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, index)

			top := stackDepths.Top()
			switch op {
			case ops.GetLocal, ops.GetGlobal:
				top++
				stackDepths.SetTop(top)
				disas.checkMaxDepth(int(top))
			case ops.SetLocal, ops.SetGlobal:
				top--
				stackDepths.SetTop(top)
			case ops.TeeLocal:
				// stack remains unchanged for tee_local
			}
		case ops.I32Const:
			i, err := leb128.ReadVarint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, i)
		case ops.I64Const:
			i, err := leb128.ReadVarint64(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, i)
		case ops.F32Const:
			var i uint32
			// TODO(vibhavp): Switch to a reflect-free method in the future
			// for reading off the bytestream.
			err := binary.Read(reader, binary.LittleEndian, &i)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, math.Float32frombits(i))
		case ops.F64Const:
			var i uint64
			// TODO(vibhavp): Switch to a reflect-free method in the future
			// for reading off the bytestream.
			err := binary.Read(reader, binary.LittleEndian, &i)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, math.Float64frombits(i))
		case ops.I32Load, ops.I64Load, ops.F32Load, ops.F64Load, ops.I32Load8s, ops.I32Load8u, ops.I32Load16s, ops.I32Load16u, ops.I64Load8s, ops.I64Load8u, ops.I64Load16s, ops.I64Load16u, ops.I64Load32s, ops.I64Load32u, ops.I32Store, ops.I64Store, ops.F32Store, ops.F64Store, ops.I32Store8, ops.I32Store16, ops.I64Store8, ops.I64Store16, ops.I64Store32:
			// read memory_immediate
			flags, err := leb128.ReadVarUint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, flags)

			offset, err := leb128.ReadVarUint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, offset)
		case ops.CurrentMemory, ops.GrowMemory:
			res, err := leb128.ReadVarUint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, res)

			curDepth := stackDepths.Top()
			switch op {
			case ops.CurrentMemory:
				top := curDepth + 1
				stackDepths.SetTop(top)
				disas.checkMaxDepth(int(top))
			case ops.GrowMemory:
				top := curDepth - 1
				stackDepths.SetTop(top)
			}
		}

		if op != ops.Return {
			lastOpReturn = false
		}

		disas.Code = append(disas.Code, instr)
		curIndex++
	}

	return disas, nil
}
<commit_msg>disasm: Handle branches to relative depths correctly.<commit_after>// Copyright 2017 The go-interpreter Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package disasm provides functions for disassembling WebAssembly bytecode.
package disasm

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
	"math"

	"github.com/go-interpreter/wagon/internal/stack"
	"github.com/go-interpreter/wagon/wasm"
	"github.com/go-interpreter/wagon/wasm/leb128"
	ops "github.com/go-interpreter/wagon/wasm/operators"
)

// Instr describes an instruction, consisting of an operator, with its
// appropriate immediate value(s).
type Instr struct {
	Op ops.Op

	// Immediates are arguments to an operator in the bytecode stream itself.
	// Valid value types are:
	// - (u)(int/float)(32/64)
	// - wasm.BlockType
	Immediates []interface{}
	NewStack   *StackInfo // non-nil if the instruction requires the current stack to be unwound.
	Block      *BlockInfo // non-nil if the instruction starts a new block.
}

// StackInfo stores details about a new stack ended by an instruction.
type StackInfo struct {
	StackTopDiff int64 // The difference between the stack depths at the end of the block
	PreserveTop  bool  // Whether the value on the top of the stack should be preserved while unwinding
}

// BlockInfo stores details about a block created or ended by an instruction.
type BlockInfo struct {
	Start     bool           // If true, this instruction starts a block. Else this instruction ends it.
	Signature wasm.BlockType // The block signature
	// The index to the accompanying control operator.
	// For 'if', this is an index to the 'else' operator
	// For else/loop/block, the index is to the 'end' operator
	PairIndex int
}

// Disassembly is the result of disassembling a WebAssembly function.
type Disassembly struct {
	Code     []Instr
	MaxDepth int // The maximum stack depth that can be reached while executing this function
}

func (d *Disassembly) checkMaxDepth(depth int) {
	if depth > d.MaxDepth {
		d.MaxDepth = depth
	}
}

var ErrStackUnderflow = errors.New("disasm: stack underflow")

// Disassemble disassembles the given function. It also takes the function's
// parent module as an argument for locating any other functions referenced by
// fn.
func Disassemble(fn wasm.Function, module *wasm.Module) (*Disassembly, error) {
	code := fn.Body.Code
	reader := bytes.NewReader(code)
	disas := &Disassembly{}

	// a stack of current execution stack depth values, so that the depth for each
	// stack is maintained indepepdently for calculating discard values
	stackDepths := &stack.Stack{}
	stackDepths.Push(0)
	blockIndices := &stack.Stack{} // a stack of indices to operators which start new blocks
	curIndex := 0
	var lastOpReturn bool

	for {
		op, err := reader.ReadByte()
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}
		logger.Printf("stack top is %d", stackDepths.Top())

		opStr, err := ops.New(op)
		if err != nil {
			return nil, err
		}
		instr := Instr{
			Op:         opStr,
			Immediates: [](interface{}){},
		}

		logger.Printf("Name is %s", opStr.Name)
		if !opStr.Polymorphic {
			top := int(stackDepths.Top())
			top -= len(opStr.Args)
			stackDepths.SetTop(uint64(top))
			if top < -1 {
				return nil, ErrStackUnderflow
			}
			if opStr.Returns != wasm.ValueType(wasm.BlockTypeEmpty) {
				top++
				stackDepths.SetTop(uint64(top))
			}
			disas.checkMaxDepth(top)
		}

		switch op {
		case ops.Drop:
			stackDepths.SetTop(stackDepths.Top() - 1)
		case ops.Select:
			stackDepths.SetTop(stackDepths.Top() - 3)
		case ops.Return:
			stackDepths.SetTop(stackDepths.Top() - uint64(len(fn.Sig.ReturnTypes)))
			lastOpReturn = true
		case ops.End, ops.Else:
			// The max depth reached while execing the current block
			curDepth := stackDepths.Top()
			blockStartIndex := blockIndices.Pop()
			blockSig := disas.Code[blockStartIndex].Block.Signature
			instr.Block = &BlockInfo{
				Start:     false,
				Signature: blockSig,
				PairIndex: int(blockStartIndex),
			}

			// The max depth reached while execing the last block
			// If the signature of the current block is not empty,
			// this will be incremented
			prevDepthIndex := stackDepths.Len() - 2
			prevDepth := stackDepths.Get(prevDepthIndex)

			if blockSig != wasm.BlockTypeEmpty {
				stackDepths.Set(prevDepthIndex, prevDepth+1)
				disas.checkMaxDepth(int(stackDepths.Get(prevDepthIndex)))
			}

			if !lastOpReturn {
				elemsDiscard := int(curDepth) - int(prevDepth)
				if elemsDiscard < -1 {
					return nil, ErrStackUnderflow
				}
				instr.NewStack = &StackInfo{
					StackTopDiff: int64(elemsDiscard),
					PreserveTop:  blockSig != wasm.BlockTypeEmpty,
				}
			} else {
				instr.NewStack = &StackInfo{}
			}

			stackDepths.Pop()
			if op == ops.Else {
				stackDepths.Push(0)
				blockIndices.Push(uint64(curIndex))
			}
		case ops.Block, ops.Loop, ops.If:
			sig, err := leb128.ReadVarint32(reader)
			if err != nil {
				return nil, err
			}
			logger.Printf("if, depth is %d", stackDepths.Top())
			stackDepths.Push(stackDepths.Top())
			instr.Block = &BlockInfo{
				Start:     true,
				Signature: wasm.BlockType(sig),
			}

			blockIndices.Push(uint64(curIndex))
			instr.Immediates = append(instr.Immediates, wasm.BlockType(sig))
		case ops.Br, ops.BrIf:
			depth, err := leb128.ReadVarUint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, depth)

			curDepth := stackDepths.Top()
			elemsDiscard := int64(int(curDepth) - int(stackDepths.Get(stackDepths.Len()-1-int(depth))))
			if elemsDiscard < 0 {
				return nil, ErrStackUnderflow
			}

			index := blockIndices.Get(blockIndices.Len() - 1 - int(depth))
			instr.NewStack = &StackInfo{
				StackTopDiff: elemsDiscard,
				PreserveTop:  disas.Code[index].Block.Signature != wasm.BlockTypeEmpty,
			}

		case ops.BrTable:
			targetCount, err := leb128.ReadVarUint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, targetCount)
			for i := uint32(0); i < targetCount; i++ {
				entry, err := leb128.ReadVarUint32(reader)
				if err != nil {
					return nil, err
				}
				instr.Immediates = append(instr.Immediates, entry)
			}

			defaultTarget, err := leb128.ReadVarUint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, defaultTarget)
			stackDepths.SetTop(stackDepths.Top() - 1)
		case ops.Call, ops.CallIndirect:
			index, err := leb128.ReadVarUint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, index)
			if op == ops.CallIndirect {
				reserved, err := leb128.ReadVarUint32(reader)
				if err != nil {
					return nil, err
				}
				instr.Immediates = append(instr.Immediates, reserved)
			}
			fn := module.GetFunction(int(index))
			top := int(stackDepths.Top())
			top -= len(fn.Sig.ParamTypes)
			top += len(fn.Sig.ReturnTypes)
			stackDepths.SetTop(uint64(top))
			disas.checkMaxDepth(top)
		case ops.GetLocal, ops.SetLocal, ops.TeeLocal, ops.GetGlobal, ops.SetGlobal:
			index, err := leb128.ReadVarUint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, index)

			top := stackDepths.Top()
			switch op {
			case ops.GetLocal, ops.GetGlobal:
				top++
				stackDepths.SetTop(top)
				disas.checkMaxDepth(int(top))
			case ops.SetLocal, ops.SetGlobal:
				top--
				stackDepths.SetTop(top)
			case ops.TeeLocal:
				// stack remains unchanged for tee_local
			}
		case ops.I32Const:
			i, err := leb128.ReadVarint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, i)
		case ops.I64Const:
			i, err := leb128.ReadVarint64(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, i)
		case ops.F32Const:
			var i uint32
			// TODO(vibhavp): Switch to a reflect-free method in the future
			// for reading off the bytestream.
			err := binary.Read(reader, binary.LittleEndian, &i)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, math.Float32frombits(i))
		case ops.F64Const:
			var i uint64
			// TODO(vibhavp): Switch to a reflect-free method in the future
			// for reading off the bytestream.
			err := binary.Read(reader, binary.LittleEndian, &i)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, math.Float64frombits(i))
		case ops.I32Load, ops.I64Load, ops.F32Load, ops.F64Load, ops.I32Load8s, ops.I32Load8u, ops.I32Load16s, ops.I32Load16u, ops.I64Load8s, ops.I64Load8u, ops.I64Load16s, ops.I64Load16u, ops.I64Load32s, ops.I64Load32u, ops.I32Store, ops.I64Store, ops.F32Store, ops.F64Store, ops.I32Store8, ops.I32Store16, ops.I64Store8, ops.I64Store16, ops.I64Store32:
			// read memory_immediate
			flags, err := leb128.ReadVarUint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, flags)

			offset, err := leb128.ReadVarUint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, offset)
		case ops.CurrentMemory, ops.GrowMemory:
			res, err := leb128.ReadVarUint32(reader)
			if err != nil {
				return nil, err
			}
			instr.Immediates = append(instr.Immediates, res)

			curDepth := stackDepths.Top()
			switch op {
			case ops.CurrentMemory:
				top := curDepth + 1
				stackDepths.SetTop(top)
				disas.checkMaxDepth(int(top))
			case ops.GrowMemory:
				top := curDepth - 1
				stackDepths.SetTop(top)
			}
		}

		if op != ops.Return {
			lastOpReturn = false
		}

		disas.Code = append(disas.Code, instr)
		curIndex++
	}

	return disas, nil
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	mailgun "github.com/riobard/go-mailgun"
	"log"
	"os"
	"strconv"
	"strings"
	"time"
)

var user_id = flag.Int("user_id", 0, "ID of user to output.")
var all_users = flag.Bool("all_users", false, "Whether to email all users.")
var send_email = flag.Bool("send_email", false, "Whether to send the email to the user.")
var mark_read = flag.Bool("mark_read", true, "Whether to mark activity sent as read. False means subsequent calls will send the same info.")

type User struct {
	Id       int
	username string
	Email    string
}

type GithubUser struct {
	Id    int
	Login string
	Url   string
}

// db thing.
type GithubRepo struct {
	Id       int
	User     string // should probably be Login
	RepoName string // should probably be Name
}

func (g *GithubRepo) FullName() string {
	return fmt.Sprintf("%s/%s", g.User, g.RepoName)
}

type GithubApiRepo struct {
	Id          int
	Owner       GithubUser
	Name        NString
	Description NString
}

func (g *GithubApiRepo) FullName() string {
	return fmt.Sprintf("%s/%s", g.Owner.Login, g.Name)
}

type Activity struct {
	Id            int
	github_id     int
	activity_type string // should be enum
	created_at    time.Time
	Username      string
	repo          GithubRepo // this isn't going to be memory efficient
	Meta          string     // full payload of json object
}

type Treeish struct {
	Label    string
	Sha      string
	Repo     GithubApiRepo
	Html_url string
	User     GithubUser
}

type Commit struct {
	Sha     string
	Message string
	Author  struct {
		Name  string
		Email string
	}
}

func (c *Commit) ShortSha() string {
	return c.Sha[0:6]
}

func (c *Commit) ShortCommit() string {
	msg := strings.Split(c.Message, "\n")[0]
	if len(msg) > 80 {
		return msg[0:77] + "..."
	}
	return msg
}

// Nullable String, per https://groups.google.com/d/msg/golang-nuts/JOFWAqrTbUs/KZIGEPclSpwJ
type NString string

func (n *NString) UnmarshalJSON(b []byte) (err error) {
	if string(b) == "null" {
		return nil
	}
	return json.Unmarshal(b, (*string)(n))
}

func getUserRepos(db *sql.DB, user_id int) ([]GithubRepo, error) {
	// guess as to initial size. Likely very few users following < 10 repos.
	var repos = make([]GithubRepo, 10)

	// TODO(justinabrahms): Should really alter the schema such
	// that the join from user <-> repo doesn't go through
	// userprofiles.
	rows, err := db.Query(
		`SELECT r.id, username, project_name
	           FROM streamer_repo r
		   JOIN streamer_userprofile_repos upr ON upr.repo_id = r.id
		   JOIN streamer_userprofile up ON up.id = upr.userprofile_id
                   WHERE up.user_id = ?;`, user_id)
	// what's a good way to handle this not working?
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var username, project_name string
	var pk int
	for rows.Next() {
		err = rows.Scan(&pk, &username, &project_name)
		if err != nil {
			fmt.Println("Error!: ", err)
		}
		var repo = GithubRepo{pk, username, project_name}
		repos = append(repos, repo)
	}

	return repos, err
}

func getUser(db *sql.DB, uid int) (u User, err error) {
	row := db.QueryRow(
		`SELECT id, username, email
		   FROM auth_user
		   WHERE id = ?`, uid)
	err = row.Scan(&u.Id, &u.username, &u.Email)
	return
}

func getUsers(db *sql.DB) (users []User, err error) {
	rows, err := db.Query(
		`SELECT id, username, email
		   FROM auth_user`)
	if err != nil {
		return
	}
	defer rows.Close()

	for rows.Next() {
		var u User
		err = rows.Scan(&u.Id, &u.username, &u.Email)
		if err != nil {
			return
		}
		users = append(users, u)
	}
	return
}

func getRepoActivity(db *sql.DB, repo *GithubRepo) (activity_list []Activity, err error) {
	activity_list = make([]Activity, 0)
	rows, err := db.Query(
		`SELECT a.id, a.event_id, a.type, a.created_at, ghu.name, r.username, r.project_name, meta
		  FROM streamer_activity a
		  JOIN streamer_repo r on r.id=a.repo_id
		  JOIN streamer_githubuser ghu on ghu.id=a.user_id
		  JOIN streamer_userprofile_repos upr on r.id=upr.repo_id
		  WHERE r.id = ?
		  AND a.created_at > DATE_SUB(NOW(), INTERVAL 5 day) -- don't send things more than a few days old. Think, new users who subscribe to rails/rails
		  AND (upr.last_sent is null -- hasn't been sent at all
		    OR a.created_at > upr.last_sent) --  or hasn't been sent since we've gotten new stuff
                  GROUP BY a.id`, // and unique-ify based on activity id to prevent dupes.
		repo.Id)
	if err != nil {
		return
	}
	defer rows.Close()

	var (
		pk, github_id, repo_id                                              int
		activity_type, username, repo_project, repo_user, meta, created_str string
		created_at                                                          time.Time
	)

	for rows.Next() {
		err = rows.Scan(&pk, &github_id, &activity_type, &created_str, &username, &repo_user, &repo_project, &meta)
		if err != nil {
			fmt.Println("ERR: ", err)
		}

		created_at, err = time.Parse("2006-01-02 15:04:05", created_str)
		if err != nil {
			fmt.Println("Can't parse the format of ", created_str)
			return nil, err
		}

		activity_list = append(activity_list,
			Activity{pk, github_id, activity_type, created_at, username,
				GithubRepo{repo_id, repo_user, repo_project}, meta})
	}

	return
}

func repoToTemplate(repo GithubRepo, activities []Activity, render_map map[string]func([]Activity, bool) string) (response string) {
	if len(activities) == 0 {
		return ""
	}
	var activity_map = make(map[string][]Activity)
	for _, activity := range activities {
		// This seems like a lot of juggling. Is there a better way?
		arr := activity_map[activity.activity_type]
		if arr == nil {
			arr = make([]Activity, 0)
			activity_map[activity.activity_type] = arr
		}
		arr = append(arr, activity)
		activity_map[activity.activity_type] = arr
	}

	// activity_map: activity_type => []activity
	for activity_type, activities := range activity_map {
		fn, ok := render_map[activity_type]
		if ok {
			response += fn(activities, true)
		} else {
			log.Print("Not sure how to render activites of type ", activity_type)
		}
	}
	if len(response) > 0 {
		response = fmt.Sprintf("\n\n%s:\n%s", repo.FullName(), response)
	}
	return

}

func repoToString(db *sql.DB, repo GithubRepo, response chan string) {
	activities, err := getRepoActivity(db, &repo)
	if err != nil {
		fmt.Println("ERR: ", err)
		os.Exit(1)
	}

	activity_type_to_renderer := map[string]func([]Activity, bool) string{
		"P":  pushRender,
		"PR": pullRequestRender,
		"D":  deleteRender,
		"C":  createRender,
		"IC": issueCommentRender,
		"Gl": wikiRender, // Gl is for Gollum, Github's wiki thing.
		"I":  issueRender,
		"Pb": publicRender,

		// Watches and forks for repos are WAYY too spammy.
		// "W":  watchRender,
		// "F":  forkRender,
	}

	response <- repoToTemplate(repo, activities, activity_type_to_renderer)
}

func markUserRepoSent(db *sql.DB, user User, repos []GithubRepo) (err error) {
	// finds the streamer_userprofile_repo row for the repo / user
	// combo, mark its last_sent as now
	ids := make([]string, 0)
	for _, repo := range repos {
		// TODO: why are they 0?
		if repo.Id != 0 {
			ids = append(ids, strconv.FormatInt(int64(repo.Id), 10))
		}
	}

	// There is likely a better way to get parameterization, but
	// it wasn't working for me with ?'s.
	str := fmt.Sprintf(
		`UPDATE streamer_userprofile_repos
		   SET last_sent=NOW()
		   WHERE repo_id IN (%s)`, strings.Join(ids, ","))
	_, err = db.Exec(str)
	return
}

// TODO: Github Users

// TODO: Need to finish the following activity types:
// - TA (team add)
// - FA (fork apply)
// - DO (Download)
// - RC (Pull Request Review Comment)
// - Fl (Follow)
// - G (Gist)
// - M (Member)
// - CC (Commit Comment)

func main() {
	// Expects the following environment variables:
	// DB_USER, DB_PASS, DB_DB  (database user, password and database)
	// MAILGUN_API_KEY  (api key from mailgun.org)

	flag.Parse()

	db, err := sql.Open("mysql", fmt.Sprintf("%s:%s@/%s?charset=utf8", os.Getenv("DB_USER"), os.Getenv("DB_PASS"), os.Getenv("DB_DB")))
	if err != nil {
		log.Fatalf("Unable to connect to mysql. ", err)
	}
	defer db.Close()

	var users []User
	if *user_id != 0 {
		users = make([]User, 1)
		user, err := getUser(db, *user_id)
		if err != nil {
			log.Fatalf("Couldn't return user %d. %s", *user_id, err)
		}
		users[0] = user
	} else if *all_users {
		users, err = getUsers(db)
		if err != nil {
			log.Fatal("Couldn't fetch all users.")
		}
	} else {
		log.Fatal("You must specify either to mail all users or a specific user id.")
	}

	for _, user := range users {
		repos, err := getUserRepos(db, user.Id)
		if err != nil {
			log.Print("Error fetching user's repos. %s", err)
			continue
		}

		response_chan := make(chan string, len(repos))
		for _, repo := range repos {
			repoToString(db, repo, response_chan)
		}

		var response string
		for num_responses := 0; num_responses < len(repos); {
			select {
			case r := <-response_chan:
				num_responses++
				response += r
			}
		}

		if *mark_read {
			fmt.Println("Marking read.")
			err = markUserRepoSent(db, user, repos)
			if err != nil {
				log.Print("Error updating repositories as sent.")
				continue
			}
		}

		if len(response) == 0 {
			log.Print("No content to email.")
			continue
		}

		if *send_email {
			mg := mailgun.Open(os.Getenv("MAILGUN_API_KEY"))
			e := &Email{
				from:    "justin@gitstreams.mailgun.org",
				to:      []string{user.Email},
				subject: "Gitstreams Digest Email",
				text:    response,
			}

			id, err := mg.Send(e)

			if err != nil {
				log.Print("Unable to send email. ", err)
			}
			log.Printf("MessageId = %s for uid:%s", id, user.Id)
		} else {
			fmt.Println("Would have sent the following email.")
			fmt.Println(response)
		}
	}

	return
}
<commit_msg>pass the error along with the log.Fatal.<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	mailgun "github.com/riobard/go-mailgun"
	"log"
	"os"
	"strconv"
	"strings"
	"time"
)

var user_id = flag.Int("user_id", 0, "ID of user to output.")
var all_users = flag.Bool("all_users", false, "Whether to email all users.")
var send_email = flag.Bool("send_email", false, "Whether to send the email to the user.")
var mark_read = flag.Bool("mark_read", true, "Whether to mark activity sent as read. False means subsequent calls will send the same info.")

type User struct {
	Id       int
	username string
	Email    string
}

type GithubUser struct {
	Id    int
	Login string
	Url   string
}

// db thing.
type GithubRepo struct {
	Id       int
	User     string // should probably be Login
	RepoName string // should probably be Name
}

func (g *GithubRepo) FullName() string {
	return fmt.Sprintf("%s/%s", g.User, g.RepoName)
}

type GithubApiRepo struct {
	Id          int
	Owner       GithubUser
	Name        NString
	Description NString
}

func (g *GithubApiRepo) FullName() string {
	return fmt.Sprintf("%s/%s", g.Owner.Login, g.Name)
}

type Activity struct {
	Id            int
	github_id     int
	activity_type string // should be enum
	created_at    time.Time
	Username      string
	repo          GithubRepo // this isn't going to be memory efficient
	Meta          string     // full payload of json object
}

type Treeish struct {
	Label    string
	Sha      string
	Repo     GithubApiRepo
	Html_url string
	User     GithubUser
}

type Commit struct {
	Sha     string
	Message string
	Author  struct {
		Name  string
		Email string
	}
}

func (c *Commit) ShortSha() string {
	return c.Sha[0:6]
}

func (c *Commit) ShortCommit() string {
	msg := strings.Split(c.Message, "\n")[0]
	if len(msg) > 80 {
		return msg[0:77] + "..."
	}
	return msg
}

// Nullable String, per https://groups.google.com/d/msg/golang-nuts/JOFWAqrTbUs/KZIGEPclSpwJ
type NString string

func (n *NString) UnmarshalJSON(b []byte) (err error) {
	if string(b) == "null" {
		return nil
	}
	return json.Unmarshal(b, (*string)(n))
}

func getUserRepos(db *sql.DB, user_id int) ([]GithubRepo, error) {
	// guess as to initial size. Likely very few users following < 10 repos.
	var repos = make([]GithubRepo, 10)

	// TODO(justinabrahms): Should really alter the schema such
	// that the join from user <-> repo doesn't go through
	// userprofiles.
	rows, err := db.Query(
		`SELECT r.id, username, project_name
	           FROM streamer_repo r
		   JOIN streamer_userprofile_repos upr ON upr.repo_id = r.id
		   JOIN streamer_userprofile up ON up.id = upr.userprofile_id
                   WHERE up.user_id = ?;`, user_id)
	// what's a good way to handle this not working?
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var username, project_name string
	var pk int
	for rows.Next() {
		err = rows.Scan(&pk, &username, &project_name)
		if err != nil {
			fmt.Println("Error!: ", err)
		}
		var repo = GithubRepo{pk, username, project_name}
		repos = append(repos, repo)
	}

	return repos, err
}

func getUser(db *sql.DB, uid int) (u User, err error) {
	row := db.QueryRow(
		`SELECT id, username, email
		   FROM auth_user
		   WHERE id = ?`, uid)
	err = row.Scan(&u.Id, &u.username, &u.Email)
	return
}

func getUsers(db *sql.DB) (users []User, err error) {
	rows, err := db.Query(
		`SELECT id, username, email
		   FROM auth_user`)
	if err != nil {
		return
	}
	defer rows.Close()

	for rows.Next() {
		var u User
		err = rows.Scan(&u.Id, &u.username, &u.Email)
		if err != nil {
			return
		}
		users = append(users, u)
	}
	return
}

func getRepoActivity(db *sql.DB, repo *GithubRepo) (activity_list []Activity, err error) {
	activity_list = make([]Activity, 0)
	rows, err := db.Query(
		`SELECT a.id, a.event_id, a.type, a.created_at, ghu.name, r.username, r.project_name, meta
		  FROM streamer_activity a
		  JOIN streamer_repo r on r.id=a.repo_id
		  JOIN streamer_githubuser ghu on ghu.id=a.user_id
		  JOIN streamer_userprofile_repos upr on r.id=upr.repo_id
		  WHERE r.id = ?
		  AND a.created_at > DATE_SUB(NOW(), INTERVAL 5 day) -- don't send things more than a few days old. Think, new users who subscribe to rails/rails
		  AND (upr.last_sent is null -- hasn't been sent at all
		    OR a.created_at > upr.last_sent) --  or hasn't been sent since we've gotten new stuff
                  GROUP BY a.id`, // and unique-ify based on activity id to prevent dupes.
		repo.Id)
	if err != nil {
		return
	}
	defer rows.Close()

	var (
		pk, github_id, repo_id                                              int
		activity_type, username, repo_project, repo_user, meta, created_str string
		created_at                                                          time.Time
	)

	for rows.Next() {
		err = rows.Scan(&pk, &github_id, &activity_type, &created_str, &username, &repo_user, &repo_project, &meta)
		if err != nil {
			fmt.Println("ERR: ", err)
		}

		created_at, err = time.Parse("2006-01-02 15:04:05", created_str)
		if err != nil {
			fmt.Println("Can't parse the format of ", created_str)
			return nil, err
		}

		activity_list = append(activity_list,
			Activity{pk, github_id, activity_type, created_at, username,
				GithubRepo{repo_id, repo_user, repo_project}, meta})
	}

	return
}

func repoToTemplate(repo GithubRepo, activities []Activity, render_map map[string]func([]Activity, bool) string) (response string) {
	if len(activities) == 0 {
		return ""
	}
	var activity_map = make(map[string][]Activity)
	for _, activity := range activities {
		// This seems like a lot of juggling. Is there a better way?
		arr := activity_map[activity.activity_type]
		if arr == nil {
			arr = make([]Activity, 0)
			activity_map[activity.activity_type] = arr
		}
		arr = append(arr, activity)
		activity_map[activity.activity_type] = arr
	}

	// activity_map: activity_type => []activity
	for activity_type, activities := range activity_map {
		fn, ok := render_map[activity_type]
		if ok {
			response += fn(activities, true)
		} else {
			log.Print("Not sure how to render activites of type ", activity_type)
		}
	}
	if len(response) > 0 {
		response = fmt.Sprintf("\n\n%s:\n%s", repo.FullName(), response)
	}
	return

}

func repoToString(db *sql.DB, repo GithubRepo, response chan string) {
	activities, err := getRepoActivity(db, &repo)
	if err != nil {
		fmt.Println("ERR: ", err)
		os.Exit(1)
	}

	activity_type_to_renderer := map[string]func([]Activity, bool) string{
		"P":  pushRender,
		"PR": pullRequestRender,
		"D":  deleteRender,
		"C":  createRender,
		"IC": issueCommentRender,
		"Gl": wikiRender, // Gl is for Gollum, Github's wiki thing.
		"I":  issueRender,
		"Pb": publicRender,

		// Watches and forks for repos are WAYY too spammy.
		// "W":  watchRender,
		// "F":  forkRender,
	}

	response <- repoToTemplate(repo, activities, activity_type_to_renderer)
}

func markUserRepoSent(db *sql.DB, user User, repos []GithubRepo) (err error) {
	// finds the streamer_userprofile_repo row for the repo / user
	// combo, mark its last_sent as now
	ids := make([]string, 0)
	for _, repo := range repos {
		// TODO: why are they 0?
		if repo.Id != 0 {
			ids = append(ids, strconv.FormatInt(int64(repo.Id), 10))
		}
	}

	// There is likely a better way to get parameterization, but
	// it wasn't working for me with ?'s.
	str := fmt.Sprintf(
		`UPDATE streamer_userprofile_repos
		   SET last_sent=NOW()
		   WHERE repo_id IN (%s)`, strings.Join(ids, ","))
	_, err = db.Exec(str)
	return
}

// TODO: Github Users

// TODO: Need to finish the following activity types:
// - TA (team add)
// - FA (fork apply)
// - DO (Download)
// - RC (Pull Request Review Comment)
// - Fl (Follow)
// - G (Gist)
// - M (Member)
// - CC (Commit Comment)

func main() {
	// Expects the following environment variables:
	// DB_USER, DB_PASS, DB_DB  (database user, password and database)
	// MAILGUN_API_KEY  (api key from mailgun.org)

	flag.Parse()

	db, err := sql.Open("mysql", fmt.Sprintf("%s:%s@/%s?charset=utf8", os.Getenv("DB_USER"), os.Getenv("DB_PASS"), os.Getenv("DB_DB")))
	if err != nil {
		log.Fatalf("Unable to connect to mysql. ", err)
	}
	defer db.Close()

	var users []User
	if *user_id != 0 {
		users = make([]User, 1)
		user, err := getUser(db, *user_id)
		if err != nil {
			log.Fatalf("Couldn't return user %d. %s", *user_id, err)
		}
		users[0] = user
	} else if *all_users {
		users, err = getUsers(db)
		if err != nil {
			log.Fatal("Couldn't fetch all users. %s", err)
		}
	} else {
		log.Fatal("You must specify either to mail all users or a specific user id.")
	}

	for _, user := range users {
		repos, err := getUserRepos(db, user.Id)
		if err != nil {
			log.Print("Error fetching user's repos. %s", err)
			continue
		}

		response_chan := make(chan string, len(repos))
		for _, repo := range repos {
			repoToString(db, repo, response_chan)
		}

		var response string
		for num_responses := 0; num_responses < len(repos); {
			select {
			case r := <-response_chan:
				num_responses++
				response += r
			}
		}

		if *mark_read {
			fmt.Println("Marking read.")
			err = markUserRepoSent(db, user, repos)
			if err != nil {
				log.Print("Error updating repositories as sent.")
				continue
			}
		}

		if len(response) == 0 {
			log.Print("No content to email.")
			continue
		}

		if *send_email {
			mg := mailgun.Open(os.Getenv("MAILGUN_API_KEY"))
			e := &Email{
				from:    "justin@gitstreams.mailgun.org",
				to:      []string{user.Email},
				subject: "Gitstreams Digest Email",
				text:    response,
			}

			id, err := mg.Send(e)

			if err != nil {
				log.Print("Unable to send email. ", err)
			}
			log.Printf("MessageId = %s for uid:%s", id, user.Id)
		} else {
			fmt.Println("Would have sent the following email.")
			fmt.Println(response)
		}
	}

	return
}
<|endoftext|>


<commit_before>package survey

import "testing"

func init() {
	// disable color output for all prompts to simplify testing
	DisableColor = true
}

func TestCanFormatMultiChoiceOptions(t *testing.T) {

	prompt := &MultiChoice{
		Options:  []string{"foo", "bar", "baz", "buz"},
		Defaults: []string{"bar", "buz"},
	}

	actual, err := runTemplate(
		MultiChoiceOptionsTemplate,
		MultiChoiceTemplateData{MultiChoice: *prompt, Selected: 2, Checked: map[int]bool{1: true, 3: true}},
	)

	if err != nil {
		t.Errorf("Failed to run template to format checkbox options: %s", err)
	}

	expected := ` ◯ foo
 ◉ bar
❯◯ baz
 ◉ buz
`

	if actual != expected {
		t.Errorf("Formatted checkbox options were not formatted correctly. Found:\n%s\nExpected:\n%s", actual, expected)
	}
}

func TestMultiChoiceFormatQuestion(t *testing.T) {

	prompt := &MultiChoice{
		Message:  "Pick your words:",
		Options:  []string{"foo", "bar", "baz", "buz"},
		Defaults: []string{"bar", "buz"},
	}

	actual, err := runTemplate(
		MultiChoiceQuestionTemplate,
		MultiChoiceTemplateData{MultiChoice: *prompt},
	)
	if err != nil {
		t.Errorf("Failed to run template to format checkbox question: %s", err)
	}

	expected := `? Pick your words: `

	if actual != expected {
		t.Errorf("Formatted checkbox question was not formatted correctly. Found:\n%s\nExpected:\n%s", actual, expected)
	}
}

func TestMultiChoiceFormatAnswer(t *testing.T) {

	prompt := &MultiChoice{
		Message:  "Pick your words:",
		Options:  []string{"foo", "bar", "baz", "buz"},
		Defaults: []string{"bar", "buz"},
	}

	actual, err := runTemplate(
		MultiChoiceQuestionTemplate,
		MultiChoiceTemplateData{MultiChoice: *prompt, Answer: []string{"foo", "buz"}},
	)
	if err != nil {
		t.Errorf("Failed to run template to format checkbox answer: %s", err)
	}

	expected := `? Pick your words: ["foo" "buz"]`

	if actual != expected {
		t.Errorf("Formatted checkbox answer was not formatted correctly. Found:\n%s\nExpected:\n%s", actual, expected)
	}
}
<commit_msg>fixed stylistic typo in tests<commit_after>package survey

import "testing"

func init() {
	// disable color output for all prompts to simplify testing
	DisableColor = true
}

func TestCanFormatMultiChoiceOptions(t *testing.T) {

	prompt := &MultiChoice{
		Options:  []string{"foo", "bar", "baz", "buz"},
		Defaults: []string{"bar", "buz"},
	}

	actual, err := runTemplate(
		MultiChoiceOptionsTemplate,
		MultiChoiceTemplateData{MultiChoice: *prompt, Selected: 2, Checked: map[int]bool{1: true, 3: true}},
	)

	if err != nil {
		t.Errorf("Failed to run template to format checkbox options: %s", err)
	}

	expected := `  ◯  foo
  ◉  bar
❯ ◯  baz
  ◉  buz
`

	if actual != expected {
		t.Errorf("Formatted checkbox options were not formatted correctly. Found:\n%s\nExpected:\n%s", actual, expected)
	}
}

func TestMultiChoiceFormatQuestion(t *testing.T) {

	prompt := &MultiChoice{
		Message:  "Pick your words:",
		Options:  []string{"foo", "bar", "baz", "buz"},
		Defaults: []string{"bar", "buz"},
	}

	actual, err := runTemplate(
		MultiChoiceQuestionTemplate,
		MultiChoiceTemplateData{MultiChoice: *prompt},
	)
	if err != nil {
		t.Errorf("Failed to run template to format checkbox question: %s", err)
	}

	expected := `? Pick your words: `

	if actual != expected {
		t.Errorf("Formatted checkbox question was not formatted correctly. Found:\n%s\nExpected:\n%s", actual, expected)
	}
}

func TestMultiChoiceFormatAnswer(t *testing.T) {

	prompt := &MultiChoice{
		Message:  "Pick your words:",
		Options:  []string{"foo", "bar", "baz", "buz"},
		Defaults: []string{"bar", "buz"},
	}

	actual, err := runTemplate(
		MultiChoiceQuestionTemplate,
		MultiChoiceTemplateData{MultiChoice: *prompt, Answer: []string{"foo", "buz"}},
	)
	if err != nil {
		t.Errorf("Failed to run template to format checkbox answer: %s", err)
	}

	expected := `? Pick your words: ["foo" "buz"]`

	if actual != expected {
		t.Errorf("Formatted checkbox answer was not formatted correctly. Found:\n%s\nExpected:\n%s", actual, expected)
	}
}
<|endoftext|>