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

package firestore

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	vkit "cloud.google.com/go/firestore/apiv1"
	"cloud.google.com/go/internal/trace"
	"cloud.google.com/go/internal/version"
	"github.com/golang/protobuf/ptypes"
	gax "github.com/googleapis/gax-go/v2"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
	"google.golang.org/api/transport"
	pb "google.golang.org/genproto/googleapis/firestore/v1"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

// resourcePrefixHeader is the name of the metadata header used to indicate
// the resource being operated on.
const resourcePrefixHeader = "google-cloud-resource-prefix"

// DetectProjectID is a sentinel value that instructs NewClient to detect the
// project ID. It is given in place of the projectID argument. NewClient will
// use the project ID from the given credentials or the default credentials
// (https://developers.google.com/accounts/docs/application-default-credentials)
// if no credentials were provided. When providing credentials, not all
// options will allow NewClient to extract the project ID. Specifically a JWT
// does not have the project ID encoded.
const DetectProjectID = "*detect-project-id*"

// A Client provides access to the Firestore service.
type Client struct {
	c          *vkit.Client
	projectID  string
	databaseID string // A client is tied to a single database.
}

// NewClient creates a new Firestore client that uses the given project.
func NewClient(ctx context.Context, projectID string, opts ...option.ClientOption) (*Client, error) {
	var o []option.ClientOption
	// Environment variables for gcloud emulator:
	// https://cloud.google.com/sdk/gcloud/reference/beta/emulators/firestore/
	if addr := os.Getenv("FIRESTORE_EMULATOR_HOST"); addr != "" {
		conn, err := grpc.Dial(addr, grpc.WithInsecure(), grpc.WithAuthority("owner"))
		if err != nil {
			return nil, fmt.Errorf("firestore: dialing address from env var FIRESTORE_EMULATOR_HOST: %v", err)
		}
		o = []option.ClientOption{option.WithGRPCConn(conn)}
	}
	o = append(o, opts...)

	if projectID == DetectProjectID {
		creds, err := transport.Creds(ctx, o...)
		if err != nil {
			return nil, fmt.Errorf("fetching creds: %v", err)
		}
		if creds.ProjectID == "" {
			return nil, errors.New("firestore: see the docs on DetectProjectID")
		}
		projectID = creds.ProjectID
	}

	vc, err := vkit.NewClient(ctx, o...)
	if err != nil {
		return nil, err
	}
	vc.SetGoogleClientInfo("gccl", version.Repo)
	c := &Client{
		c:          vc,
		projectID:  projectID,
		databaseID: "(default)", // always "(default)", for now
	}
	return c, nil

}

// Close closes any resources held by the client.
//
// Close need not be called at program exit.
func (c *Client) Close() error {
	return c.c.Close()
}

func (c *Client) path() string {
	return fmt.Sprintf("projects/%s/databases/%s", c.projectID, c.databaseID)
}

func withResourceHeader(ctx context.Context, resource string) context.Context {
	md, _ := metadata.FromOutgoingContext(ctx)
	md = md.Copy()
	md[resourcePrefixHeader] = []string{resource}
	return metadata.NewOutgoingContext(ctx, md)
}

// Collection creates a reference to a collection with the given path.
// A path is a sequence of IDs separated by slashes.
//
// Collection returns nil if path contains an even number of IDs or any ID is empty.
func (c *Client) Collection(path string) *CollectionRef {
	coll, _ := c.idsToRef(strings.Split(path, "/"), c.path())
	return coll
}

// Doc creates a reference to a document with the given path.
// A path is a sequence of IDs separated by slashes.
//
// Doc returns nil if path contains an odd number of IDs or any ID is empty.
func (c *Client) Doc(path string) *DocumentRef {
	_, doc := c.idsToRef(strings.Split(path, "/"), c.path())
	return doc
}

// CollectionGroup creates a reference to a group of collections that include
// the given ID, regardless of parent document.
//
// For example, consider:
// France/Cities/Paris = {population: 100}
// Canada/Cities/Montreal = {population: 90}
//
// CollectionGroup can be used to query across all "Cities" regardless of
// its parent "Countries". See ExampleCollectionGroup for a complete example.
func (c *Client) CollectionGroup(collectionID string) *CollectionGroupRef {
	return newCollectionGroupRef(c, c.path(), collectionID)
}

func (c *Client) idsToRef(IDs []string, dbPath string) (*CollectionRef, *DocumentRef) {
	if len(IDs) == 0 {
		return nil, nil
	}
	for _, id := range IDs {
		if id == "" {
			return nil, nil
		}
	}
	coll := newTopLevelCollRef(c, dbPath, IDs[0])
	i := 1
	for i < len(IDs) {
		doc := newDocRef(coll, IDs[i])
		i++
		if i == len(IDs) {
			return nil, doc
		}
		coll = newCollRefWithParent(c, doc, IDs[i])
		i++
	}
	return coll, nil
}

// GetAll retrieves multiple documents with a single call. The
// DocumentSnapshots are returned in the order of the given DocumentRefs.
// The return value will always contain the same number of DocumentSnapshots
// as the number of DocumentRefs in the input.
//
// If the same DocumentRef is specified multiple times in the input, the return
// value will contain the same number of DocumentSnapshots referencing the same
// document.
//
// If a document is not present, the corresponding DocumentSnapshot's Exists
// method will return false.
func (c *Client) GetAll(ctx context.Context, docRefs []*DocumentRef) (_ []*DocumentSnapshot, err error) {
	ctx = trace.StartSpan(ctx, "cloud.google.com/go/firestore.GetAll")
	defer func() { trace.EndSpan(ctx, err) }()

	return c.getAll(ctx, docRefs, nil)
}

func (c *Client) getAll(ctx context.Context, docRefs []*DocumentRef, tid []byte) ([]*DocumentSnapshot, error) {
	var docNames []string
	docIndices := map[string][]int{} // doc name to positions in docRefs
	for i, dr := range docRefs {
		if dr == nil {
			return nil, errNilDocRef
		}
		docNames = append(docNames, dr.Path)
		docIndices[dr.Path] = append(docIndices[dr.Path], i)
	}
	req := &pb.BatchGetDocumentsRequest{
		Database:  c.path(),
		Documents: docNames,
	}
	if tid != nil {
		req.ConsistencySelector = &pb.BatchGetDocumentsRequest_Transaction{tid}
	}
	streamClient, err := c.c.BatchGetDocuments(withResourceHeader(ctx, req.Database), req)
	if err != nil {
		return nil, err
	}

	// Read and remember all results from the stream.
	var resps []*pb.BatchGetDocumentsResponse
	for {
		resp, err := streamClient.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}
		resps = append(resps, resp)
	}

	// Results may arrive out of order. Put each at the right indices.
	docs := make([]*DocumentSnapshot, len(docNames))
	for _, resp := range resps {
		var (
			indices []int
			doc     *pb.Document
			err     error
		)
		switch r := resp.Result.(type) {
		case *pb.BatchGetDocumentsResponse_Found:
			indices = docIndices[r.Found.Name]
			doc = r.Found
		case *pb.BatchGetDocumentsResponse_Missing:
			indices = docIndices[r.Missing]
			doc = nil
		default:
			return nil, errors.New("firestore: unknown BatchGetDocumentsResponse result type")
		}
		for _, index := range indices {
			if docs[index] != nil {
				return nil, fmt.Errorf("firestore: %q seen twice", docRefs[index].Path)
			}
			docs[index], err = newDocumentSnapshot(docRefs[index], doc, c, resp.ReadTime)
			if err != nil {
				return nil, err
			}
		}
	}
	return docs, nil
}

// Collections returns an interator over the top-level collections.
func (c *Client) Collections(ctx context.Context) *CollectionIterator {
	it := &CollectionIterator{
		client: c,
		it: c.c.ListCollectionIds(
			withResourceHeader(ctx, c.path()),
			&pb.ListCollectionIdsRequest{Parent: c.path() + "/documents"}),
	}
	it.pageInfo, it.nextFunc = iterator.NewPageInfo(
		it.fetch,
		func() int { return len(it.items) },
		func() interface{} { b := it.items; it.items = nil; return b })
	return it
}

// Batch returns a WriteBatch.
func (c *Client) Batch() *WriteBatch {
	return &WriteBatch{c: c}
}

// commit calls the Commit RPC outside of a transaction.
func (c *Client) commit(ctx context.Context, ws []*pb.Write) ([]*WriteResult, error) {
	req := &pb.CommitRequest{
		Database: c.path(),
		Writes:   ws,
	}
	res, err := c.c.Commit(withResourceHeader(ctx, req.Database), req)
	if err != nil {
		return nil, err
	}
	if len(res.WriteResults) == 0 {
		return nil, errors.New("firestore: missing WriteResult")
	}
	var wrs []*WriteResult
	for _, pwr := range res.WriteResults {
		wr, err := writeResultFromProto(pwr)
		if err != nil {
			return nil, err
		}
		wrs = append(wrs, wr)
	}
	return wrs, nil
}

func (c *Client) commit1(ctx context.Context, ws []*pb.Write) (*WriteResult, error) {
	wrs, err := c.commit(ctx, ws)
	if err != nil {
		return nil, err
	}
	return wrs[0], nil
}

// A WriteResult is returned by methods that write documents.
type WriteResult struct {
	// The time at which the document was updated, or created if it did not
	// previously exist. Writes that do not actually change the document do
	// not change the update time.
	UpdateTime time.Time
}

func writeResultFromProto(wr *pb.WriteResult) (*WriteResult, error) {
	t, err := ptypes.Timestamp(wr.UpdateTime)
	if err != nil {
		t = time.Time{}
		// TODO(jba): Follow up if Delete is supposed to return a nil timestamp.
	}
	return &WriteResult{UpdateTime: t}, nil
}

func sleep(ctx context.Context, dur time.Duration) error {
	switch err := gax.Sleep(ctx, dur); err {
	case context.Canceled:
		return status.Error(codes.Canceled, "context canceled")
	case context.DeadlineExceeded:
		return status.Error(codes.DeadlineExceeded, "context deadline exceeded")
	default:
		return err
	}
}
<commit_msg>Set admin credentials when using the emulator<commit_after>// Copyright 2017 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package firestore

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	vkit "cloud.google.com/go/firestore/apiv1"
	"cloud.google.com/go/internal/trace"
	"cloud.google.com/go/internal/version"
	"github.com/golang/protobuf/ptypes"
	gax "github.com/googleapis/gax-go/v2"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
	"google.golang.org/api/transport"
	pb "google.golang.org/genproto/googleapis/firestore/v1"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

// resourcePrefixHeader is the name of the metadata header used to indicate
// the resource being operated on.
const resourcePrefixHeader = "google-cloud-resource-prefix"

// DetectProjectID is a sentinel value that instructs NewClient to detect the
// project ID. It is given in place of the projectID argument. NewClient will
// use the project ID from the given credentials or the default credentials
// (https://developers.google.com/accounts/docs/application-default-credentials)
// if no credentials were provided. When providing credentials, not all
// options will allow NewClient to extract the project ID. Specifically a JWT
// does not have the project ID encoded.
const DetectProjectID = "*detect-project-id*"

// A Client provides access to the Firestore service.
type Client struct {
	c          *vkit.Client
	projectID  string
	databaseID string // A client is tied to a single database.
}

// NewClient creates a new Firestore client that uses the given project.
func NewClient(ctx context.Context, projectID string, opts ...option.ClientOption) (*Client, error) {
	var o []option.ClientOption
	// If this environment variable is defined, configure the client to talk to the emulator.
	if addr := os.Getenv("FIRESTORE_EMULATOR_HOST"); addr != "" {
		conn, err := grpc.Dial(addr, grpc.WithInsecure(), grpc.WithPerRPCCredentials(emulatorCreds{}))
		if err != nil {
			return nil, fmt.Errorf("firestore: dialing address from env var FIRESTORE_EMULATOR_HOST: %s", err)
		}
		o = []option.ClientOption{option.WithGRPCConn(conn)}
	}
	o = append(o, opts...)

	if projectID == DetectProjectID {
		creds, err := transport.Creds(ctx, o...)
		if err != nil {
			return nil, fmt.Errorf("fetching creds: %v", err)
		}
		if creds.ProjectID == "" {
			return nil, errors.New("firestore: see the docs on DetectProjectID")
		}
		projectID = creds.ProjectID
	}

	vc, err := vkit.NewClient(ctx, o...)
	if err != nil {
		return nil, err
	}
	vc.SetGoogleClientInfo("gccl", version.Repo)
	c := &Client{
		c:          vc,
		projectID:  projectID,
		databaseID: "(default)", // always "(default)", for now
	}
	return c, nil

}

// Close closes any resources held by the client.
//
// Close need not be called at program exit.
func (c *Client) Close() error {
	return c.c.Close()
}

func (c *Client) path() string {
	return fmt.Sprintf("projects/%s/databases/%s", c.projectID, c.databaseID)
}

func withResourceHeader(ctx context.Context, resource string) context.Context {
	md, _ := metadata.FromOutgoingContext(ctx)
	md = md.Copy()
	md[resourcePrefixHeader] = []string{resource}
	return metadata.NewOutgoingContext(ctx, md)
}

// Collection creates a reference to a collection with the given path.
// A path is a sequence of IDs separated by slashes.
//
// Collection returns nil if path contains an even number of IDs or any ID is empty.
func (c *Client) Collection(path string) *CollectionRef {
	coll, _ := c.idsToRef(strings.Split(path, "/"), c.path())
	return coll
}

// Doc creates a reference to a document with the given path.
// A path is a sequence of IDs separated by slashes.
//
// Doc returns nil if path contains an odd number of IDs or any ID is empty.
func (c *Client) Doc(path string) *DocumentRef {
	_, doc := c.idsToRef(strings.Split(path, "/"), c.path())
	return doc
}

// CollectionGroup creates a reference to a group of collections that include
// the given ID, regardless of parent document.
//
// For example, consider:
// France/Cities/Paris = {population: 100}
// Canada/Cities/Montreal = {population: 90}
//
// CollectionGroup can be used to query across all "Cities" regardless of
// its parent "Countries". See ExampleCollectionGroup for a complete example.
func (c *Client) CollectionGroup(collectionID string) *CollectionGroupRef {
	return newCollectionGroupRef(c, c.path(), collectionID)
}

func (c *Client) idsToRef(IDs []string, dbPath string) (*CollectionRef, *DocumentRef) {
	if len(IDs) == 0 {
		return nil, nil
	}
	for _, id := range IDs {
		if id == "" {
			return nil, nil
		}
	}
	coll := newTopLevelCollRef(c, dbPath, IDs[0])
	i := 1
	for i < len(IDs) {
		doc := newDocRef(coll, IDs[i])
		i++
		if i == len(IDs) {
			return nil, doc
		}
		coll = newCollRefWithParent(c, doc, IDs[i])
		i++
	}
	return coll, nil
}

// GetAll retrieves multiple documents with a single call. The
// DocumentSnapshots are returned in the order of the given DocumentRefs.
// The return value will always contain the same number of DocumentSnapshots
// as the number of DocumentRefs in the input.
//
// If the same DocumentRef is specified multiple times in the input, the return
// value will contain the same number of DocumentSnapshots referencing the same
// document.
//
// If a document is not present, the corresponding DocumentSnapshot's Exists
// method will return false.
func (c *Client) GetAll(ctx context.Context, docRefs []*DocumentRef) (_ []*DocumentSnapshot, err error) {
	ctx = trace.StartSpan(ctx, "cloud.google.com/go/firestore.GetAll")
	defer func() { trace.EndSpan(ctx, err) }()

	return c.getAll(ctx, docRefs, nil)
}

func (c *Client) getAll(ctx context.Context, docRefs []*DocumentRef, tid []byte) ([]*DocumentSnapshot, error) {
	var docNames []string
	docIndices := map[string][]int{} // doc name to positions in docRefs
	for i, dr := range docRefs {
		if dr == nil {
			return nil, errNilDocRef
		}
		docNames = append(docNames, dr.Path)
		docIndices[dr.Path] = append(docIndices[dr.Path], i)
	}
	req := &pb.BatchGetDocumentsRequest{
		Database:  c.path(),
		Documents: docNames,
	}
	if tid != nil {
		req.ConsistencySelector = &pb.BatchGetDocumentsRequest_Transaction{tid}
	}
	streamClient, err := c.c.BatchGetDocuments(withResourceHeader(ctx, req.Database), req)
	if err != nil {
		return nil, err
	}

	// Read and remember all results from the stream.
	var resps []*pb.BatchGetDocumentsResponse
	for {
		resp, err := streamClient.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}
		resps = append(resps, resp)
	}

	// Results may arrive out of order. Put each at the right indices.
	docs := make([]*DocumentSnapshot, len(docNames))
	for _, resp := range resps {
		var (
			indices []int
			doc     *pb.Document
			err     error
		)
		switch r := resp.Result.(type) {
		case *pb.BatchGetDocumentsResponse_Found:
			indices = docIndices[r.Found.Name]
			doc = r.Found
		case *pb.BatchGetDocumentsResponse_Missing:
			indices = docIndices[r.Missing]
			doc = nil
		default:
			return nil, errors.New("firestore: unknown BatchGetDocumentsResponse result type")
		}
		for _, index := range indices {
			if docs[index] != nil {
				return nil, fmt.Errorf("firestore: %q seen twice", docRefs[index].Path)
			}
			docs[index], err = newDocumentSnapshot(docRefs[index], doc, c, resp.ReadTime)
			if err != nil {
				return nil, err
			}
		}
	}
	return docs, nil
}

// Collections returns an interator over the top-level collections.
func (c *Client) Collections(ctx context.Context) *CollectionIterator {
	it := &CollectionIterator{
		client: c,
		it: c.c.ListCollectionIds(
			withResourceHeader(ctx, c.path()),
			&pb.ListCollectionIdsRequest{Parent: c.path() + "/documents"}),
	}
	it.pageInfo, it.nextFunc = iterator.NewPageInfo(
		it.fetch,
		func() int { return len(it.items) },
		func() interface{} { b := it.items; it.items = nil; return b })
	return it
}

// Batch returns a WriteBatch.
func (c *Client) Batch() *WriteBatch {
	return &WriteBatch{c: c}
}

// commit calls the Commit RPC outside of a transaction.
func (c *Client) commit(ctx context.Context, ws []*pb.Write) ([]*WriteResult, error) {
	req := &pb.CommitRequest{
		Database: c.path(),
		Writes:   ws,
	}
	res, err := c.c.Commit(withResourceHeader(ctx, req.Database), req)
	if err != nil {
		return nil, err
	}
	if len(res.WriteResults) == 0 {
		return nil, errors.New("firestore: missing WriteResult")
	}
	var wrs []*WriteResult
	for _, pwr := range res.WriteResults {
		wr, err := writeResultFromProto(pwr)
		if err != nil {
			return nil, err
		}
		wrs = append(wrs, wr)
	}
	return wrs, nil
}

func (c *Client) commit1(ctx context.Context, ws []*pb.Write) (*WriteResult, error) {
	wrs, err := c.commit(ctx, ws)
	if err != nil {
		return nil, err
	}
	return wrs[0], nil
}

// A WriteResult is returned by methods that write documents.
type WriteResult struct {
	// The time at which the document was updated, or created if it did not
	// previously exist. Writes that do not actually change the document do
	// not change the update time.
	UpdateTime time.Time
}

func writeResultFromProto(wr *pb.WriteResult) (*WriteResult, error) {
	t, err := ptypes.Timestamp(wr.UpdateTime)
	if err != nil {
		t = time.Time{}
		// TODO(jba): Follow up if Delete is supposed to return a nil timestamp.
	}
	return &WriteResult{UpdateTime: t}, nil
}

func sleep(ctx context.Context, dur time.Duration) error {
	switch err := gax.Sleep(ctx, dur); err {
	case context.Canceled:
		return status.Error(codes.Canceled, "context canceled")
	case context.DeadlineExceeded:
		return status.Error(codes.DeadlineExceeded, "context deadline exceeded")
	default:
		return err
	}
}

// emulatorCreds is an instance of grpc.PerRPCCredentials that will configure a
// client to act as an admin for the Firestore emulator. It always hardcodes
// the "authorization" metadata field to contain "Bearer owner", which the
// Firestore emulator accepts as valid admin credentials.
type emulatorCreds struct{}

func (ec emulatorCreds) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
	return map[string]string{"authorization": "Bearer owner"}, nil
}
func (ec emulatorCreds) RequireTransportSecurity() bool {
	return false
}
<|endoftext|>


<commit_before>package modules

import (
	"fmt"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/pajbot/pajbot2/pkg"
	"github.com/pajbot/pajbot2/pkg/commands"
	mbase "github.com/pajbot/pajbot2/pkg/modules/base"
	"github.com/pajbot/pajbot2/pkg/twitchactions"
)

func init() {
	Register("nuke", func() pkg.ModuleSpec {
		return &Spec{
			id:    "nuke",
			name:  "Nuke",
			maker: newNuke,

			enabledByDefault: true,
		}
	})
}

const garbageCollectionInterval = 1 * time.Minute
const maxMessageAge = 5 * time.Minute

type nukeModule struct {
	mbase.Base

	messages      map[string][]nukeMessage
	messagesMutex sync.Mutex

	commands pkg.CommandsManager

	ticker *time.Ticker
}

type nukeMessage struct {
	user      pkg.User
	message   pkg.Message
	timestamp time.Time
}

func newNuke(b *mbase.Base) pkg.Module {
	m := &nukeModule{
		Base: *b,

		messages: make(map[string][]nukeMessage),

		commands: commands.NewCommands(),
	}

	m.commands.Register([]string{"!nuke"}, m)

	m.ticker = time.NewTicker(garbageCollectionInterval)

	go func() {
		for range m.ticker.C {
			m.garbageCollect()
		}
	}()

	return m
}

func (m *nukeModule) Trigger(parts []string, event pkg.MessageEvent) pkg.Actions {
	if !(event.User.IsModerator() || event.User.HasChannelPermission(m.BotChannel().Channel(), pkg.PermissionModeration)) {
		return nil
	}

	if len(parts) == 1 {
		return twitchactions.Mention(event.User, "usage: !nuke bad phrase 1m 10m")
	}
	phrase := strings.Join(parts[1:len(parts)-2], " ")
	scrollbackLength, err := time.ParseDuration(parts[len(parts)-2])
	if err != nil {
		return twitchactions.Mention(event.User, "usage: !nuke bad phrase 1m 10m")
	}
	if scrollbackLength < 0 {
		return twitchactions.Mention(event.User, "usage: !nuke bad phrase 1m 10m")
	}
	timeoutDuration, err := time.ParseDuration(parts[len(parts)-1])
	if err != nil {
		return twitchactions.Mention(event.User, "usage: !nuke bad phrase 1m 10m")
	}
	if timeoutDuration < 0 {
		return twitchactions.Mention(event.User, "usage: !nuke bad phrase 1m 10m")
	}

	m.nuke(event.User, m.BotChannel(), phrase, scrollbackLength, timeoutDuration)

	return nil
}

func (m *nukeModule) OnWhisper(event pkg.MessageEvent) pkg.Actions {
	return m.commands.OnMessage(event)
}

func (m *nukeModule) OnMessage(event pkg.MessageEvent) pkg.Actions {
	defer func() {
		m.addMessage(m.BotChannel().Channel(), event.User, event.Message)
	}()

	return m.commands.OnMessage(event)
}

func (m *nukeModule) garbageCollect() {
	m.messagesMutex.Lock()
	defer m.messagesMutex.Unlock()

	now := time.Now()

	for channelID := range m.messages {
		for i := 0; i < len(m.messages[channelID]); i++ {
			diff := now.Sub(m.messages[channelID][i].timestamp)
			if diff < maxMessageAge {
				m.messages[channelID] = m.messages[channelID][i:]
				break
			}
		}
	}
}

func (m *nukeModule) nuke(source pkg.User, bot pkg.BotChannel, phrase string, scrollbackLength, timeoutDuration time.Duration) {
	if timeoutDuration > 72*time.Hour {
		timeoutDuration = 72 * time.Hour
	}

	lowercasePhrase := strings.ToLower(phrase)

	matcher := func(msg *nukeMessage) bool {
		return strings.Contains(strings.ToLower(msg.message.GetText()), lowercasePhrase)
	}

	reason := "Nuked '" + phrase + "'"

	if strings.HasPrefix(phrase, "/") && strings.HasSuffix(phrase, "/") {
		regex, err := regexp.Compile(phrase[1 : len(phrase)-1])
		if err == nil {
			reason = "Nuked r'" + phrase[1:len(phrase)-1] + "'"
			matcher = func(msg *nukeMessage) bool {
				return regex.MatchString(msg.message.GetText())
			}
		}
		// parse as regex
	}

	now := time.Now()
	timeoutDurationInSeconds := int(timeoutDuration.Seconds())

	if timeoutDurationInSeconds < 1 {
		// Timeout duration too short
		return
	}

	targets := make(map[string]pkg.User)

	m.messagesMutex.Lock()
	defer m.messagesMutex.Unlock()

	messages := m.messages[bot.Channel().GetID()]

	for i := len(messages) - 1; i >= 0; i-- {
		diff := now.Sub(messages[i].timestamp)
		if diff > scrollbackLength {
			// We've gone far enough in the buffer, time to exit
			break
		}

		if matcher(&messages[i]) {
			targets[messages[i].user.GetID()] = messages[i].user
		}
	}

	for _, user := range targets {
		bot.SingleTimeout(user, timeoutDurationInSeconds, reason)
	}

	fmt.Printf("%s nuked %d users for the phrase %s in the last %s for %s\n", source.GetName(), len(targets), phrase, scrollbackLength, timeoutDuration)
}

func (m *nukeModule) addMessage(channel pkg.Channel, user pkg.User, message pkg.Message) {
	m.messagesMutex.Lock()
	defer m.messagesMutex.Unlock()

	m.messages[channel.GetID()] = append(m.messages[channel.GetID()], nukeMessage{
		user:      user,
		message:   message,
		timestamp: time.Now(),
	})
}
<commit_msg>fix nuke argument parsing xd<commit_after>package modules

import (
	"fmt"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/pajbot/pajbot2/pkg"
	"github.com/pajbot/pajbot2/pkg/commands"
	mbase "github.com/pajbot/pajbot2/pkg/modules/base"
	"github.com/pajbot/pajbot2/pkg/twitchactions"
)

func init() {
	Register("nuke", func() pkg.ModuleSpec {
		return &Spec{
			id:    "nuke",
			name:  "Nuke",
			maker: newNuke,

			enabledByDefault: true,
		}
	})
}

const garbageCollectionInterval = 1 * time.Minute
const maxMessageAge = 5 * time.Minute

type nukeModule struct {
	mbase.Base

	messages      map[string][]nukeMessage
	messagesMutex sync.Mutex

	commands pkg.CommandsManager

	ticker *time.Ticker
}

type nukeMessage struct {
	user      pkg.User
	message   pkg.Message
	timestamp time.Time
}

func newNuke(b *mbase.Base) pkg.Module {
	m := &nukeModule{
		Base: *b,

		messages: make(map[string][]nukeMessage),

		commands: commands.NewCommands(),
	}

	m.commands.Register([]string{"!nuke"}, m)

	m.ticker = time.NewTicker(garbageCollectionInterval)

	go func() {
		for range m.ticker.C {
			m.garbageCollect()
		}
	}()

	return m
}

func (m *nukeModule) Trigger(parts []string, event pkg.MessageEvent) pkg.Actions {
	if !(event.User.IsModerator() || event.User.HasChannelPermission(m.BotChannel().Channel(), pkg.PermissionModeration)) {
		return nil
	}

	if len(parts) < 3 {
		return twitchactions.Mention(event.User, "usage: !nuke bad phrase 1m 10m")
	}
	phrase := strings.Join(parts[1:len(parts)-2], " ")
	scrollbackLength, err := time.ParseDuration(parts[len(parts)-2])
	if err != nil {
		return twitchactions.Mention(event.User, "usage: !nuke bad phrase 1m 10m")
	}
	if scrollbackLength < 0 {
		return twitchactions.Mention(event.User, "usage: !nuke bad phrase 1m 10m")
	}
	timeoutDuration, err := time.ParseDuration(parts[len(parts)-1])
	if err != nil {
		return twitchactions.Mention(event.User, "usage: !nuke bad phrase 1m 10m")
	}
	if timeoutDuration < 0 {
		return twitchactions.Mention(event.User, "usage: !nuke bad phrase 1m 10m")
	}

	m.nuke(event.User, m.BotChannel(), phrase, scrollbackLength, timeoutDuration)

	return nil
}

func (m *nukeModule) OnWhisper(event pkg.MessageEvent) pkg.Actions {
	return m.commands.OnMessage(event)
}

func (m *nukeModule) OnMessage(event pkg.MessageEvent) pkg.Actions {
	defer func() {
		m.addMessage(m.BotChannel().Channel(), event.User, event.Message)
	}()

	return m.commands.OnMessage(event)
}

func (m *nukeModule) garbageCollect() {
	m.messagesMutex.Lock()
	defer m.messagesMutex.Unlock()

	now := time.Now()

	for channelID := range m.messages {
		for i := 0; i < len(m.messages[channelID]); i++ {
			diff := now.Sub(m.messages[channelID][i].timestamp)
			if diff < maxMessageAge {
				m.messages[channelID] = m.messages[channelID][i:]
				break
			}
		}
	}
}

func (m *nukeModule) nuke(source pkg.User, bot pkg.BotChannel, phrase string, scrollbackLength, timeoutDuration time.Duration) {
	if timeoutDuration > 72*time.Hour {
		timeoutDuration = 72 * time.Hour
	}

	lowercasePhrase := strings.ToLower(phrase)

	matcher := func(msg *nukeMessage) bool {
		return strings.Contains(strings.ToLower(msg.message.GetText()), lowercasePhrase)
	}

	reason := "Nuked '" + phrase + "'"

	if strings.HasPrefix(phrase, "/") && strings.HasSuffix(phrase, "/") {
		regex, err := regexp.Compile(phrase[1 : len(phrase)-1])
		if err == nil {
			reason = "Nuked r'" + phrase[1:len(phrase)-1] + "'"
			matcher = func(msg *nukeMessage) bool {
				return regex.MatchString(msg.message.GetText())
			}
		}
		// parse as regex
	}

	now := time.Now()
	timeoutDurationInSeconds := int(timeoutDuration.Seconds())

	if timeoutDurationInSeconds < 1 {
		// Timeout duration too short
		return
	}

	targets := make(map[string]pkg.User)

	m.messagesMutex.Lock()
	defer m.messagesMutex.Unlock()

	messages := m.messages[bot.Channel().GetID()]

	for i := len(messages) - 1; i >= 0; i-- {
		diff := now.Sub(messages[i].timestamp)
		if diff > scrollbackLength {
			// We've gone far enough in the buffer, time to exit
			break
		}

		if matcher(&messages[i]) {
			targets[messages[i].user.GetID()] = messages[i].user
		}
	}

	for _, user := range targets {
		bot.SingleTimeout(user, timeoutDurationInSeconds, reason)
	}

	fmt.Printf("%s nuked %d users for the phrase %s in the last %s for %s\n", source.GetName(), len(targets), phrase, scrollbackLength, timeoutDuration)
}

func (m *nukeModule) addMessage(channel pkg.Channel, user pkg.User, message pkg.Message) {
	m.messagesMutex.Lock()
	defer m.messagesMutex.Unlock()

	m.messages[channel.GetID()] = append(m.messages[channel.GetID()], nukeMessage{
		user:      user,
		message:   message,
		timestamp: time.Now(),
	})
}
<|endoftext|>


<commit_before>package moq

// taken from https://github.com/ernesto-jimenez/gogen
// Copyright (c) 2015 Ernesto Jiménez

import (
	"fmt"
	"go/ast"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"
)

type customImporter struct {
	source        string
	imported      map[string]*types.Package
	base          types.Importer
	skipTestFiles bool
}

func (i *customImporter) Import(path string) (*types.Package, error) {
	var err error
	if path == "" || path[0] == '.' {
		path, err = filepath.Abs(filepath.Clean(path))
		if err != nil {
			return nil, err
		}
		path = stripGopath(path)
	}
	if pkg, ok := i.imported[path]; ok {
		return pkg, nil
	}
	pkg, err := i.fsPkg(path)
	if err != nil {
		return nil, err
	}
	i.imported[path] = pkg
	return pkg, nil
}

func gopathDir(source, pkg string) (string, error) {
	// check vendor directory
	vendorPath, found := vendorPath(source, pkg)
	if found {
		return vendorPath, nil
	}
	for _, gopath := range gopaths() {
		absPath, err := filepath.Abs(path.Join(gopath, "src", pkg))
		if err != nil {
			return "", err
		}
		if dir, err := os.Stat(absPath); err == nil && dir.IsDir() {
			return absPath, nil
		}
	}
	return "", fmt.Errorf("%s not in $GOPATH or %s", pkg, path.Join(source, "vendor"))
}

func vendorPath(source, pkg string) (string, bool) {
	for {
		if isGopath(source) {
			return "", false
		}
		vendorPath, err := filepath.Abs(path.Join(source, "vendor", pkg))
		if err != nil {
			return "", false
		}
		if dir, err := os.Stat(vendorPath); err == nil && dir.IsDir() {
			return vendorPath, true
		}
		source = filepath.Dir(source)
		if source == "." {
			return "", false
		}
	}
}

func removeGopath(p string) string {
	for _, gopath := range gopaths() {
		p = strings.Replace(p, path.Join(gopath, "src")+"/", "", 1)
	}
	return p
}

func gopaths() []string {
	return strings.Split(os.Getenv("GOPATH"), string(filepath.ListSeparator))
}

func isGopath(path string) bool {
	for _, p := range gopaths() {
		if p == path {
			return true
		}
	}
	return false
}

func (i *customImporter) fsPkg(pkg string) (*types.Package, error) {
	dir, err := gopathDir(i.source, pkg)
	if err != nil {
		return importOrErr(i.base, pkg, err)
	}

	dirFiles, err := ioutil.ReadDir(dir)
	if err != nil {
		return importOrErr(i.base, pkg, err)
	}

	fset := token.NewFileSet()
	var files []*ast.File
	for _, fileInfo := range dirFiles {
		if fileInfo.IsDir() {
			continue
		}
		n := fileInfo.Name()
		if path.Ext(fileInfo.Name()) != ".go" {
			continue
		}
		if i.skipTestFiles && strings.Contains(fileInfo.Name(), "_test.go") {
			continue
		}
		file := path.Join(dir, n)
		src, err := ioutil.ReadFile(file)
		if err != nil {
			return nil, err
		}
		f, err := parser.ParseFile(fset, file, src, 0)
		if err != nil {
			return nil, err
		}
		files = append(files, f)
	}
	conf := types.Config{
		Importer: i,
	}
	p, err := conf.Check(pkg, fset, files, nil)

	if err != nil {
		return importOrErr(i.base, pkg, err)
	}
	return p, nil
}

func importOrErr(base types.Importer, pkg string, err error) (*types.Package, error) {
	p, impErr := base.Import(pkg)
	if impErr != nil {
		return nil, err
	}
	return p, nil
}

// newImporter returns an importer that will try to import code from gopath before using go/importer.Default and skipping test files
func newImporter(source string) types.Importer {
	return &customImporter{
		source:        source,
		imported:      make(map[string]*types.Package),
		base:          importer.Default(),
		skipTestFiles: true,
	}
}

// // DefaultWithTestFiles same as Default but it parses test files too
// func DefaultWithTestFiles() types.Importer {
// 	return &customImporter{
// 		imported:      make(map[string]*types.Package),
// 		base:          importer.Default(),
// 		skipTestFiles: false,
// 	}
// }

// stripGopath teks the directory to a package and remove the gopath to get the
// cannonical package name
func stripGopath(p string) string {
	for _, gopath := range gopaths() {
		p = strings.TrimPrefix(p, path.Join(gopath, "src")+"/")
	}
	return p
}
<commit_msg>better path parsing<commit_after>package moq

// taken from https://github.com/ernesto-jimenez/gogen
// Copyright (c) 2015 Ernesto Jiménez

import (
	"fmt"
	"go/ast"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"
)

type customImporter struct {
	source        string
	imported      map[string]*types.Package
	base          types.Importer
	skipTestFiles bool
}

func (i *customImporter) Import(path string) (*types.Package, error) {
	var err error
	if path == "" || path[0] == '.' {
		path, err = filepath.Abs(filepath.Clean(path))
		if err != nil {
			return nil, err
		}
		path = stripGopath(path)
	}
	if pkg, ok := i.imported[path]; ok {
		return pkg, nil
	}
	pkg, err := i.fsPkg(path)
	if err != nil {
		return nil, err
	}
	i.imported[path] = pkg
	return pkg, nil
}

func gopathDir(source, pkg string) (string, error) {
	// check vendor directory
	vendorPath, found := vendorPath(source, pkg)
	if found {
		return vendorPath, nil
	}
	for _, gopath := range gopaths() {
		absPath, err := filepath.Abs(path.Join(gopath, "src", pkg))
		if err != nil {
			return "", err
		}
		if dir, err := os.Stat(absPath); err == nil && dir.IsDir() {
			return absPath, nil
		}
	}
	return "", fmt.Errorf("%s not in $GOPATH or %s", pkg, path.Join(source, "vendor"))
}

func vendorPath(source, pkg string) (string, bool) {
	for {
		if isGopath(source) {
			return "", false
		}
		var err error
		source, err = filepath.Abs(source)
		if err != nil {
			return "", false
		}
		vendorPath, err := filepath.Abs(path.Join(source, "vendor", pkg))
		if err != nil {
			return "", false
		}
		if dir, err := os.Stat(vendorPath); err == nil && dir.IsDir() {
			return vendorPath, true
		}
		source = filepath.Dir(source)
	}
}

func removeGopath(p string) string {
	for _, gopath := range gopaths() {
		p = strings.Replace(p, path.Join(gopath, "src")+"/", "", 1)
	}
	return p
}

func gopaths() []string {
	return strings.Split(os.Getenv("GOPATH"), string(filepath.ListSeparator))
}

func isGopath(path string) bool {
	for _, p := range gopaths() {
		if p == path {
			return true
		}
	}
	return false
}

func (i *customImporter) fsPkg(pkg string) (*types.Package, error) {
	dir, err := gopathDir(i.source, pkg)
	if err != nil {
		return importOrErr(i.base, pkg, err)
	}

	dirFiles, err := ioutil.ReadDir(dir)
	if err != nil {
		return importOrErr(i.base, pkg, err)
	}

	fset := token.NewFileSet()
	var files []*ast.File
	for _, fileInfo := range dirFiles {
		if fileInfo.IsDir() {
			continue
		}
		n := fileInfo.Name()
		if path.Ext(fileInfo.Name()) != ".go" {
			continue
		}
		if i.skipTestFiles && strings.Contains(fileInfo.Name(), "_test.go") {
			continue
		}
		file := path.Join(dir, n)
		src, err := ioutil.ReadFile(file)
		if err != nil {
			return nil, err
		}
		f, err := parser.ParseFile(fset, file, src, 0)
		if err != nil {
			return nil, err
		}
		files = append(files, f)
	}
	conf := types.Config{
		Importer: i,
	}
	p, err := conf.Check(pkg, fset, files, nil)

	if err != nil {
		return importOrErr(i.base, pkg, err)
	}
	return p, nil
}

func importOrErr(base types.Importer, pkg string, err error) (*types.Package, error) {
	p, impErr := base.Import(pkg)
	if impErr != nil {
		return nil, err
	}
	return p, nil
}

// newImporter returns an importer that will try to import code from gopath before using go/importer.Default and skipping test files
func newImporter(source string) types.Importer {
	return &customImporter{
		source:        source,
		imported:      make(map[string]*types.Package),
		base:          importer.Default(),
		skipTestFiles: true,
	}
}

// // DefaultWithTestFiles same as Default but it parses test files too
// func DefaultWithTestFiles() types.Importer {
// 	return &customImporter{
// 		imported:      make(map[string]*types.Package),
// 		base:          importer.Default(),
// 		skipTestFiles: false,
// 	}
// }

// stripGopath teks the directory to a package and remove the gopath to get the
// cannonical package name
func stripGopath(p string) string {
	for _, gopath := range gopaths() {
		p = strings.TrimPrefix(p, path.Join(gopath, "src")+"/")
	}
	return p
}
<|endoftext|>


<commit_before>// Copyright 2018 the u-root 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 qemu

import (
	"fmt"
	"net"
	"sync/atomic"
)

// Device is a QEMU device to expose to a VM.
type Device interface {
	// Cmdline returns arguments to append to the QEMU command line for this device.
	Cmdline() []string
}

// Network is a Device that can connect multiple QEMU VMs to each other.
//
// Network uses the QEMU socket mechanism to connect multiple VMs with a simple
// TCP socket.
type Network struct {
	port uint16

	// numVMs must be atomically accessed so VMs can be started in parallel
	// in goroutines.
	numVMs uint32
}

func NewNetwork() *Network {
	return &Network{
		port: 1234,
	}
}

// Cmdline implements Device.
func (n *Network) Cmdline() []string {
	if n == nil {
		return nil
	}

	newNum := atomic.AddUint32(&n.numVMs, 1)
	num := newNum - 1

	// MAC for the virtualized NIC.
	//
	// This is from the range of locally administered address ranges.
	mac := net.HardwareAddr{0x0e, 0x00, 0x00, 0x00, 0x00, byte(num)}

	args := []string{"-net", fmt.Sprintf("nic,macaddr=%s", mac)}
	if num != 0 {
		args = append(args, "-net", fmt.Sprintf("socket,connect=:%d", n.port))
	} else {
		args = append(args, "-net", fmt.Sprintf("socket,listen=:%d", n.port))
	}
	return args
}

// ReadOnlyDirectory is a Device that exposes a directory as a /dev/sda1
// readonly vfat partition in the VM.
type ReadOnlyDirectory struct {
	// Dir is the directory to expose as a read-only vfat partition.
	Dir string
}

// Cmdline implements Device.
func (rod ReadOnlyDirectory) Cmdline() []string {
	if len(rod.Dir) == 0 {
		return nil
	}

	// Expose the temp directory to QEMU as /dev/sda1
	return []string{
		"-drive", fmt.Sprintf("file=fat:ro:%s,if=none,id=tmpdir", rod.Dir),
		"-device", "ich9-ahci,id=ahci",
		"-device", "ide-drive,drive=tmpdir,bus=ahci.0",
	}
}

// VirtioRandom exposes a PCI random number generator Device to the QEMU VM.
type VirtioRandom struct{}

// Cmdline implements Device.
func (VirtioRandom) Cmdline() []string {
	return []string{"-device", "virtio-rng-pci"}
}

// ArbitraryArgs allows users to add arbitrary arguments to the QEMU command
// line.
type ArbitraryArgs []string

// Cmdline implements Device.
func (aa ArbitraryArgs) Cmdline() []string {
	return aa
}
<commit_msg>qemu: workaround for some kind of read-only block device issue.<commit_after>// Copyright 2018 the u-root 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 qemu

import (
	"fmt"
	"net"
	"sync/atomic"
)

// Device is a QEMU device to expose to a VM.
type Device interface {
	// Cmdline returns arguments to append to the QEMU command line for this device.
	Cmdline() []string
}

// Network is a Device that can connect multiple QEMU VMs to each other.
//
// Network uses the QEMU socket mechanism to connect multiple VMs with a simple
// TCP socket.
type Network struct {
	port uint16

	// numVMs must be atomically accessed so VMs can be started in parallel
	// in goroutines.
	numVMs uint32
}

func NewNetwork() *Network {
	return &Network{
		port: 1234,
	}
}

// Cmdline implements Device.
func (n *Network) Cmdline() []string {
	if n == nil {
		return nil
	}

	newNum := atomic.AddUint32(&n.numVMs, 1)
	num := newNum - 1

	// MAC for the virtualized NIC.
	//
	// This is from the range of locally administered address ranges.
	mac := net.HardwareAddr{0x0e, 0x00, 0x00, 0x00, 0x00, byte(num)}

	args := []string{"-net", fmt.Sprintf("nic,macaddr=%s", mac)}
	if num != 0 {
		args = append(args, "-net", fmt.Sprintf("socket,connect=:%d", n.port))
	} else {
		args = append(args, "-net", fmt.Sprintf("socket,listen=:%d", n.port))
	}
	return args
}

// ReadOnlyDirectory is a Device that exposes a directory as a /dev/sda1
// readonly vfat partition in the VM.
type ReadOnlyDirectory struct {
	// Dir is the directory to expose as a read-only vfat partition.
	Dir string
}

// Cmdline implements Device.
func (rod ReadOnlyDirectory) Cmdline() []string {
	if len(rod.Dir) == 0 {
		return nil
	}

	// Expose the temp directory to QEMU as /dev/sda1
	return []string{
		"-drive", fmt.Sprintf("file=fat:rw:%s,if=none,id=tmpdir", rod.Dir),
		"-device", "ich9-ahci,id=ahci",
		"-device", "ide-drive,drive=tmpdir,bus=ahci.0",
	}
}

// VirtioRandom exposes a PCI random number generator Device to the QEMU VM.
type VirtioRandom struct{}

// Cmdline implements Device.
func (VirtioRandom) Cmdline() []string {
	return []string{"-device", "virtio-rng-pci"}
}

// ArbitraryArgs allows users to add arbitrary arguments to the QEMU command
// line.
type ArbitraryArgs []string

// Cmdline implements Device.
func (aa ArbitraryArgs) Cmdline() []string {
	return aa
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google LLC

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

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

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

package util

import (
	"bufio"
	pkgutil "github.com/GoogleCloudPlatform/container-diff/pkg/util"
	"github.com/GoogleCloudPlatform/k8s-container-builder/pkg/constants"
	"github.com/containers/image/docker"
	"github.com/sirupsen/logrus"
	"io"
	"os"
	"strings"
)

// ExtractFileSystemFromImage pulls an image and unpacks it to a file system at root
func ExtractFileSystemFromImage(img string) error {
	ref, err := docker.ParseReference("//" + img)
	if err != nil {
		return err
	}
	imgSrc, err := ref.NewImageSource(nil)
	if err != nil {
		return err
	}
	whitelist, err := fileSystemWhitelist(constants.WhitelistPath)
	if err != nil {
		return err
	}
	logrus.Infof("Whitelisted directories are %s", whitelist)
	return pkgutil.GetFileSystemFromReference(ref, imgSrc, constants.RootDir, whitelist)
}

// Get whitelist from roots of mounted files
// Each line of /proc/self/mountinfo is in the form:
// 36 35 98:0 /mnt1 /mnt2 rw,noatime master:1 - ext3 /dev/root rw,errors=continue
// (1)(2)(3)   (4)   (5)      (6)      (7)   (8) (9)   (10)         (11)
// Where (5) is the mount point relative to the process's root
// From: https://www.kernel.org/doc/Documentation/filesystems/proc.txt
func fileSystemWhitelist(path string) ([]string, error) {
	whitelist := []string{"/work-dir", "/dockerfile"}
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	reader := bufio.NewReader(f)
	for {
		line, err := reader.ReadString('\n')
		if err != nil && err != io.EOF {
			return nil, err
		}
		lineArr := strings.Split(line, " ")
		if len(lineArr) < 5 {
			if err == io.EOF {
				break
			}
			continue
		}
		if lineArr[4] != constants.RootDir {
			whitelist = append(whitelist, lineArr[4])
		}
		if err == io.EOF {
			break
		}
	}
	return whitelist, nil
}
<commit_msg>Added logging statements<commit_after>/*
Copyright 2018 Google LLC

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

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

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

package util

import (
	"bufio"
	pkgutil "github.com/GoogleCloudPlatform/container-diff/pkg/util"
	"github.com/GoogleCloudPlatform/k8s-container-builder/pkg/constants"
	"github.com/containers/image/docker"
	"github.com/sirupsen/logrus"
	"io"
	"os"
	"strings"
)

var whitelist = []string{"/work-dir", "/dockerfile"}

// ExtractFileSystemFromImage pulls an image and unpacks it to a file system at root
func ExtractFileSystemFromImage(img string) error {
	ref, err := docker.ParseReference("//" + img)
	if err != nil {
		return err
	}
	imgSrc, err := ref.NewImageSource(nil)
	if err != nil {
		return err
	}
	whitelist, err := fileSystemWhitelist(constants.WhitelistPath)
	if err != nil {
		return err
	}
	logrus.Infof("Whitelisted directories are %s", whitelist)
	return pkgutil.GetFileSystemFromReference(ref, imgSrc, constants.RootDir, whitelist)
}

// Get whitelist from roots of mounted files
// Each line of /proc/self/mountinfo is in the form:
// 36 35 98:0 /mnt1 /mnt2 rw,noatime master:1 - ext3 /dev/root rw,errors=continue
// (1)(2)(3)   (4)   (5)      (6)      (7)   (8) (9)   (10)         (11)
// Where (5) is the mount point relative to the process's root
// From: https://www.kernel.org/doc/Documentation/filesystems/proc.txt
func fileSystemWhitelist(path string) ([]string, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	reader := bufio.NewReader(f)
	for {
		line, err := reader.ReadString('\n')
		logrus.Debugf("Read the following line from %s: %s", path, line)
		if err != nil && err != io.EOF {
			return nil, err
		}
		lineArr := strings.Split(line, " ")
		if len(lineArr) < 5 {
			if err == io.EOF {
				logrus.Debugf("Reached end of file %s", path)
				break
			}
			continue
		}
		if lineArr[4] != constants.RootDir {
			logrus.Debugf("Appending %s from line: %s", lineArr[4], line)
			whitelist = append(whitelist, lineArr[4])
		}
		if err == io.EOF {
			logrus.Debugf("Reached end of file %s", path)
			break
		}
	}
	return whitelist, nil
}
<|endoftext|>


<commit_before>package peco

import (
	"bufio"
	"bytes"
	"os/exec"
	"regexp"
	"sort"
	"sync"

	"github.com/lestrrat/go-pdebug"
	"github.com/peco/peco/hub"
	"github.com/peco/peco/internal/util"
	"github.com/peco/peco/pipeline"
	"github.com/pkg/errors"
	"golang.org/x/net/context"
)

func (fx *FilterSet) Reset() {
	fx.current = 0
}

func (fs *FilterSet) Size() int {
	return len(fs.filters)
}

func (fs *FilterSet) Add(lf LineFilter) error {
	fs.filters = append(fs.filters, lf)
	return nil
}

func (fs *FilterSet) Rotate() {
	fs.current++
	if fs.current >= len(fs.filters) {
		fs.current = 0
	}
	if pdebug.Enabled {
		pdebug.Printf("FilterSet.Rotate: now filter in effect is %s", fs.filters[fs.current])
	}
}

func (fs *FilterSet) SetCurrentByName(name string) error {
	for i, f := range fs.filters {
		if f.String() == name {
			fs.current = i
			return nil
		}
	}
	return ErrFilterNotFound
}

func (fs *FilterSet) Current() LineFilter {
	return fs.filters[fs.current]
}

func NewFilter(state *Peco) *Filter {
	return &Filter{
		state: state,
	}
}

// Work is the actual work horse that that does the matching
// in a goroutine of its own. It wraps Matcher.Match().
func (f *Filter) Work(ctx context.Context, q hub.Payload) {
	defer q.Done()

	query, ok := q.Data().(string)
	if !ok {
		return
	}

	state := f.state
	if query == "" {
		state.ResetCurrentLineBuffer()
	} else {
		// Create a new pipeline
		p := pipeline.New()
		p.SetSource(state.Source())
		f := state.Filters().Current().Clone()
		f.SetQuery(query)
		p.Add(f)

		buf := NewMemoryBuffer()
		p.SetDestination(buf)
		state.SetCurrentLineBuffer(buf)

		go func() {
			defer state.Hub().SendDraw(true)
			if err := p.Run(ctx); err != nil {
				state.Hub().SendStatusMsg(err.Error())
			}
		}()

		go func() {
			if pdebug.Enabled {
				pdebug.Printf("waiting for query to finish")
				defer pdebug.Printf("Filter.Work: finished running query")
			}
			<-p.Done()
			state.Hub().SendStatusMsg("")
		}()
	}

	if !state.config.StickySelection {
		state.Selection().Reset()
	}
}

// Loop keeps watching for incoming queries, and upon receiving
// a query, spawns a goroutine to do the heavy work. It also
// checks for previously running queries, so we can avoid
// running many goroutines doing the grep at the same time
func (f *Filter) Loop(ctx context.Context, cancel func()) error {
	defer cancel()

	// previous holds the function that can cancel the previous
	// query. This is used when multiple queries come in succession
	// and the previous query is discarded anyway
	var mutex sync.Mutex
	var previous func()
	previous = func() {} // no op func
	for {
		select {
		case <-ctx.Done():
			return nil
		case q := <-f.state.Hub().QueryCh():
			workctx, workcancel := context.WithCancel(ctx)

			mutex.Lock()
			previous()
			previous = workcancel
			mutex.Unlock()

			f.state.Hub().SendStatusMsg("Running query...")

			go func() {
				f.Work(workctx, q)
			}()
		}
	}
}

func NewRegexpFilter() *RegexpFilter {
	return &RegexpFilter{
		flags: regexpFlagList(defaultFlags),
		name:  "Regexp",
		outCh: pipeline.OutputChannel(make(chan interface{})),
	}
}

func (rf RegexpFilter) OutCh() <-chan interface{} {
	return rf.outCh
}

func (rf RegexpFilter) Clone() LineFilter {
	return &RegexpFilter{
		flags:     rf.flags,
		quotemeta: rf.quotemeta,
		query:     rf.query,
		name:      rf.name,
		outCh:     pipeline.OutputChannel(make(chan interface{})),
	}
}

func (rf *RegexpFilter) Accept(ctx context.Context, p pipeline.Producer) {
	if pdebug.Enabled {
		g := pdebug.Marker("RegexpFilter.Accept")
		defer g.End()
	}
	defer rf.outCh.SendEndMark("end of RegexpFilter")
	for {
		select {
		case <-ctx.Done():
			if pdebug.Enabled {
				pdebug.Printf("RegexpFilter received done")
			}
			return
		case v := <-p.OutCh():
			switch v.(type) {
			case error:
				if pipeline.IsEndMark(v.(error)) {
					if pdebug.Enabled {
						pdebug.Printf("RegexpFilter received end mark")
					}
					return
				}
			case Line:
				if l, err := rf.filter(v.(Line)); err == nil {
					rf.outCh.Send(l)
				}
			}
		}
	}
}

func (rf *RegexpFilter) filter(l Line) (Line, error) {
	regexps, err := rf.getQueryAsRegexps()
	if err != nil {
		return nil, errors.Wrap(err, "failed to compile queries as regular expression")
	}
	v := l.DisplayString()
	allMatched := true
	matches := [][]int{}
TryRegexps:
	for _, rx := range regexps {
		match := rx.FindAllStringSubmatchIndex(v, -1)
		if match == nil {
			allMatched = false
			break TryRegexps
		}
		matches = append(matches, match...)
	}

	if !allMatched {
		return nil, errors.New("filter did not match against given line")
	}

	sort.Sort(byMatchStart(matches))

	// We need to "dedupe" the results. For example, if we matched the
	// same region twice, we don't want that to be drawn

	deduped := make([][]int, 0, len(matches))

	for i, m := range matches {
		// Always push the first one
		if i == 0 {
			deduped = append(deduped, m)
			continue
		}

		prev := deduped[len(deduped)-1]
		switch {
		case matchContains(prev, m):
			// If the previous match contains this one, then
			// don't do anything
			continue
		case matchOverlaps(prev, m):
			// If the previous match overlaps with this one,
			// merge the results and make it a bigger one
			deduped[len(deduped)-1] = mergeMatches(prev, m)
		default:
			deduped = append(deduped, m)
		}
	}
	return NewMatchedLine(l, deduped), nil
}

func (rf *RegexpFilter) getQueryAsRegexps() ([]*regexp.Regexp, error) {
	if q := rf.compiledQuery; q != nil {
		return q, nil
	}
	q, err := queryToRegexps(rf.flags, rf.quotemeta, rf.query)
	if err != nil {
		return nil, errors.Wrap(err, "failed to compile queries as regular expression")
	}

	rf.compiledQuery = q
	return q, nil
}

func (rf *RegexpFilter) SetQuery(q string) {
	rf.query = q
	rf.compiledQuery = nil
}

func (rf RegexpFilter) String() string {
	return rf.name
}

var ErrFilterNotFound = errors.New("specified filter was not found")

func NewIgnoreCaseFilter() *RegexpFilter {
	rf := NewRegexpFilter()
	rf.flags = ignoreCaseFlags
	rf.quotemeta = true
	rf.name = "IgnoreCase"
	return rf
}

func NewCaseSensitiveFilter() *RegexpFilter {
	rf := NewRegexpFilter()
	rf.quotemeta = true
	rf.name = "CaseSensitive"
	return rf
}

// SmartCaseFilter turns ON the ignore-case flag in the regexp
// if the query contains a upper-case character
func NewSmartCaseFilter() *RegexpFilter {
	rf := NewRegexpFilter()
	rf.quotemeta = true
	rf.name = "SmartCase"
	rf.flags = regexpFlagFunc(func(q string) []string {
		if util.ContainsUpper(q) {
			return defaultFlags
		}
		return []string{"i"}
	})
	return rf
}

func NewExternalCmdFilter(name, cmd string, args []string, threshold int, enableSep bool) *ExternalCmdFilter {
	if len(args) == 0 {
		args = []string{"$QUERY"}
	}

	return &ExternalCmdFilter{
		enableSep:       enableSep,
		cmd:             cmd,
		args:            args,
		name:            name,
		thresholdBufsiz: threshold,
		outCh:           pipeline.OutputChannel(make(chan interface{})),
	}
}

func (ecf ExternalCmdFilter) Clone() LineFilter {
	return &ExternalCmdFilter{
		enableSep:       ecf.enableSep,
		cmd:             ecf.cmd,
		args:            ecf.args,
		name:            ecf.name,
		thresholdBufsiz: ecf.thresholdBufsiz,
		outCh:           pipeline.OutputChannel(make(chan interface{})),
	}
}

func (ecf *ExternalCmdFilter) Verify() error {
	if ecf.cmd == "" {
		return errors.Errorf("no executable specified for custom matcher '%s'", ecf.name)
	}

	if _, err := exec.LookPath(ecf.cmd); err != nil {
		return errors.Wrap(err, "failed to locate command")
	}
	return nil
}

func (ecf *ExternalCmdFilter) Accept(ctx context.Context, p pipeline.Producer) {
	if pdebug.Enabled {
		g := pdebug.Marker("ExternalCmdFilter.Accept")
		defer g.End()
	}
	defer ecf.outCh.SendEndMark("end of ExternalCmdFilter")

	buf := make([]Line, 0, ecf.thresholdBufsiz)
	for {
		select {
		case <-ctx.Done():
			if pdebug.Enabled {
				pdebug.Printf("ExternalCmdFilter received done")
			}
			return
		case v := <-p.OutCh():
			switch v.(type) {
			case error:
				if pipeline.IsEndMark(v.(error)) {
					if pdebug.Enabled {
						pdebug.Printf("ExternalCmdFilter received end mark")
					}
					if len(buf) > 0 {
						ecf.launchExternalCmd(ctx, buf)
					}
				}
				return
			case Line:
				if pdebug.Enabled {
					pdebug.Printf("ExternalCmdFilter received new line")
				}
				buf = append(buf, v.(Line))
				if len(buf) < ecf.thresholdBufsiz {
					continue
				}

				ecf.launchExternalCmd(ctx, buf)
				buf = buf[0:0]
			}
		}
	}
}

func (ecf ExternalCmdFilter) OutCh() <-chan interface{} {
	return ecf.outCh
}

func (ecf *ExternalCmdFilter) SetQuery(q string) {
	ecf.query = q
}

func (ecf ExternalCmdFilter) String() string {
	return ecf.name
}

func (ecf *ExternalCmdFilter) launchExternalCmd(ctx context.Context, buf []Line) {
	defer func() { recover() }() // ignore errors
	if pdebug.Enabled {
		g := pdebug.Marker("ExternalCmdFilter.launchExternalCmd")
		defer g.End()
	}

	args := append([]string(nil), ecf.args...)
	for i, v := range args {
		if v == "$QUERY" {
			args[i] = ecf.query
		}
	}
	cmd := exec.Command(ecf.cmd, args...)

	inbuf := &bytes.Buffer{}
	for _, l := range buf {
		inbuf.WriteString(l.DisplayString() + "\n")
	}

	cmd.Stdin = inbuf
	r, err := cmd.StdoutPipe()
	if err != nil {
		return
	}

	err = cmd.Start()
	if err != nil {
		return
	}

	go cmd.Wait()

	cmdCh := make(chan Line)
	go func(cmdCh chan Line, rdr *bufio.Reader) {
		defer func() { recover() }()
		defer close(cmdCh)
		for {
			b, _, err := rdr.ReadLine()
			if len(b) > 0 {
				// TODO: need to redo the spec for custom matchers
				// This is the ONLY location where we need to actually
				// RECREATE a RawLine, and thus the only place where
				// ctx.enableSep is required.
				cmdCh <- NewMatchedLine(NewRawLine(string(b), ecf.enableSep), nil)
			}
			if err != nil {
				break
			}
		}
	}(cmdCh, bufio.NewReader(r))

	defer func() {
		if p := cmd.Process; p != nil {
			p.Kill()
		}
	}()

	for {
		select {
		case <-ctx.Done():
			return
		case l, ok := <-cmdCh:
			if l == nil || !ok {
				return
			}
			ecf.outCh.Send(l)
		}
	}
}
<commit_msg>properly use the default buffer threshold<commit_after>package peco

import (
	"bufio"
	"bytes"
	"os/exec"
	"regexp"
	"sort"
	"sync"

	"github.com/lestrrat/go-pdebug"
	"github.com/peco/peco/hub"
	"github.com/peco/peco/internal/util"
	"github.com/peco/peco/pipeline"
	"github.com/pkg/errors"
	"golang.org/x/net/context"
)

func (fx *FilterSet) Reset() {
	fx.current = 0
}

func (fs *FilterSet) Size() int {
	return len(fs.filters)
}

func (fs *FilterSet) Add(lf LineFilter) error {
	fs.filters = append(fs.filters, lf)
	return nil
}

func (fs *FilterSet) Rotate() {
	fs.current++
	if fs.current >= len(fs.filters) {
		fs.current = 0
	}
	if pdebug.Enabled {
		pdebug.Printf("FilterSet.Rotate: now filter in effect is %s", fs.filters[fs.current])
	}
}

func (fs *FilterSet) SetCurrentByName(name string) error {
	for i, f := range fs.filters {
		if f.String() == name {
			fs.current = i
			return nil
		}
	}
	return ErrFilterNotFound
}

func (fs *FilterSet) Current() LineFilter {
	return fs.filters[fs.current]
}

func NewFilter(state *Peco) *Filter {
	return &Filter{
		state: state,
	}
}

// Work is the actual work horse that that does the matching
// in a goroutine of its own. It wraps Matcher.Match().
func (f *Filter) Work(ctx context.Context, q hub.Payload) {
	defer q.Done()

	query, ok := q.Data().(string)
	if !ok {
		return
	}

	if pdebug.Enabled {
		pdebug.Printf("New query '%s'", query)
	}

	state := f.state
	if query == "" {
		state.ResetCurrentLineBuffer()
		if !state.config.StickySelection {
			state.Selection().Reset()
		}
		return
	}

	// Create a new pipeline
	p := pipeline.New()
	p.SetSource(state.Source())
	thisf := state.Filters().Current().Clone()
	thisf.SetQuery(query)
	p.Add(thisf)

	buf := NewMemoryBuffer()
	p.SetDestination(buf)
	state.SetCurrentLineBuffer(buf)

	go func() {
		defer state.Hub().SendDraw(true)
		if err := p.Run(ctx); err != nil {
			state.Hub().SendStatusMsg(err.Error())
		}
	}()

	go func() {
		if pdebug.Enabled {
			pdebug.Printf("waiting for query to finish")
			defer pdebug.Printf("Filter.Work: finished running query")
		}
		<-p.Done()
		state.Hub().SendStatusMsg("")
	}()

	if !state.config.StickySelection {
		state.Selection().Reset()
	}
}

// Loop keeps watching for incoming queries, and upon receiving
// a query, spawns a goroutine to do the heavy work. It also
// checks for previously running queries, so we can avoid
// running many goroutines doing the grep at the same time
func (f *Filter) Loop(ctx context.Context, cancel func()) error {
	defer cancel()

	// previous holds the function that can cancel the previous
	// query. This is used when multiple queries come in succession
	// and the previous query is discarded anyway
	var mutex sync.Mutex
	var previous func()
	previous = func() {} // no op func
	for {
		select {
		case <-ctx.Done():
			return nil
		case q := <-f.state.Hub().QueryCh():
			workctx, workcancel := context.WithCancel(ctx)

			mutex.Lock()
			previous()
			previous = workcancel
			mutex.Unlock()

			f.state.Hub().SendStatusMsg("Running query...")

			go func() {
				f.Work(workctx, q)
			}()
		}
	}
}

func NewRegexpFilter() *RegexpFilter {
	return &RegexpFilter{
		flags: regexpFlagList(defaultFlags),
		name:  "Regexp",
		outCh: pipeline.OutputChannel(make(chan interface{})),
	}
}

func (rf RegexpFilter) OutCh() <-chan interface{} {
	return rf.outCh
}

func (rf RegexpFilter) Clone() LineFilter {
	return &RegexpFilter{
		flags:     rf.flags,
		quotemeta: rf.quotemeta,
		query:     rf.query,
		name:      rf.name,
		outCh:     pipeline.OutputChannel(make(chan interface{})),
	}
}

func (rf *RegexpFilter) Accept(ctx context.Context, p pipeline.Producer) {
	if pdebug.Enabled {
		g := pdebug.Marker("RegexpFilter.Accept")
		defer g.End()
	}
	defer rf.outCh.SendEndMark("end of RegexpFilter")
	for {
		select {
		case <-ctx.Done():
			if pdebug.Enabled {
				pdebug.Printf("RegexpFilter received done")
			}
			return
		case v := <-p.OutCh():
			switch v.(type) {
			case error:
				if pipeline.IsEndMark(v.(error)) {
					if pdebug.Enabled {
						pdebug.Printf("RegexpFilter received end mark")
					}
					return
				}
			case Line:
				if l, err := rf.filter(v.(Line)); err == nil {
					rf.outCh.Send(l)
				}
			}
		}
	}
}

func (rf *RegexpFilter) filter(l Line) (Line, error) {
	regexps, err := rf.getQueryAsRegexps()
	if err != nil {
		return nil, errors.Wrap(err, "failed to compile queries as regular expression")
	}
	v := l.DisplayString()
	allMatched := true
	matches := [][]int{}
TryRegexps:
	for _, rx := range regexps {
		match := rx.FindAllStringSubmatchIndex(v, -1)
		if match == nil {
			allMatched = false
			break TryRegexps
		}
		matches = append(matches, match...)
	}

	if !allMatched {
		return nil, errors.New("filter did not match against given line")
	}

	sort.Sort(byMatchStart(matches))

	// We need to "dedupe" the results. For example, if we matched the
	// same region twice, we don't want that to be drawn

	deduped := make([][]int, 0, len(matches))

	for i, m := range matches {
		// Always push the first one
		if i == 0 {
			deduped = append(deduped, m)
			continue
		}

		prev := deduped[len(deduped)-1]
		switch {
		case matchContains(prev, m):
			// If the previous match contains this one, then
			// don't do anything
			continue
		case matchOverlaps(prev, m):
			// If the previous match overlaps with this one,
			// merge the results and make it a bigger one
			deduped[len(deduped)-1] = mergeMatches(prev, m)
		default:
			deduped = append(deduped, m)
		}
	}
	return NewMatchedLine(l, deduped), nil
}

func (rf *RegexpFilter) getQueryAsRegexps() ([]*regexp.Regexp, error) {
	if q := rf.compiledQuery; q != nil {
		return q, nil
	}
	q, err := queryToRegexps(rf.flags, rf.quotemeta, rf.query)
	if err != nil {
		return nil, errors.Wrap(err, "failed to compile queries as regular expression")
	}

	rf.compiledQuery = q
	return q, nil
}

func (rf *RegexpFilter) SetQuery(q string) {
	rf.query = q
	rf.compiledQuery = nil
}

func (rf RegexpFilter) String() string {
	return rf.name
}

var ErrFilterNotFound = errors.New("specified filter was not found")

func NewIgnoreCaseFilter() *RegexpFilter {
	rf := NewRegexpFilter()
	rf.flags = ignoreCaseFlags
	rf.quotemeta = true
	rf.name = "IgnoreCase"
	return rf
}

func NewCaseSensitiveFilter() *RegexpFilter {
	rf := NewRegexpFilter()
	rf.quotemeta = true
	rf.name = "CaseSensitive"
	return rf
}

// SmartCaseFilter turns ON the ignore-case flag in the regexp
// if the query contains a upper-case character
func NewSmartCaseFilter() *RegexpFilter {
	rf := NewRegexpFilter()
	rf.quotemeta = true
	rf.name = "SmartCase"
	rf.flags = regexpFlagFunc(func(q string) []string {
		if util.ContainsUpper(q) {
			return defaultFlags
		}
		return []string{"i"}
	})
	return rf
}

func NewExternalCmdFilter(name string, cmd string, args []string, threshold int, enableSep bool) *ExternalCmdFilter {
	if len(args) == 0 {
		args = []string{"$QUERY"}
	}

	if threshold <= 0 {
		threshold = DefaultCustomFilterBufferThreshold
	}

	return &ExternalCmdFilter{
		args:            args,
		cmd:             cmd,
		enableSep:       enableSep,
		name:            name,
		outCh:           pipeline.OutputChannel(make(chan interface{})),
		thresholdBufsiz: threshold,
	}
}

func (ecf ExternalCmdFilter) Clone() LineFilter {
	return &ExternalCmdFilter{
		args:            ecf.args,
		cmd:             ecf.cmd,
		enableSep:       ecf.enableSep,
		name:            ecf.name,
		outCh:           pipeline.OutputChannel(make(chan interface{})),
		thresholdBufsiz: ecf.thresholdBufsiz,
	}
}

func (ecf *ExternalCmdFilter) Verify() error {
	if ecf.cmd == "" {
		return errors.Errorf("no executable specified for custom matcher '%s'", ecf.name)
	}

	if _, err := exec.LookPath(ecf.cmd); err != nil {
		return errors.Wrap(err, "failed to locate command")
	}
	return nil
}

func (ecf *ExternalCmdFilter) Accept(ctx context.Context, p pipeline.Producer) {
	if pdebug.Enabled {
		g := pdebug.Marker("ExternalCmdFilter.Accept")
		defer g.End()
	}
	defer ecf.outCh.SendEndMark("end of ExternalCmdFilter")

	buf := make([]Line, 0, ecf.thresholdBufsiz)
	for {
		select {
		case <-ctx.Done():
			if pdebug.Enabled {
				pdebug.Printf("ExternalCmdFilter received done")
			}
			return
		case v := <-p.OutCh():
			switch v.(type) {
			case error:
				if pipeline.IsEndMark(v.(error)) {
					if pdebug.Enabled {
						pdebug.Printf("ExternalCmdFilter received end mark")
					}
					if len(buf) > 0 {
						ecf.launchExternalCmd(ctx, buf)
					}
				}
				return
			case Line:
				if pdebug.Enabled {
					pdebug.Printf("ExternalCmdFilter received new line")
				}
				buf = append(buf, v.(Line))
				if len(buf) < ecf.thresholdBufsiz {
					continue
				}

				ecf.launchExternalCmd(ctx, buf)
				buf = buf[0:0]
			}
		}
	}
}

func (ecf ExternalCmdFilter) OutCh() <-chan interface{} {
	return ecf.outCh
}

func (ecf *ExternalCmdFilter) SetQuery(q string) {
	ecf.query = q
}

func (ecf ExternalCmdFilter) String() string {
	return ecf.name
}

func (ecf *ExternalCmdFilter) launchExternalCmd(ctx context.Context, buf []Line) {
	defer func() { recover() }() // ignore errors
	if pdebug.Enabled {
		g := pdebug.Marker("ExternalCmdFilter.launchExternalCmd")
		defer g.End()
	}

	args := append([]string(nil), ecf.args...)
	for i, v := range args {
		if v == "$QUERY" {
			args[i] = ecf.query
		}
	}
	cmd := exec.Command(ecf.cmd, args...)
	if pdebug.Enabled {
		pdebug.Printf("Executing command %s %v", cmd.Path, cmd.Args)
	}

	inbuf := &bytes.Buffer{}
	for _, l := range buf {
		inbuf.WriteString(l.DisplayString() + "\n")
	}

	cmd.Stdin = inbuf
	r, err := cmd.StdoutPipe()
	if err != nil {
		return
	}

	err = cmd.Start()
	if err != nil {
		return
	}

	go cmd.Wait()

	cmdCh := make(chan Line)
	go func(cmdCh chan Line, rdr *bufio.Reader) {
		defer func() { recover() }()
		defer close(cmdCh)
		for {
			b, _, err := rdr.ReadLine()
			if len(b) > 0 {
				// TODO: need to redo the spec for custom matchers
				// This is the ONLY location where we need to actually
				// RECREATE a RawLine, and thus the only place where
				// ctx.enableSep is required.
				cmdCh <- NewMatchedLine(NewRawLine(string(b), ecf.enableSep), nil)
			}
			if err != nil {
				break
			}
		}
	}(cmdCh, bufio.NewReader(r))

	defer func() {
		if p := cmd.Process; p != nil {
			p.Kill()
		}
	}()

	for {
		select {
		case <-ctx.Done():
			return
		case l, ok := <-cmdCh:
			if l == nil || !ok {
				return
			}
			ecf.outCh.Send(l)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package goroutines

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"runtime"
	"strconv"
	"strings"
)

func getStackBuffer(all bool) []byte {
	for i := 4096; ; i *= 2 {
		buf := make([]byte, i)
		if n := runtime.Stack(buf, all); n < i {
			return buf
		}
	}
}

// parseGoStackHeader parses a stack header that looks like:
// goroutine 643 [runnable]:\n
// And returns the goroutine ID, and the state.
func parseGoStackHeader(line string) (goroutineID int, state string) {
	line = strings.TrimSuffix(line, ":\n")
	parts := strings.SplitN(line, " ", 3)
	if len(parts) != 3 {
		panic(fmt.Sprintf("unexpected stack header format: %v", line))
	}

	id, err := strconv.Atoi(parts[1])
	if err != nil {
		panic(fmt.Sprintf("failed to parse goroutine ID: %v", parts[1]))
	}

	state = strings.TrimSuffix(strings.TrimPrefix(parts[2], "["), "]")
	return id, state
}

// Stack represents a single Goroutine's stack.
type Stack struct {
	id        int
	state     string
	fullStack *bytes.Buffer
}

// ID returns the goroutine ID.
func (s Stack) ID() int {
	return s.id
}

// State returns the Goroutine's state.
func (s Stack) State() string {
	return s.state
}

// Full returns the full stack trace for this goroutine.
func (s Stack) Full() []byte {
	return s.fullStack.Bytes()
}

func (s Stack) String() string {
	return fmt.Sprintf("Goroutine %v in state %v:\n%s", s.id, s.state, s.Full())
}

// GetAll returns the stacks for all running goroutines.
func GetAll() []Stack {
	var stacks []Stack

	var curStack *Stack
	stackReader := bufio.NewReader(bytes.NewReader(getStackBuffer(true /* all */)))
	for {
		line, err := stackReader.ReadString('\n')
		if err == io.EOF {
			break
		}
		if err != nil {
			panic("stack reader failed")
		}

		// If we see the goroutine header, start a new stack.
		if strings.HasPrefix(line, "goroutine ") {
			// flush any previous stack
			if curStack != nil {
				stacks = append(stacks, *curStack)
			}
			id, goState := parseGoStackHeader(line)
			curStack = &Stack{
				id:        id,
				state:     goState,
				fullStack: &bytes.Buffer{},
			}
		}
		curStack.fullStack.WriteString(line)
	}

	if curStack != nil {
		stacks = append(stacks, *curStack)
	}
	return stacks
}
<commit_msg>Re-order stacks parsing code<commit_after>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package goroutines

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"runtime"
	"strconv"
	"strings"
)

// Stack represents a single Goroutine's stack.
type Stack struct {
	id        int
	state     string
	fullStack *bytes.Buffer
}

// ID returns the goroutine ID.
func (s Stack) ID() int {
	return s.id
}

// State returns the Goroutine's state.
func (s Stack) State() string {
	return s.state
}

// Full returns the full stack trace for this goroutine.
func (s Stack) Full() []byte {
	return s.fullStack.Bytes()
}

func (s Stack) String() string {
	return fmt.Sprintf("Goroutine %v in state %v:\n%s", s.id, s.state, s.Full())
}

// GetAll returns the stacks for all running goroutines.
func GetAll() []Stack {
	var stacks []Stack

	var curStack *Stack
	stackReader := bufio.NewReader(bytes.NewReader(getStackBuffer(true /* all */)))
	for {
		line, err := stackReader.ReadString('\n')
		if err == io.EOF {
			break
		}
		if err != nil {
			panic("stack reader failed")
		}

		// If we see the goroutine header, start a new stack.
		if strings.HasPrefix(line, "goroutine ") {
			// flush any previous stack
			if curStack != nil {
				stacks = append(stacks, *curStack)
			}
			id, goState := parseGoStackHeader(line)
			curStack = &Stack{
				id:        id,
				state:     goState,
				fullStack: &bytes.Buffer{},
			}
		}
		curStack.fullStack.WriteString(line)
	}

	if curStack != nil {
		stacks = append(stacks, *curStack)
	}
	return stacks
}

func getStackBuffer(all bool) []byte {
	for i := 4096; ; i *= 2 {
		buf := make([]byte, i)
		if n := runtime.Stack(buf, all); n < i {
			return buf
		}
	}
}

// parseGoStackHeader parses a stack header that looks like:
// goroutine 643 [runnable]:\n
// And returns the goroutine ID, and the state.
func parseGoStackHeader(line string) (goroutineID int, state string) {
	line = strings.TrimSuffix(line, ":\n")
	parts := strings.SplitN(line, " ", 3)
	if len(parts) != 3 {
		panic(fmt.Sprintf("unexpected stack header format: %v", line))
	}

	id, err := strconv.Atoi(parts[1])
	if err != nil {
		panic(fmt.Sprintf("failed to parse goroutine ID: %v", parts[1]))
	}

	state = strings.TrimSuffix(strings.TrimPrefix(parts[2], "["), "]")
	return id, state
}
<|endoftext|>


<commit_before>package processors

import (
	"bytes"
	"fmt"
	"io"
	"strings"

	"github.com/linuxerwang/goats-html/runtime"
	"github.com/linuxerwang/goats-html/util"
	"golang.org/x/net/html"
)

/**
 * TODO: Fix me. Logic is out of date.
 * The logic to handle tag processing:
 * if tag is the first:
 *     if caller has omitTag:
 *         if caller omitTag evals to false:
 *             call genTagWithDynamicAttrs()
 *         else:
 *             // Do nothing
 *     else: (caller has not omitTag)
 *         if tag has omitTag:
 *             if eval(omitTag) == false:
 *                 call genTagWithDynamicAttrs()
 *         else:
 *             call genTagWithDynamicAttrs()
 * else: (tag is subsequent)
 *     if omitTag == "":
 *         omitTag <- "false"
 *     if omitTag == "true":
 *         // Do nothing
 *     else if omitTag == "false":
 *         if has only static attrs:
 *             call genTagWithStaticAttrs()
 *         else:
 *             call genTagWithDynamicAttrs()
 *     else: (has omitTag expression)
 *         if eval(omitTag) == true:
 *             call genReturn()
 *         else:
 *             call genTagWithDynamicAttrs()
 * processChildrenTags()
 *
 */

type GoTagProcessor struct {
	BaseProcessor
	tagName         string
	omitTag         string
	hasOmitTag      bool
	firstTag        bool // Whether it's the first tag of the template.
	needsClosing    bool
	childProcessors []Processor
	attrs           []html.Attribute
}

func (t *GoTagProcessor) Process(writer io.Writer, ctx *TagContext) {
	originalAutoEscape := ctx.AutoEscape
	// There might be imports for constants or enums
	for _, attr := range t.attrs {
		if attr.Key == "go:autoescape" {
			ctx.AutoEscape = (attr.Val == "true")
			break
		}
	}

	// Start local scope.
	io.WriteString(writer, "{\n")

	switch ctx.OutputFormat {
	case "closure":
		io.WriteString(writer, "var __element = null;\n")
		io.WriteString(writer, "var __attrs = new goats.runtime.TagAttrs();\n")
	}

	if t.firstTag {
		t.processFirstTag(writer, ctx)
	} else {
		t.processSubseqTag(writer, ctx)
	}

	switch ctx.OutputFormat {
	case "closure":
		io.WriteString(writer, "if (__element) {\n")
		io.WriteString(writer, "goog.dom.appendChild(__tag_stack[__tag_stack.length-1], __element);\n")
		io.WriteString(writer, "__tag_stack.push(__element);\n")
		io.WriteString(writer, "}\n")
	}

	t.processChildrenTags(writer, ctx)

	switch ctx.OutputFormat {
	case "closure":
		io.WriteString(writer, "if (__element) {\n")
		io.WriteString(writer, "__tag_stack.pop();\n")
		io.WriteString(writer, "}\n")
	}

	t.maybeCloseTag(writer, ctx)

	// End local scope.
	io.WriteString(writer, "}\n")

	ctx.AutoEscape = originalAutoEscape
}

func (t *GoTagProcessor) processFirstTag(writer io.Writer, ctx *TagContext) {
	switch ctx.OutputFormat {
	case "go":
		// Vars for attributes
		io.WriteString(writer, "var __omitTag bool = false\n")
		io.WriteString(writer, "var __attrs = runtime.TagAttrs{}\n")
		// Local attributes
		t.genLocalAttrs(writer, ctx)
		// Caller attributes may overwrite local attrs.
		io.WriteString(writer, "if __callerAttrsFunc := __impl.GetCallerAttrsFunc(); __callerAttrsFunc != nil {\n")
		io.WriteString(writer, "  var __callerAttrs runtime.TagAttrs\n")
		io.WriteString(writer, "  __callerAttrs, __callerHasOmitTag, __callerOmitTag := __callerAttrsFunc()\n")
		io.WriteString(writer, "  __attrs.MergeFrom(__callerAttrs)\n")
		io.WriteString(writer, "  if __callerHasOmitTag {\n")
		io.WriteString(writer, "    if __callerOmitTag {\n")
		io.WriteString(writer, "      __omitTag = true\n")
		io.WriteString(writer, "    } else {\n")
		io.WriteString(writer, fmt.Sprintf("    __attrs.GenTagAndAttrs(__impl.GetWriter(), \"%s\")\n", t.tagName))
		io.WriteString(writer, "    }\n")
		io.WriteString(writer, "  } else {\n")
		if t.hasOmitTag {
			t.genConditionalTag(writer, ctx)
		} else {
			io.WriteString(writer, fmt.Sprintf("  __attrs.GenTagAndAttrs(__impl.GetWriter(), \"%s\")\n", t.tagName))
		}
		io.WriteString(writer, "  }\n")
		io.WriteString(writer, "} else {\n")
		if !t.hasOmitTag || t.omitTag == "false" {
			io.WriteString(writer, fmt.Sprintf("  __attrs.GenTagAndAttrs(__impl.GetWriter(), \"%s\")\n", t.tagName))
		} else if t.omitTag == "true" {
			io.WriteString(writer, "  __omitTag = true\n")
			// Do not output tag.
		} else {
			t.genConditionalTag(writer, ctx)
		}
		io.WriteString(writer, "}\n")
	case "closure":
		// Vars for attributes
		io.WriteString(writer, "var __omitTag = false;\n")
		// Local attributes
		t.genLocalAttrs(writer, ctx)
		// Caller attributes may overwrite local attrs.
		io.WriteString(writer, "if (this.callerAttrsFunc_) {\n")
		io.WriteString(writer, "  var __callerAttrs = this.callerAttrsFunc_();\n")
		io.WriteString(writer, "  __attrs.mergeFrom(__callerAttrs.attrs);\n")
		io.WriteString(writer, "  if (__callerAttrs.hasOmitTag) {\n")
		io.WriteString(writer, "    if (__callerAttrs.omitTag) {\n")
		io.WriteString(writer, "      __omitTag = true;\n")
		io.WriteString(writer, "    }\n")
		io.WriteString(writer, "  } else {\n")
		if t.hasOmitTag {
			io.WriteString(writer, fmt.Sprintf("__omitTag = %s;\n", t.omitTag))
		}
		io.WriteString(writer, "  }\n")
		io.WriteString(writer, "}\n")
		io.WriteString(writer, "if (!__omitTag) {\n")
		io.WriteString(writer, fmt.Sprintf("  __element = goog.dom.createDom(\"%s\", __attrs);\n", t.tagName))
		io.WriteString(writer, "}\n")
	}
}

func (t *GoTagProcessor) processSubseqTag(writer io.Writer, ctx *TagContext) {
	switch t.omitTag {
	case "true":
		// Do nothing
	case "", "false":
		if t.hasOnlyStaticAttrs() {
			attrs := runtime.TagAttrs{}
			for _, attr := range t.attrs {
				if !strings.HasPrefix(attr.Key, "go:") {
					attrs.AddAttr(attr.Key, attr.Val)
				}
			}
			switch ctx.OutputFormat {
			case "go":
				var tagBuffer bytes.Buffer
				attrs.GenTagAndAttrs(&tagBuffer, t.tagName)
				io.WriteString(writer,
					fmt.Sprintf("__impl.WriteString(\"%s\")\n",
						strings.Replace(tagBuffer.String(), "\"", "\\\"", -1)))
			case "closure":
				io.WriteString(writer, fmt.Sprintf("__element = goog.dom.createElement(\"%s\");\n", t.tagName))
				for k, v := range map[string]string(attrs) {
					io.WriteString(writer, fmt.Sprintf("__element.setAttribute(\"%s\", \"%s\");\n", k, v))
				}
			}
		} else {
			switch ctx.OutputFormat {
			case "go":
				io.WriteString(writer, "  var __attrs = runtime.TagAttrs{}\n")
				t.genLocalAttrs(writer, ctx)
				io.WriteString(writer, fmt.Sprintf("  __attrs.GenTagAndAttrs(__impl.GetWriter(), \"%s\")\n", t.tagName))
			case "closure":
				t.genLocalAttrs(writer, ctx)
				io.WriteString(writer, fmt.Sprintf("__element = goog.dom.createDom(\"%s\", __attrs);\n", t.tagName))
			}
		}
	default:
		switch ctx.OutputFormat {
		case "go":
			io.WriteString(writer, fmt.Sprintf("if !%s {\n", t.omitTag))
			io.WriteString(writer, "var __omitTag bool = false\n")
			io.WriteString(writer, "  var __attrs = runtime.TagAttrs{}\n")
			io.WriteString(writer, fmt.Sprintf("  __attrs.GenTagAndAttrs(__impl.GetWriter(), \"%s\")\n", t.tagName))
			t.genLocalAttrs(writer, ctx)
			t.genConditionalTag(writer, ctx)
			io.WriteString(writer, "} else {\n")
			io.WriteString(writer, "  __omitTag = true\n")
			// Do nothing
			io.WriteString(writer, "}\n")
		case "closure":
			io.WriteString(writer, fmt.Sprintf("if (!%s) {\n", t.omitTag))
			io.WriteString(writer, "var __omitTag = false;\n")
			io.WriteString(writer, "  var __attrs = {};\n")
			io.WriteString(writer, fmt.Sprintf("__element = goog.dom.createDom(\"%s\", __attrs);\n", t.tagName))
			t.genLocalAttrs(writer, ctx)
			t.genConditionalTag(writer, ctx)
			io.WriteString(writer, "} else {\n")
			io.WriteString(writer, "  __omitTag = true;\n")
			// Do nothing
			io.WriteString(writer, "}\n")
		}
	}
}

func (t *GoTagProcessor) genConditionalTag(writer io.Writer, ctx *TagContext) {
	switch ctx.OutputFormat {
	case "go":
		io.WriteString(writer, fmt.Sprintf("__omitTag = %s\n", t.omitTag))
		io.WriteString(writer, "if !__omitTag {\n")
		io.WriteString(writer, fmt.Sprintf("  __attrs.GenTagAndAttrs(__impl.GetWriter(), \"%s\")\n", t.tagName))
		io.WriteString(writer, "}\n")
	case "closure":
		io.WriteString(writer, fmt.Sprintf("__omitTag = %s;\n", t.omitTag))
		io.WriteString(writer, "if (!__omitTag) {\n")
		io.WriteString(writer, fmt.Sprintf("  __element = goog.dom.createDom(\"%s\", __attrs);\n", t.tagName))
		io.WriteString(writer, "}\n")
	}
}

func (t *GoTagProcessor) genLocalAttrs(writer io.Writer, ctx *TagContext) {
	for _, attr := range t.attrs {
		if attr.Key == "go:attr" {
			varName, varVal := util.SplitVarDef(attr.Val)
			ctx.ExprParser.Evaluate(varVal, writer, func(expr string) {
				switch ctx.OutputFormat {
				case "go":
					io.WriteString(writer, fmt.Sprintf("__attrs.AddAttr(\"%s\", %s)\n", varName, expr))
				case "closure":
					io.WriteString(writer, fmt.Sprintf("__attrs.add(\"%s\", %s);\n", varName, expr))
				}
			})
		} else if !strings.HasPrefix(attr.Key, "go:") {
			// Static attrs
			switch ctx.OutputFormat {
			case "go":
				io.WriteString(writer,
					fmt.Sprintf("__attrs.AddAttr(\"%s\", \"%s\")\n",
						util.TrimWhiteSpaces(attr.Key), util.TrimWhiteSpaces(attr.Val)))
			case "closure":
				io.WriteString(writer, fmt.Sprintf("__attrs.add(\"%s\", \"%s\");\n", util.TrimWhiteSpaces(attr.Key), util.TrimWhiteSpaces(attr.Val)))
			}
		}
	}
}

func (t *GoTagProcessor) hasOnlyStaticAttrs() bool {
	for _, attr := range t.attrs {
		if attr.Key == "go:attr" {
			return false
		}
	}
	return true
}

func (t *GoTagProcessor) processChildrenTags(writer io.Writer, ctx *TagContext) {
	originalAutoEscape := ctx.AutoEscape
	if t.tagName == "script" || t.tagName == "style" {
		ctx.AutoEscape = false
	}

	if t.next != nil {
		t.next.Process(writer, ctx)
	}

	for _, child := range t.childProcessors {
		child.Process(writer, ctx)
	}

	ctx.AutoEscape = originalAutoEscape
}

func (t *GoTagProcessor) maybeCloseTag(writer io.Writer, ctx *TagContext) {
	switch ctx.OutputFormat {
	case "go":
		if t.needsClosing {
			if t.firstTag || t.isDynamicOmitTag() {
				io.WriteString(writer, "if (!__omitTag) {\n")
				t.closeTag(writer)
				io.WriteString(writer, "}\n")
			} else if t.omitTag == "" || t.omitTag == "false" {
				t.closeTag(writer)
			}
		}
	}
}

func (t *GoTagProcessor) closeTag(writer io.Writer) {
	io.WriteString(writer, fmt.Sprintf("__impl.WriteString(\"</%s>\")\n", t.tagName))
}

func (t *GoTagProcessor) isDynamicOmitTag() bool {
	return t.omitTag != "" && t.omitTag != "true" && t.omitTag != "false"
}

func (t *GoTagProcessor) getTail(processor Processor) Processor {
	tail := processor
	for tail.HasNext() {
		tail = tail.GetNext()
	}
	return tail
}

func (t *GoTagProcessor) AddChild(child Processor) {
	t.childProcessors = append(t.childProcessors, child)
}

func NewTagProcessor(tagName string, omitTag string, firstTag bool, needsClosing bool,
	attrs []html.Attribute) *GoTagProcessor {
	processor := &GoTagProcessor{
		tagName:      tagName,
		omitTag:      omitTag,
		hasOmitTag:   omitTag != "",
		firstTag:     firstTag,
		needsClosing: needsClosing,
		attrs:        attrs,
	}
	return processor
}
<commit_msg>Fixes for javascript generation: 1) correctly scope javascript variables for go:tag; 2) fix attrs.<commit_after>package processors

import (
	"bytes"
	"fmt"
	"io"
	"strings"

	"github.com/linuxerwang/goats-html/runtime"
	"github.com/linuxerwang/goats-html/util"
	"golang.org/x/net/html"
)

/**
 * TODO: Fix me. Logic is out of date.
 * The logic to handle tag processing:
 * if tag is the first:
 *     if caller has omitTag:
 *         if caller omitTag evals to false:
 *             call genTagWithDynamicAttrs()
 *         else:
 *             // Do nothing
 *     else: (caller has not omitTag)
 *         if tag has omitTag:
 *             if eval(omitTag) == false:
 *                 call genTagWithDynamicAttrs()
 *         else:
 *             call genTagWithDynamicAttrs()
 * else: (tag is subsequent)
 *     if omitTag == "":
 *         omitTag <- "false"
 *     if omitTag == "true":
 *         // Do nothing
 *     else if omitTag == "false":
 *         if has only static attrs:
 *             call genTagWithStaticAttrs()
 *         else:
 *             call genTagWithDynamicAttrs()
 *     else: (has omitTag expression)
 *         if eval(omitTag) == true:
 *             call genReturn()
 *         else:
 *             call genTagWithDynamicAttrs()
 * processChildrenTags()
 *
 */

type GoTagProcessor struct {
	BaseProcessor
	tagName         string
	omitTag         string
	hasOmitTag      bool
	firstTag        bool // Whether it's the first tag of the template.
	needsClosing    bool
	childProcessors []Processor
	attrs           []html.Attribute
}

func (t *GoTagProcessor) Process(writer io.Writer, ctx *TagContext) {
	originalAutoEscape := ctx.AutoEscape
	// There might be imports for constants or enums
	for _, attr := range t.attrs {
		if attr.Key == "go:autoescape" {
			ctx.AutoEscape = (attr.Val == "true")
			break
		}
	}

	// Start local scope.
	switch ctx.OutputFormat {
	case "go":
		io.WriteString(writer, "{\n")
	case "closure":
		io.WriteString(writer, "(function(){\n")
	}

	switch ctx.OutputFormat {
	case "closure":
		io.WriteString(writer, "var __element = null;\n")
		io.WriteString(writer, "var __attrs = new goats.runtime.TagAttrs();\n")
	}

	if t.firstTag {
		t.processFirstTag(writer, ctx)
	} else {
		t.processSubseqTag(writer, ctx)
	}

	switch ctx.OutputFormat {
	case "closure":
		io.WriteString(writer, "if (__element) {\n")
		io.WriteString(writer, "goog.dom.appendChild(__tag_stack[__tag_stack.length-1], __element);\n")
		io.WriteString(writer, "__tag_stack.push(__element);\n")
		io.WriteString(writer, "}\n")
	}

	t.processChildrenTags(writer, ctx)

	switch ctx.OutputFormat {
	case "closure":
		io.WriteString(writer, "if (__element) {\n")
		io.WriteString(writer, "__tag_stack.pop();\n")
		io.WriteString(writer, "}\n")
	}

	t.maybeCloseTag(writer, ctx)

	// End local scope.
	switch ctx.OutputFormat {
	case "go":
		io.WriteString(writer, "}\n")
	case "closure":
		io.WriteString(writer, "})();\n")
	}

	ctx.AutoEscape = originalAutoEscape
}

func (t *GoTagProcessor) processFirstTag(writer io.Writer, ctx *TagContext) {
	switch ctx.OutputFormat {
	case "go":
		// Vars for attributes
		io.WriteString(writer, "var __omitTag bool = false\n")
		io.WriteString(writer, "var __attrs = runtime.TagAttrs{}\n")
		// Local attributes
		t.genLocalAttrs(writer, ctx)
		// Caller attributes may overwrite local attrs.
		io.WriteString(writer, "if __callerAttrsFunc := __impl.GetCallerAttrsFunc(); __callerAttrsFunc != nil {\n")
		io.WriteString(writer, "  var __callerAttrs runtime.TagAttrs\n")
		io.WriteString(writer, "  __callerAttrs, __callerHasOmitTag, __callerOmitTag := __callerAttrsFunc()\n")
		io.WriteString(writer, "  __attrs.MergeFrom(__callerAttrs)\n")
		io.WriteString(writer, "  if __callerHasOmitTag {\n")
		io.WriteString(writer, "    if __callerOmitTag {\n")
		io.WriteString(writer, "      __omitTag = true\n")
		io.WriteString(writer, "    } else {\n")
		io.WriteString(writer, fmt.Sprintf("    __attrs.GenTagAndAttrs(__impl.GetWriter(), \"%s\")\n", t.tagName))
		io.WriteString(writer, "    }\n")
		io.WriteString(writer, "  } else {\n")
		if t.hasOmitTag {
			t.genConditionalTag(writer, ctx)
		} else {
			io.WriteString(writer, fmt.Sprintf("  __attrs.GenTagAndAttrs(__impl.GetWriter(), \"%s\")\n", t.tagName))
		}
		io.WriteString(writer, "  }\n")
		io.WriteString(writer, "} else {\n")
		if !t.hasOmitTag || t.omitTag == "false" {
			io.WriteString(writer, fmt.Sprintf("  __attrs.GenTagAndAttrs(__impl.GetWriter(), \"%s\")\n", t.tagName))
		} else if t.omitTag == "true" {
			io.WriteString(writer, "  __omitTag = true\n")
			// Do not output tag.
		} else {
			t.genConditionalTag(writer, ctx)
		}
		io.WriteString(writer, "}\n")
	case "closure":
		// Vars for attributes
		io.WriteString(writer, "var __omitTag = false;\n")
		// Local attributes
		t.genLocalAttrs(writer, ctx)
		// Caller attributes may overwrite local attrs.
		io.WriteString(writer, "if (this.callerAttrsFunc_) {\n")
		io.WriteString(writer, "  var __callerAttrs = this.callerAttrsFunc_();\n")
		io.WriteString(writer, "  __attrs.mergeFrom(__callerAttrs.attrs);\n")
		io.WriteString(writer, "  if (__callerAttrs.hasOmitTag) {\n")
		io.WriteString(writer, "    if (__callerAttrs.omitTag) {\n")
		io.WriteString(writer, "      __omitTag = true;\n")
		io.WriteString(writer, "    }\n")
		io.WriteString(writer, "  } else {\n")
		if t.hasOmitTag {
			io.WriteString(writer, fmt.Sprintf("__omitTag = %s;\n", t.omitTag))
		}
		io.WriteString(writer, "  }\n")
		io.WriteString(writer, "}\n")
		io.WriteString(writer, "if (!__omitTag) {\n")
		io.WriteString(writer, fmt.Sprintf("  __element = goog.dom.createDom(\"%s\", __attrs.get());\n", t.tagName))
		io.WriteString(writer, "}\n")
	}
}

func (t *GoTagProcessor) processSubseqTag(writer io.Writer, ctx *TagContext) {
	switch t.omitTag {
	case "true":
		// Do nothing
	case "", "false":
		if t.hasOnlyStaticAttrs() {
			attrs := runtime.TagAttrs{}
			for _, attr := range t.attrs {
				if !strings.HasPrefix(attr.Key, "go:") {
					attrs.AddAttr(attr.Key, attr.Val)
				}
			}
			switch ctx.OutputFormat {
			case "go":
				var tagBuffer bytes.Buffer
				attrs.GenTagAndAttrs(&tagBuffer, t.tagName)
				io.WriteString(writer,
					fmt.Sprintf("__impl.WriteString(\"%s\")\n",
						strings.Replace(tagBuffer.String(), "\"", "\\\"", -1)))
			case "closure":
				io.WriteString(writer, fmt.Sprintf("__element = goog.dom.createElement(\"%s\");\n", t.tagName))
				for k, v := range map[string]string(attrs) {
					io.WriteString(writer, fmt.Sprintf("__element.setAttribute(\"%s\", \"%s\");\n", k, v))
				}
			}
		} else {
			switch ctx.OutputFormat {
			case "go":
				io.WriteString(writer, "  var __attrs = runtime.TagAttrs{}\n")
				t.genLocalAttrs(writer, ctx)
				io.WriteString(writer, fmt.Sprintf("  __attrs.GenTagAndAttrs(__impl.GetWriter(), \"%s\")\n", t.tagName))
			case "closure":
				t.genLocalAttrs(writer, ctx)
				io.WriteString(writer, fmt.Sprintf("__element = goog.dom.createDom(\"%s\", __attrs.get());\n", t.tagName))
			}
		}
	default:
		switch ctx.OutputFormat {
		case "go":
			io.WriteString(writer, fmt.Sprintf("if !%s {\n", t.omitTag))
			io.WriteString(writer, "var __omitTag bool = false\n")
			io.WriteString(writer, "  var __attrs = runtime.TagAttrs{}\n")
			io.WriteString(writer, fmt.Sprintf("  __attrs.GenTagAndAttrs(__impl.GetWriter(), \"%s\")\n", t.tagName))
			t.genLocalAttrs(writer, ctx)
			t.genConditionalTag(writer, ctx)
			io.WriteString(writer, "} else {\n")
			io.WriteString(writer, "  __omitTag = true\n")
			// Do nothing
			io.WriteString(writer, "}\n")
		case "closure":
			io.WriteString(writer, fmt.Sprintf("if (!%s) {\n", t.omitTag))
			io.WriteString(writer, "var __omitTag = false;\n")
			io.WriteString(writer, "  var __attrs = {};\n")
			io.WriteString(writer, fmt.Sprintf("__element = goog.dom.createDom(\"%s\", __attrs.get());\n", t.tagName))
			t.genLocalAttrs(writer, ctx)
			t.genConditionalTag(writer, ctx)
			io.WriteString(writer, "} else {\n")
			io.WriteString(writer, "  __omitTag = true;\n")
			// Do nothing
			io.WriteString(writer, "}\n")
		}
	}
}

func (t *GoTagProcessor) genConditionalTag(writer io.Writer, ctx *TagContext) {
	switch ctx.OutputFormat {
	case "go":
		io.WriteString(writer, fmt.Sprintf("__omitTag = %s\n", t.omitTag))
		io.WriteString(writer, "if !__omitTag {\n")
		io.WriteString(writer, fmt.Sprintf("  __attrs.GenTagAndAttrs(__impl.GetWriter(), \"%s\")\n", t.tagName))
		io.WriteString(writer, "}\n")
	case "closure":
		io.WriteString(writer, fmt.Sprintf("__omitTag = %s;\n", t.omitTag))
		io.WriteString(writer, "if (!__omitTag) {\n")
		io.WriteString(writer, fmt.Sprintf("  __element = goog.dom.createDom(\"%s\", __attrs.get());\n", t.tagName))
		io.WriteString(writer, "}\n")
	}
}

func (t *GoTagProcessor) genLocalAttrs(writer io.Writer, ctx *TagContext) {
	for _, attr := range t.attrs {
		if attr.Key == "go:attr" {
			varName, varVal := util.SplitVarDef(attr.Val)
			ctx.ExprParser.Evaluate(varVal, writer, func(expr string) {
				switch ctx.OutputFormat {
				case "go":
					io.WriteString(writer, fmt.Sprintf("__attrs.AddAttr(\"%s\", %s)\n", varName, expr))
				case "closure":
					io.WriteString(writer, fmt.Sprintf("__attrs.add(\"%s\", %s);\n", varName, expr))
				}
			})
		} else if !strings.HasPrefix(attr.Key, "go:") {
			// Static attrs
			switch ctx.OutputFormat {
			case "go":
				io.WriteString(writer,
					fmt.Sprintf("__attrs.AddAttr(\"%s\", \"%s\")\n",
						util.TrimWhiteSpaces(attr.Key), util.TrimWhiteSpaces(attr.Val)))
			case "closure":
				io.WriteString(writer, fmt.Sprintf("__attrs.add(\"%s\", \"%s\");\n", util.TrimWhiteSpaces(attr.Key), util.TrimWhiteSpaces(attr.Val)))
			}
		}
	}
}

func (t *GoTagProcessor) hasOnlyStaticAttrs() bool {
	for _, attr := range t.attrs {
		if attr.Key == "go:attr" {
			return false
		}
	}
	return true
}

func (t *GoTagProcessor) processChildrenTags(writer io.Writer, ctx *TagContext) {
	originalAutoEscape := ctx.AutoEscape
	if t.tagName == "script" || t.tagName == "style" {
		ctx.AutoEscape = false
	}

	if t.next != nil {
		t.next.Process(writer, ctx)
	}

	for _, child := range t.childProcessors {
		child.Process(writer, ctx)
	}

	ctx.AutoEscape = originalAutoEscape
}

func (t *GoTagProcessor) maybeCloseTag(writer io.Writer, ctx *TagContext) {
	switch ctx.OutputFormat {
	case "go":
		if t.needsClosing {
			if t.firstTag || t.isDynamicOmitTag() {
				io.WriteString(writer, "if (!__omitTag) {\n")
				t.closeTag(writer)
				io.WriteString(writer, "}\n")
			} else if t.omitTag == "" || t.omitTag == "false" {
				t.closeTag(writer)
			}
		}
	}
}

func (t *GoTagProcessor) closeTag(writer io.Writer) {
	io.WriteString(writer, fmt.Sprintf("__impl.WriteString(\"</%s>\")\n", t.tagName))
}

func (t *GoTagProcessor) isDynamicOmitTag() bool {
	return t.omitTag != "" && t.omitTag != "true" && t.omitTag != "false"
}

func (t *GoTagProcessor) getTail(processor Processor) Processor {
	tail := processor
	for tail.HasNext() {
		tail = tail.GetNext()
	}
	return tail
}

func (t *GoTagProcessor) AddChild(child Processor) {
	t.childProcessors = append(t.childProcessors, child)
}

func NewTagProcessor(tagName string, omitTag string, firstTag bool, needsClosing bool,
	attrs []html.Attribute) *GoTagProcessor {
	processor := &GoTagProcessor{
		tagName:      tagName,
		omitTag:      omitTag,
		hasOmitTag:   omitTag != "",
		firstTag:     firstTag,
		needsClosing: needsClosing,
		attrs:        attrs,
	}
	return processor
}
<|endoftext|>