<commit_before>package sharing

import (
	"encoding/json"
	"fmt"
	"runtime"
	"strings"

	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/couchdb/mango"
	"github.com/cozy/cozy-stack/pkg/instance"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/lock"
	"github.com/cozy/cozy-stack/pkg/vfs"
)

// SetupReceiver is used on the receivers' cozy to make sure the cozy can
// receive the shared documents.
func (s *Sharing) SetupReceiver(inst *instance.Instance) error {
	inst.Logger().WithField("nspace", "sharing").
		Debugf("Setup receiver on %#v", inst)

	if err := couchdb.EnsureDBExist(inst, consts.Shared); err != nil {
		return err
	}
	if err := s.AddTrackTriggers(inst); err != nil {
		return err
	}
	if !s.ReadOnly() {
		if err := s.AddReplicateTrigger(inst); err != nil {
			return err
		}
		if s.FirstFilesRule() != nil {
			if err := s.AddUploadTrigger(inst); err != nil {
				return err
			}
		}
	}
	return nil
}

// Setup is used when a member accepts a sharing to prepare the io.cozy.shared
// database and start an initial replication. It is meant to be used in a new
// goroutine and, as such, does not return errors but log them.
func (s *Sharing) Setup(inst *instance.Instance, m *Member) {
	// Don't do the setup for most tests
	if !inst.OnboardingFinished {
		return
	}
	inst.Logger().WithField("nspace", "sharing").
		Debugf("Setup for member %#v on %#v", m, inst)

	defer func() {
		if r := recover(); r != nil {
			var err error
			switch r := r.(type) {
			case error:
				err = r
			default:
				err = fmt.Errorf("%v", r)
			}
			stack := make([]byte, 4<<10) // 4 KB
			length := runtime.Stack(stack, false)
			log := inst.Logger().WithField("panic", true).WithField("nspace", "sharing")
			log.Errorf("PANIC RECOVER %s: %s", err.Error(), stack[:length])
		}
	}()

	mu := lock.ReadWrite(inst, "sharings/"+s.SID)
	mu.Lock()
	defer mu.Unlock()

	if err := couchdb.EnsureDBExist(inst, consts.Shared); err != nil {
		inst.Logger().WithField("nspace", "sharing").
			Warnf("Can't ensure io.cozy.shared exists (%s): %s", s.SID, err)
	}
	if rule := s.FirstFilesRule(); rule != nil && rule.Selector != couchdb.SelectorReferencedBy {
		if err := s.AddReferenceForSharingDir(inst, rule); err != nil {
			inst.Logger().WithField("nspace", "sharing").
				Warnf("Error on referenced_by for the sharing dir (%s): %s", s.SID, err)
		}
	}

	if err := s.AddTrackTriggers(inst); err != nil {
		inst.Logger().WithField("nspace", "sharing").
			Warnf("Error on setup of track triggers (%s): %s", s.SID, err)
	}
	for i, rule := range s.Rules {
		if err := s.InitialCopy(inst, rule, i); err != nil {
			inst.Logger().Warnf("Error on initial copy for %s (%s): %s", rule.Title, s.SID, err)
		}
	}
	if err := s.AddReplicateTrigger(inst); err != nil {
		inst.Logger().WithField("nspace", "sharing").
			Warnf("Error on setup replicate trigger (%s): %s", s.SID, err)
	}
	if pending, err := s.ReplicateTo(inst, m, true); err != nil {
		inst.Logger().WithField("nspace", "sharing").
			Warnf("Error on initial replication (%s): %s", s.SID, err)
		s.retryWorker(inst, "share-replicate", 0)
	} else {
		if pending {
			s.pushJob(inst, "share-replicate")
		}
		if s.FirstFilesRule() == nil {
			return
		}
		if err := s.AddUploadTrigger(inst); err != nil {
			inst.Logger().WithField("nspace", "sharing").
				Warnf("Error on setup upload trigger (%s): %s", s.SID, err)
		}
		if err := s.InitialUpload(inst, m); err != nil {
			inst.Logger().WithField("nspace", "sharing").
				Warnf("Error on initial upload (%s): %s", s.SID, err)
			s.retryWorker(inst, "share-upload", 0)
		}
	}

	go s.NotifyRecipients(inst, m)
}

// AddTrackTriggers creates the share-track triggers for each rule of the
// sharing that will update the io.cozy.shared database.
func (s *Sharing) AddTrackTriggers(inst *instance.Instance) error {
	if s.Triggers.TrackID != "" {
		return nil
	}
	sched := jobs.System()
	for i, rule := range s.Rules {
		args := rule.TriggerArgs()
		if args == "" {
			continue
		}
		msg := &TrackMessage{
			SharingID: s.SID,
			RuleIndex: i,
			DocType:   rule.DocType,
		}
		t, err := jobs.NewTrigger(inst, jobs.TriggerInfos{
			Type:       "@event",
			WorkerType: "share-track",
			Arguments:  args,
		}, msg)
		if err != nil {
			return err
		}
		if err = sched.AddTrigger(t); err != nil {
			return err
		}
		s.Triggers.TrackID = t.ID()
	}
	return couchdb.UpdateDoc(inst, s)
}

// AddReplicateTrigger creates the share-replicate trigger for this sharing:
// it will starts the replicator when some changes are made to the
// io.cozy.shared database.
func (s *Sharing) AddReplicateTrigger(inst *instance.Instance) error {
	if s.Triggers.ReplicateID != "" {
		return nil
	}
	msg := &ReplicateMsg{
		SharingID: s.SID,
		Errors:    0,
	}
	args := consts.Shared + ":CREATED,UPDATED:" + s.SID + ":sharing"
	t, err := jobs.NewTrigger(inst, jobs.TriggerInfos{
		Domain:     inst.ContextualDomain(),
		Type:       "@event",
		WorkerType: "share-replicate",
		Arguments:  args,
		Debounce:   "5s",
	}, msg)
	inst.Logger().WithField("nspace", "sharing").Infof("Create trigger %#v", t)
	if err != nil {
		return err
	}
	sched := jobs.System()
	if err = sched.AddTrigger(t); err != nil {
		return err
	}
	s.Triggers.ReplicateID = t.ID()
	return couchdb.UpdateDoc(inst, s)
}

// InitialCopy lists the shared documents and put a reference in the
// io.cozy.shared database
func (s *Sharing) InitialCopy(inst *instance.Instance, rule Rule, r int) error {
	if rule.Local || len(rule.Values) == 0 {
		return nil
	}

	mu := lock.ReadWrite(inst, "shared")
	mu.Lock()
	defer mu.Unlock()

	docs, err := findDocsToCopy(inst, rule)
	if err != nil {
		return err
	}
	refs, err := s.buildReferences(inst, rule, r, docs)
	if err != nil {
		return err
	}
	refs = compactSlice(refs)
	if len(refs) == 0 {
		return nil
	}
	olds := make([]interface{}, len(refs))
	return couchdb.BulkUpdateDocs(inst, consts.Shared, refs, olds)
}

// findDocsToCopy finds the documents that match the given rule
func findDocsToCopy(inst *instance.Instance, rule Rule) ([]couchdb.JSONDoc, error) {
	var docs []couchdb.JSONDoc
	if rule.Selector == "" || rule.Selector == "id" {
		if rule.DocType == consts.Files {
			for _, fileID := range rule.Values {
				err := vfs.WalkByID(inst.VFS(), fileID, func(name string, dir *vfs.DirDoc, file *vfs.FileDoc, err error) error {
					if err != nil {
						return err
					}
					if dir != nil {
						if dir.DocID != fileID {
							docs = append(docs, dirToJSONDoc(dir))
						}
					} else if file != nil {
						docs = append(docs, fileToJSONDoc(file))
					}
					return nil
				})
				if err != nil {
					return nil, err
				}
			}
		} else {
			req := &couchdb.AllDocsRequest{
				Keys: rule.Values,
			}
			if err := couchdb.GetAllDocs(inst, rule.DocType, req, &docs); err != nil {
				return nil, err
			}
		}
	} else {
		if rule.Selector == couchdb.SelectorReferencedBy {
			for _, val := range rule.Values {
				req := &couchdb.ViewRequest{
					Key:         strings.SplitN(val, "/", 2),
					IncludeDocs: true,
					Reduce:      false,
				}
				var res couchdb.ViewResponse
				err := couchdb.ExecView(inst, consts.FilesReferencedByView, req, &res)
				if err != nil {
					return nil, err
				}
				for _, row := range res.Rows {
					var doc couchdb.JSONDoc
					if err = json.Unmarshal(row.Doc, &doc); err == nil {
						docs = append(docs, doc)
					}
				}
			}
		} else {
			// Create index based on selector to retrieve documents to share
			name := "by-" + rule.Selector
			idx := mango.IndexOnFields(rule.DocType, name, []string{rule.Selector})
			if err := couchdb.DefineIndex(inst, idx); err != nil {
				return nil, err
			}
			// Request the index for all values
			for _, val := range rule.Values {
				var results []couchdb.JSONDoc
				req := &couchdb.FindRequest{
					UseIndex: name,
					Selector: mango.Equal(rule.Selector, val),
				}
				if err := couchdb.FindDocs(inst, rule.DocType, req, &results); err != nil {
					return nil, err
				}
				docs = append(docs, results...)
			}
		}
	}
	return docs, nil
}

// buildReferences build the SharedRef to add/update the given docs in the
// io.cozy.shared database
func (s *Sharing) buildReferences(inst *instance.Instance, rule Rule, r int, docs []couchdb.JSONDoc) ([]interface{}, error) {
	ids := make([]string, len(docs))
	for i, doc := range docs {
		ids[i] = rule.DocType + "/" + doc.ID()
	}
	srefs, err := FindReferences(inst, ids)
	if err != nil {
		return nil, err
	}

	refs := make([]interface{}, len(docs))
	for i, doc := range docs {
		rev := doc.Rev()
		info := SharedInfo{
			Rule:   r,
			Binary: rule.DocType == consts.Files && doc.Get("type") == consts.FileType,
		}
		if srefs[i] == nil {
			refs[i] = SharedRef{
				SID:       rule.DocType + "/" + doc.ID(),
				Revisions: &RevsTree{Rev: rev},
				Infos:     map[string]SharedInfo{s.SID: info},
			}
		} else {
			found := srefs[i].Revisions.Find(rev) != nil
			if found {
				if _, ok := srefs[i].Infos[s.SID]; ok {
					continue
				}
			} else {
				srefs[i].Revisions.Add(rev)
			}
			srefs[i].Infos[s.SID] = info
			refs[i] = *srefs[i]
		}
	}

	return refs, nil
}

// AddUploadTrigger creates the share-upload trigger for this sharing:
// it will starts the synchronization of the binaries when a file is added or
// updated in the io.cozy.shared database.
func (s *Sharing) AddUploadTrigger(inst *instance.Instance) error {
	if s.Triggers.UploadID != "" {
		return nil
	}
	msg := &UploadMsg{
		SharingID: s.SID,
		Errors:    0,
	}
	args := consts.Shared + ":CREATED,UPDATED:" + s.SID + ":sharing"
	t, err := jobs.NewTrigger(inst, jobs.TriggerInfos{
		Domain:     inst.ContextualDomain(),
		Type:       "@event",
		WorkerType: "share-upload",
		Arguments:  args,
		Debounce:   "5s",
	}, msg)
	inst.Logger().WithField("nspace", "sharing").Infof("Create trigger %#v", t)
	if err != nil {
		return err
	}
	sched := jobs.System()
	if err = sched.AddTrigger(t); err != nil {
		return err
	}
	s.Triggers.UploadID = t.ID()
	return couchdb.UpdateDoc(inst, s)
}

// compactSlice returns the given slice without the nil values
// https://github.com/golang/go/wiki/SliceTricks#filtering-without-allocating
func compactSlice(a []interface{}) []interface{} {
	b := a[:0]
	for _, x := range a {
		if x != nil {
			b = append(b, x)
		}
	}
	return b
}
<commit_msg>Fix sharing an empty directory<commit_after>package sharing

import (
	"encoding/json"
	"fmt"
	"runtime"
	"strings"

	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/couchdb/mango"
	"github.com/cozy/cozy-stack/pkg/instance"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/lock"
	"github.com/cozy/cozy-stack/pkg/vfs"
)

// SetupReceiver is used on the receivers' cozy to make sure the cozy can
// receive the shared documents.
func (s *Sharing) SetupReceiver(inst *instance.Instance) error {
	inst.Logger().WithField("nspace", "sharing").
		Debugf("Setup receiver on %#v", inst)

	if err := couchdb.EnsureDBExist(inst, consts.Shared); err != nil {
		return err
	}
	if err := s.AddTrackTriggers(inst); err != nil {
		return err
	}
	if !s.ReadOnly() {
		if err := s.AddReplicateTrigger(inst); err != nil {
			return err
		}
		if s.FirstFilesRule() != nil {
			if err := s.AddUploadTrigger(inst); err != nil {
				return err
			}
			// The sharing directory is created when the stack receives the
			// first file (it allows to not create it if it's just a file that
			// is shared, not a directory or an album). But, for an empty
			// directory, no files are sent, so we need to create it now.
			if s.NbFiles == 0 {
				if _, err := s.GetSharingDir(inst); err != nil {
					return err
				}
			}
		}
	}
	return nil
}

// Setup is used when a member accepts a sharing to prepare the io.cozy.shared
// database and start an initial replication. It is meant to be used in a new
// goroutine and, as such, does not return errors but log them.
func (s *Sharing) Setup(inst *instance.Instance, m *Member) {
	// Don't do the setup for most tests
	if !inst.OnboardingFinished {
		return
	}
	inst.Logger().WithField("nspace", "sharing").
		Debugf("Setup for member %#v on %#v", m, inst)

	defer func() {
		if r := recover(); r != nil {
			var err error
			switch r := r.(type) {
			case error:
				err = r
			default:
				err = fmt.Errorf("%v", r)
			}
			stack := make([]byte, 4<<10) // 4 KB
			length := runtime.Stack(stack, false)
			log := inst.Logger().WithField("panic", true).WithField("nspace", "sharing")
			log.Errorf("PANIC RECOVER %s: %s", err.Error(), stack[:length])
		}
	}()

	mu := lock.ReadWrite(inst, "sharings/"+s.SID)
	mu.Lock()
	defer mu.Unlock()

	if err := couchdb.EnsureDBExist(inst, consts.Shared); err != nil {
		inst.Logger().WithField("nspace", "sharing").
			Warnf("Can't ensure io.cozy.shared exists (%s): %s", s.SID, err)
	}
	if rule := s.FirstFilesRule(); rule != nil && rule.Selector != couchdb.SelectorReferencedBy {
		if err := s.AddReferenceForSharingDir(inst, rule); err != nil {
			inst.Logger().WithField("nspace", "sharing").
				Warnf("Error on referenced_by for the sharing dir (%s): %s", s.SID, err)
		}
	}

	if err := s.AddTrackTriggers(inst); err != nil {
		inst.Logger().WithField("nspace", "sharing").
			Warnf("Error on setup of track triggers (%s): %s", s.SID, err)
	}
	for i, rule := range s.Rules {
		if err := s.InitialCopy(inst, rule, i); err != nil {
			inst.Logger().Warnf("Error on initial copy for %s (%s): %s", rule.Title, s.SID, err)
		}
	}
	if err := s.AddReplicateTrigger(inst); err != nil {
		inst.Logger().WithField("nspace", "sharing").
			Warnf("Error on setup replicate trigger (%s): %s", s.SID, err)
	}
	if pending, err := s.ReplicateTo(inst, m, true); err != nil {
		inst.Logger().WithField("nspace", "sharing").
			Warnf("Error on initial replication (%s): %s", s.SID, err)
		s.retryWorker(inst, "share-replicate", 0)
	} else {
		if pending {
			s.pushJob(inst, "share-replicate")
		}
		if s.FirstFilesRule() == nil {
			return
		}
		if err := s.AddUploadTrigger(inst); err != nil {
			inst.Logger().WithField("nspace", "sharing").
				Warnf("Error on setup upload trigger (%s): %s", s.SID, err)
		}
		if err := s.InitialUpload(inst, m); err != nil {
			inst.Logger().WithField("nspace", "sharing").
				Warnf("Error on initial upload (%s): %s", s.SID, err)
			s.retryWorker(inst, "share-upload", 0)
		}
	}

	go s.NotifyRecipients(inst, m)
}

// AddTrackTriggers creates the share-track triggers for each rule of the
// sharing that will update the io.cozy.shared database.
func (s *Sharing) AddTrackTriggers(inst *instance.Instance) error {
	if s.Triggers.TrackID != "" {
		return nil
	}
	sched := jobs.System()
	for i, rule := range s.Rules {
		args := rule.TriggerArgs()
		if args == "" {
			continue
		}
		msg := &TrackMessage{
			SharingID: s.SID,
			RuleIndex: i,
			DocType:   rule.DocType,
		}
		t, err := jobs.NewTrigger(inst, jobs.TriggerInfos{
			Type:       "@event",
			WorkerType: "share-track",
			Arguments:  args,
		}, msg)
		if err != nil {
			return err
		}
		if err = sched.AddTrigger(t); err != nil {
			return err
		}
		s.Triggers.TrackID = t.ID()
	}
	return couchdb.UpdateDoc(inst, s)
}

// AddReplicateTrigger creates the share-replicate trigger for this sharing:
// it will starts the replicator when some changes are made to the
// io.cozy.shared database.
func (s *Sharing) AddReplicateTrigger(inst *instance.Instance) error {
	if s.Triggers.ReplicateID != "" {
		return nil
	}
	msg := &ReplicateMsg{
		SharingID: s.SID,
		Errors:    0,
	}
	args := consts.Shared + ":CREATED,UPDATED:" + s.SID + ":sharing"
	t, err := jobs.NewTrigger(inst, jobs.TriggerInfos{
		Domain:     inst.ContextualDomain(),
		Type:       "@event",
		WorkerType: "share-replicate",
		Arguments:  args,
		Debounce:   "5s",
	}, msg)
	inst.Logger().WithField("nspace", "sharing").Infof("Create trigger %#v", t)
	if err != nil {
		return err
	}
	sched := jobs.System()
	if err = sched.AddTrigger(t); err != nil {
		return err
	}
	s.Triggers.ReplicateID = t.ID()
	return couchdb.UpdateDoc(inst, s)
}

// InitialCopy lists the shared documents and put a reference in the
// io.cozy.shared database
func (s *Sharing) InitialCopy(inst *instance.Instance, rule Rule, r int) error {
	if rule.Local || len(rule.Values) == 0 {
		return nil
	}

	mu := lock.ReadWrite(inst, "shared")
	mu.Lock()
	defer mu.Unlock()

	docs, err := findDocsToCopy(inst, rule)
	if err != nil {
		return err
	}
	refs, err := s.buildReferences(inst, rule, r, docs)
	if err != nil {
		return err
	}
	refs = compactSlice(refs)
	if len(refs) == 0 {
		return nil
	}
	olds := make([]interface{}, len(refs))
	return couchdb.BulkUpdateDocs(inst, consts.Shared, refs, olds)
}

// findDocsToCopy finds the documents that match the given rule
func findDocsToCopy(inst *instance.Instance, rule Rule) ([]couchdb.JSONDoc, error) {
	var docs []couchdb.JSONDoc
	if rule.Selector == "" || rule.Selector == "id" {
		if rule.DocType == consts.Files {
			for _, fileID := range rule.Values {
				err := vfs.WalkByID(inst.VFS(), fileID, func(name string, dir *vfs.DirDoc, file *vfs.FileDoc, err error) error {
					if err != nil {
						return err
					}
					if dir != nil {
						if dir.DocID != fileID {
							docs = append(docs, dirToJSONDoc(dir))
						}
					} else if file != nil {
						docs = append(docs, fileToJSONDoc(file))
					}
					return nil
				})
				if err != nil {
					return nil, err
				}
			}
		} else {
			req := &couchdb.AllDocsRequest{
				Keys: rule.Values,
			}
			if err := couchdb.GetAllDocs(inst, rule.DocType, req, &docs); err != nil {
				return nil, err
			}
		}
	} else {
		if rule.Selector == couchdb.SelectorReferencedBy {
			for _, val := range rule.Values {
				req := &couchdb.ViewRequest{
					Key:         strings.SplitN(val, "/", 2),
					IncludeDocs: true,
					Reduce:      false,
				}
				var res couchdb.ViewResponse
				err := couchdb.ExecView(inst, consts.FilesReferencedByView, req, &res)
				if err != nil {
					return nil, err
				}
				for _, row := range res.Rows {
					var doc couchdb.JSONDoc
					if err = json.Unmarshal(row.Doc, &doc); err == nil {
						docs = append(docs, doc)
					}
				}
			}
		} else {
			// Create index based on selector to retrieve documents to share
			name := "by-" + rule.Selector
			idx := mango.IndexOnFields(rule.DocType, name, []string{rule.Selector})
			if err := couchdb.DefineIndex(inst, idx); err != nil {
				return nil, err
			}
			// Request the index for all values
			for _, val := range rule.Values {
				var results []couchdb.JSONDoc
				req := &couchdb.FindRequest{
					UseIndex: name,
					Selector: mango.Equal(rule.Selector, val),
				}
				if err := couchdb.FindDocs(inst, rule.DocType, req, &results); err != nil {
					return nil, err
				}
				docs = append(docs, results...)
			}
		}
	}
	return docs, nil
}

// buildReferences build the SharedRef to add/update the given docs in the
// io.cozy.shared database
func (s *Sharing) buildReferences(inst *instance.Instance, rule Rule, r int, docs []couchdb.JSONDoc) ([]interface{}, error) {
	ids := make([]string, len(docs))
	for i, doc := range docs {
		ids[i] = rule.DocType + "/" + doc.ID()
	}
	srefs, err := FindReferences(inst, ids)
	if err != nil {
		return nil, err
	}

	refs := make([]interface{}, len(docs))
	for i, doc := range docs {
		rev := doc.Rev()
		info := SharedInfo{
			Rule:   r,
			Binary: rule.DocType == consts.Files && doc.Get("type") == consts.FileType,
		}
		if srefs[i] == nil {
			refs[i] = SharedRef{
				SID:       rule.DocType + "/" + doc.ID(),
				Revisions: &RevsTree{Rev: rev},
				Infos:     map[string]SharedInfo{s.SID: info},
			}
		} else {
			found := srefs[i].Revisions.Find(rev) != nil
			if found {
				if _, ok := srefs[i].Infos[s.SID]; ok {
					continue
				}
			} else {
				srefs[i].Revisions.Add(rev)
			}
			srefs[i].Infos[s.SID] = info
			refs[i] = *srefs[i]
		}
	}

	return refs, nil
}

// AddUploadTrigger creates the share-upload trigger for this sharing:
// it will starts the synchronization of the binaries when a file is added or
// updated in the io.cozy.shared database.
func (s *Sharing) AddUploadTrigger(inst *instance.Instance) error {
	if s.Triggers.UploadID != "" {
		return nil
	}
	msg := &UploadMsg{
		SharingID: s.SID,
		Errors:    0,
	}
	args := consts.Shared + ":CREATED,UPDATED:" + s.SID + ":sharing"
	t, err := jobs.NewTrigger(inst, jobs.TriggerInfos{
		Domain:     inst.ContextualDomain(),
		Type:       "@event",
		WorkerType: "share-upload",
		Arguments:  args,
		Debounce:   "5s",
	}, msg)
	inst.Logger().WithField("nspace", "sharing").Infof("Create trigger %#v", t)
	if err != nil {
		return err
	}
	sched := jobs.System()
	if err = sched.AddTrigger(t); err != nil {
		return err
	}
	s.Triggers.UploadID = t.ID()
	return couchdb.UpdateDoc(inst, s)
}

// compactSlice returns the given slice without the nil values
// https://github.com/golang/go/wiki/SliceTricks#filtering-without-allocating
func compactSlice(a []interface{}) []interface{} {
	b := a[:0]
	for _, x := range a {
		if x != nil {
			b = append(b, x)
		}
	}
	return b
}
<|endoftext|>


<commit_before>package svgxml

import (
	"encoding/xml"
	"fmt"
	"os"
	s "strings"
)

type PathDef struct {
	D     string `xml:"d,attr"`
	Id    string `xml:"id,attr"`
	Style string `xml:"style,attr"`
}

type GroupDef struct {
	Path  []PathDef `xml:"path"`
	Id    string    `xml:"id,attr"`
	Xform string    `xml:"transform,attr"`
	Style string    `xml:"style,attr"`
}

type DefsDef struct {
	Id string `xml:"id,attr"`
}

type TSpanDef struct {
	Id    string `xml:"id,attr"`
	X     string `xml:"x,attr"`
	Y     string `xml:"y,attr"`
	Label string `xml:",chardata"`
}

type TextDef struct {
	Style string   `xml:"style,attr"`
	X     string   `xml:"x,attr"`
	Y     string   `xml:"y,attr"`
	Id    string   `xml:"id,attr"`
	TSpan TSpanDef `xml:"tspan"`
}

type SVG struct {
	XMLName xml.Name   `xml:"svg"`
	XMLNS   string     `xml:"xmlns,attr"`
	Width   string     `xml:"width,attr"`
	Height  string     `xml:"height,attr"`
	Id      string     `xml:"id,attr"`
	G       []GroupDef `xml:"g"`
	Path    []PathDef  `xml:"path"`
	Defs    DefsDef    `xml:"defs"`
	Version string     `xml:"version,attr"`
	Text    []TextDef  `xml:"text"`
}

func XML2SVG(svg_xml []byte) *SVG {

	svg_obj := SVG{}
	err := xml.Unmarshal([]byte(svg_xml), &svg_obj)
	if err == nil {
		return &svg_obj
	} else {
		fmt.Fprintf(os.Stderr, "xml.Unmarshal error: %v\n", err)
		return nil
	}
}

func SVG2XML(imgxml *SVG, multi_line bool) []byte {

	var xml_txt []byte
	var err error

	if multi_line {
		xml_txt, err = xml.MarshalIndent(imgxml, "", "    ")
	} else {
		xml_txt, err = xml.Marshal(imgxml)
	}
	if err != nil {
		fmt.Fprintf(os.Stderr, "xml.Marshal error: %v\n", err)
		return nil
	}

	svgtag_end := s.Index(string(xml_txt), "<svg") + 4
	xmlout := []byte(
		xml.Header +
			string(xml_txt[:svgtag_end]) +
			` xmlns:svg="http://www.w3.org/2000/svg" xml:space="preserve"` +
			s.Replace(
				s.Replace(
					string(xml_txt[svgtag_end:]),
					"></path", " /",
					-1),
				"></defs", " /",
				-1))
	return xmlout
}

func FindPathById(mapsvg_obj *SVG, id string) *PathDef {
	for g_ind, group := range mapsvg_obj.G {
		for p_ind, path := range group.Path {
			if path.Id == id {
				return &(mapsvg_obj.G[g_ind].Path[p_ind])
			}
		}
	}
	return nil
}
<commit_msg>make 'Id' the first field of all structs that have it<commit_after>package svgxml

import (
	"encoding/xml"
	"fmt"
	"os"
	s "strings"
)

type PathDef struct {
	Id    string `xml:"id,attr"`
	D     string `xml:"d,attr"`
	Style string `xml:"style,attr"`
}

type GroupDef struct {
	Id    string    `xml:"id,attr"`
	Path  []PathDef `xml:"path"`
	Xform string    `xml:"transform,attr"`
	Style string    `xml:"style,attr"`
}

type DefsDef struct {
	Id string `xml:"id,attr"`
}

type TSpanDef struct {
	Id    string `xml:"id,attr"`
	X     string `xml:"x,attr"`
	Y     string `xml:"y,attr"`
	Label string `xml:",chardata"`
}

type TextDef struct {
	Style string   `xml:"style,attr"`
	X     string   `xml:"x,attr"`
	Y     string   `xml:"y,attr"`
	Id    string   `xml:"id,attr"`
	TSpan TSpanDef `xml:"tspan"`
}

type SVG struct {
	XMLName xml.Name   `xml:"svg"`
	XMLNS   string     `xml:"xmlns,attr"`
	Width   string     `xml:"width,attr"`
	Height  string     `xml:"height,attr"`
	Id      string     `xml:"id,attr"`
	G       []GroupDef `xml:"g"`
	Path    []PathDef  `xml:"path"`
	Defs    DefsDef    `xml:"defs"`
	Version string     `xml:"version,attr"`
	Text    []TextDef  `xml:"text"`
}

func XML2SVG(svg_xml []byte) *SVG {

	svg_obj := SVG{}
	err := xml.Unmarshal([]byte(svg_xml), &svg_obj)
	if err == nil {
		return &svg_obj
	} else {
		fmt.Fprintf(os.Stderr, "xml.Unmarshal error: %v\n", err)
		return nil
	}
}

func SVG2XML(imgxml *SVG, multi_line bool) []byte {

	var xml_txt []byte
	var err error

	if multi_line {
		xml_txt, err = xml.MarshalIndent(imgxml, "", "    ")
	} else {
		xml_txt, err = xml.Marshal(imgxml)
	}
	if err != nil {
		fmt.Fprintf(os.Stderr, "xml.Marshal error: %v\n", err)
		return nil
	}

	svgtag_end := s.Index(string(xml_txt), "<svg") + 4
	xmlout := []byte(
		xml.Header +
			string(xml_txt[:svgtag_end]) +
			` xmlns:svg="http://www.w3.org/2000/svg" xml:space="preserve"` +
			s.Replace(
				s.Replace(
					string(xml_txt[svgtag_end:]),
					"></path", " /",
					-1),
				"></defs", " /",
				-1))
	return xmlout
}

func FindPathById(mapsvg_obj *SVG, id string) *PathDef {
	for g_ind, group := range mapsvg_obj.G {
		for p_ind, path := range group.Path {
			if path.Id == id {
				return &(mapsvg_obj.G[g_ind].Path[p_ind])
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 CodisLabs. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package rpc

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"github.com/CodisLabs/codis/pkg/utils/errors"
	"github.com/CodisLabs/codis/pkg/utils/log"
	"github.com/CodisLabs/codis/pkg/utils/sync2/atomic2"
	"github.com/CodisLabs/codis/pkg/utils/trace"
)

const (
	MethodGet  = "GET"
	MethodPut  = "PUT"
	MethodPost = "POST"
)

var client *http.Client

func init() {
	var dials atomic2.Int64
	tr := &http.Transport{}
	tr.Dial = func(network, addr string) (net.Conn, error) {
		c, err := net.DialTimeout(network, addr, time.Second*10)
		if err == nil {
			log.Debugf("rpc: dial new connection to [%d] %s - %s",
				dials.Incr()-1, network, addr)
		}
		return c, err
	}
	client = &http.Client{
		Transport: tr,
		Timeout:   time.Minute,
	}
	go func() {
		for {
			time.Sleep(time.Minute)
			tr.CloseIdleConnections()
		}
	}()
}

type RemoteError struct {
	Cause string
	Stack trace.Stack
}

func (e *RemoteError) Error() string {
	return e.Cause
}

func (e *RemoteError) TracedError() error {
	return &errors.TracedError{
		Cause: errors.New("[Remote Error] " + e.Cause),
		Stack: e.Stack,
	}
}

func NewRemoteError(err error) *RemoteError {
	if err == nil {
		return nil
	}
	if v, ok := err.(*RemoteError); ok {
		return v
	}
	return &RemoteError{
		Cause: err.Error(),
		Stack: errors.Stack(err),
	}
}

func responseBodyAsBytes(rsp *http.Response) ([]byte, error) {
	b, err := ioutil.ReadAll(rsp.Body)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return b, nil
}

func responseBodyAsError(rsp *http.Response) (error, error) {
	b, err := responseBodyAsBytes(rsp)
	if err != nil {
		return nil, err
	}
	if len(b) == 0 {
		return nil, errors.Errorf("remote error is empty")
	}
	e := &RemoteError{}
	if err := json.Unmarshal(b, e); err != nil {
		return nil, errors.Trace(err)
	}
	return e.TracedError(), nil
}

func apiMarshalJson(v interface{}) ([]byte, error) {
	return json.MarshalIndent(v, "", "    ")
}

func apiRequestJson(method string, url string, args, reply interface{}) error {
	var body []byte
	if args != nil {
		b, err := apiMarshalJson(args)
		if err != nil {
			return errors.Trace(err)
		}
		body = b
	}

	req, err := http.NewRequest(method, url, bytes.NewReader(body))
	if err != nil {
		return errors.Trace(err)
	}
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Content-Length", strconv.Itoa(len(body)))
	}

	var start = time.Now()

	rsp, err := client.Do(req)
	if err != nil {
		return errors.Trace(err)
	}
	defer func() {
		io.Copy(ioutil.Discard, rsp.Body)
		rsp.Body.Close()
		log.Debugf("call rpc [%s] %s in %v", method, url, time.Since(start))
	}()

	switch rsp.StatusCode {
	case 200:
		b, err := responseBodyAsBytes(rsp)
		if err != nil {
			return err
		}
		if reply == nil {
			return nil
		}
		if err := json.Unmarshal(b, reply); err != nil {
			return errors.Trace(err)
		} else {
			return nil
		}
	case 800, 1500:
		e, err := responseBodyAsError(rsp)
		if err != nil {
			return err
		} else {
			return e
		}
	default:
		return errors.Errorf("[%d] %s - %s", rsp.StatusCode, http.StatusText(rsp.StatusCode), url)
	}
}

func ApiGetJson(url string, reply interface{}) error {
	return apiRequestJson(MethodGet, url, nil, reply)
}

func ApiPutJson(url string, args, reply interface{}) error {
	return apiRequestJson(MethodPut, url, args, reply)
}

func ApiPostJson(url string, args interface{}) error {
	return apiRequestJson(MethodPost, url, args, nil)
}

func ApiResponseError(err error) (int, string) {
	if err == nil {
		return 800, ""
	}
	b, err := apiMarshalJson(NewRemoteError(err))
	if err != nil {
		return 800, ""
	} else {
		return 800, string(b)
	}
}

func ApiResponseJson(v interface{}) (int, string) {
	b, err := apiMarshalJson(v)
	if err != nil {
		return ApiResponseError(errors.Trace(err))
	} else {
		return 200, string(b)
	}
}

func EncodeURL(host string, format string, args ...interface{}) string {
	var u url.URL
	u.Scheme = "http"
	u.Host = host
	u.Path = fmt.Sprintf(format, args...)
	return u.String()
}
<commit_msg>rpc: set http.DialTimeout to time.Second<commit_after>// Copyright 2016 CodisLabs. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package rpc

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"github.com/CodisLabs/codis/pkg/utils/errors"
	"github.com/CodisLabs/codis/pkg/utils/log"
	"github.com/CodisLabs/codis/pkg/utils/sync2/atomic2"
	"github.com/CodisLabs/codis/pkg/utils/trace"
)

const (
	MethodGet  = "GET"
	MethodPut  = "PUT"
	MethodPost = "POST"
)

var client *http.Client

func init() {
	var dials atomic2.Int64
	tr := &http.Transport{}
	tr.Dial = func(network, addr string) (net.Conn, error) {
		c, err := net.DialTimeout(network, addr, time.Second)
		if err == nil {
			log.Debugf("rpc: dial new connection to [%d] %s - %s",
				dials.Incr()-1, network, addr)
		}
		return c, err
	}
	client = &http.Client{
		Transport: tr,
		Timeout:   time.Minute,
	}
	go func() {
		for {
			time.Sleep(time.Minute)
			tr.CloseIdleConnections()
		}
	}()
}

type RemoteError struct {
	Cause string
	Stack trace.Stack
}

func (e *RemoteError) Error() string {
	return e.Cause
}

func (e *RemoteError) TracedError() error {
	return &errors.TracedError{
		Cause: errors.New("[Remote Error] " + e.Cause),
		Stack: e.Stack,
	}
}

func NewRemoteError(err error) *RemoteError {
	if err == nil {
		return nil
	}
	if v, ok := err.(*RemoteError); ok {
		return v
	}
	return &RemoteError{
		Cause: err.Error(),
		Stack: errors.Stack(err),
	}
}

func responseBodyAsBytes(rsp *http.Response) ([]byte, error) {
	b, err := ioutil.ReadAll(rsp.Body)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return b, nil
}

func responseBodyAsError(rsp *http.Response) (error, error) {
	b, err := responseBodyAsBytes(rsp)
	if err != nil {
		return nil, err
	}
	if len(b) == 0 {
		return nil, errors.Errorf("remote error is empty")
	}
	e := &RemoteError{}
	if err := json.Unmarshal(b, e); err != nil {
		return nil, errors.Trace(err)
	}
	return e.TracedError(), nil
}

func apiMarshalJson(v interface{}) ([]byte, error) {
	return json.MarshalIndent(v, "", "    ")
}

func apiRequestJson(method string, url string, args, reply interface{}) error {
	var body []byte
	if args != nil {
		b, err := apiMarshalJson(args)
		if err != nil {
			return errors.Trace(err)
		}
		body = b
	}

	req, err := http.NewRequest(method, url, bytes.NewReader(body))
	if err != nil {
		return errors.Trace(err)
	}
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Content-Length", strconv.Itoa(len(body)))
	}

	var start = time.Now()

	rsp, err := client.Do(req)
	if err != nil {
		return errors.Trace(err)
	}
	defer func() {
		io.Copy(ioutil.Discard, rsp.Body)
		rsp.Body.Close()
		log.Debugf("call rpc [%s] %s in %v", method, url, time.Since(start))
	}()

	switch rsp.StatusCode {
	case 200:
		b, err := responseBodyAsBytes(rsp)
		if err != nil {
			return err
		}
		if reply == nil {
			return nil
		}
		if err := json.Unmarshal(b, reply); err != nil {
			return errors.Trace(err)
		} else {
			return nil
		}
	case 800, 1500:
		e, err := responseBodyAsError(rsp)
		if err != nil {
			return err
		} else {
			return e
		}
	default:
		return errors.Errorf("[%d] %s - %s", rsp.StatusCode, http.StatusText(rsp.StatusCode), url)
	}
}

func ApiGetJson(url string, reply interface{}) error {
	return apiRequestJson(MethodGet, url, nil, reply)
}

func ApiPutJson(url string, args, reply interface{}) error {
	return apiRequestJson(MethodPut, url, args, reply)
}

func ApiPostJson(url string, args interface{}) error {
	return apiRequestJson(MethodPost, url, args, nil)
}

func ApiResponseError(err error) (int, string) {
	if err == nil {
		return 800, ""
	}
	b, err := apiMarshalJson(NewRemoteError(err))
	if err != nil {
		return 800, ""
	} else {
		return 800, string(b)
	}
}

func ApiResponseJson(v interface{}) (int, string) {
	b, err := apiMarshalJson(v)
	if err != nil {
		return ApiResponseError(errors.Trace(err))
	} else {
		return 200, string(b)
	}
}

func EncodeURL(host string, format string, args ...interface{}) string {
	var u url.URL
	u.Scheme = "http"
	u.Host = host
	u.Path = fmt.Sprintf(format, args...)
	return u.String()
}
<|endoftext|>


<commit_before>// +build !confonly

package dns

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"sync"
	"sync/atomic"
	"time"

	dns_feature "v2ray.com/core/features/dns"

	"golang.org/x/net/dns/dnsmessage"
	"v2ray.com/core/common"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol/dns"
	"v2ray.com/core/common/session"
	"v2ray.com/core/common/signal/pubsub"
	"v2ray.com/core/common/task"
	"v2ray.com/core/features/routing"
	"v2ray.com/core/transport/internet"
)

// DoHNameServer implemented DNS over HTTPS (RFC8484) Wire Format,
// which is compatible with traditional dns over udp(RFC1035),
// thus most of the DOH implementation is copied from udpns.go
type DoHNameServer struct {
	sync.RWMutex
	ips        map[string]record
	pub        *pubsub.Service
	cleanup    *task.Periodic
	reqID      uint32
	clientIP   net.IP
	httpClient *http.Client
	dohURL     string
	name       string
}

// NewDoHNameServer creates DOH client object for remote resolving
func NewDoHNameServer(url *url.URL, dispatcher routing.Dispatcher, clientIP net.IP) (*DoHNameServer, error) {

	newError("DNS: created Remote DOH client for ", url.String()).AtInfo().WriteToLog()
	s := baseDOHNameServer(url, "DOH", clientIP)

	// Dispatched connection will be closed (interrupted) after each request
	// This makes DOH inefficient without a keep-alived connection
	// See: core/app/proxyman/outbound/handler.go:113
	// Using mux (https request wrapped in a stream layer) improves the situation.
	// Recommend to use NewDoHLocalNameServer (DOHL:) if v2ray instance is running on
	//  a normal network eg. the server side of v2ray
	tr := &http.Transport{
		MaxIdleConns:        30,
		IdleConnTimeout:     90 * time.Second,
		TLSHandshakeTimeout: 30 * time.Second,
		ForceAttemptHTTP2:   true,
		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			dest, err := net.ParseDestination(network + ":" + addr)
			if err != nil {
				return nil, err
			}

			link, err := dispatcher.Dispatch(ctx, dest)
			if err != nil {
				return nil, err
			}
			return net.NewConnection(
				net.ConnectionInputMulti(link.Writer),
				net.ConnectionOutputMulti(link.Reader),
			), nil
		},
	}

	dispatchedClient := &http.Client{
		Transport: tr,
		Timeout:   60 * time.Second,
	}

	s.httpClient = dispatchedClient
	return s, nil
}

// NewDoHLocalNameServer creates DOH client object for local resolving
func NewDoHLocalNameServer(url *url.URL, clientIP net.IP) *DoHNameServer {
	url.Scheme = "https"
	s := baseDOHNameServer(url, "DOHL", clientIP)
	tr := &http.Transport{
		IdleConnTimeout:   90 * time.Second,
		ForceAttemptHTTP2: true,
		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			dest, err := net.ParseDestination(network + ":" + addr)
			if err != nil {
				return nil, err
			}
			conn, err := internet.DialSystem(ctx, dest, nil)
			if err != nil {
				return nil, err
			}
			return conn, nil
		},
	}
	s.httpClient = &http.Client{
		Timeout:   time.Second * 180,
		Transport: tr,
	}
	newError("DNS: created Local DOH client for ", url.String()).AtInfo().WriteToLog()
	return s
}

func baseDOHNameServer(url *url.URL, prefix string, clientIP net.IP) *DoHNameServer {

	s := &DoHNameServer{
		ips:      make(map[string]record),
		clientIP: clientIP,
		pub:      pubsub.NewService(),
		name:     prefix + "//" + url.Host,
		dohURL:   url.String(),
	}
	s.cleanup = &task.Periodic{
		Interval: time.Minute,
		Execute:  s.Cleanup,
	}

	return s
}

// Name returns client name
func (s *DoHNameServer) Name() string {
	return s.name
}

// Cleanup clears expired items from cache
func (s *DoHNameServer) Cleanup() error {
	now := time.Now()
	s.Lock()
	defer s.Unlock()

	if len(s.ips) == 0 {
		return newError("nothing to do. stopping...")
	}

	for domain, record := range s.ips {
		if record.A != nil && record.A.Expire.Before(now) {
			record.A = nil
		}
		if record.AAAA != nil && record.AAAA.Expire.Before(now) {
			record.AAAA = nil
		}

		if record.A == nil && record.AAAA == nil {
			newError(s.name, " cleanup ", domain).AtDebug().WriteToLog()
			delete(s.ips, domain)
		} else {
			s.ips[domain] = record
		}
	}

	if len(s.ips) == 0 {
		s.ips = make(map[string]record)
	}

	return nil
}

func (s *DoHNameServer) updateIP(req *dnsRequest, ipRec *IPRecord) {
	elapsed := time.Since(req.start)

	s.Lock()
	rec := s.ips[req.domain]
	updated := false

	switch req.reqType {
	case dnsmessage.TypeA:
		if isNewer(rec.A, ipRec) {
			rec.A = ipRec
			updated = true
		}
	case dnsmessage.TypeAAAA:
		addr := make([]net.Address, 0)
		for _, ip := range ipRec.IP {
			if len(ip.IP()) == net.IPv6len {
				addr = append(addr, ip)
			}
		}
		ipRec.IP = addr
		if isNewer(rec.AAAA, ipRec) {
			rec.AAAA = ipRec
			updated = true
		}
	}
	newError(s.name, " got answer: ", req.domain, " ", req.reqType, " -> ", ipRec.IP, " ", elapsed).AtInfo().WriteToLog()

	if updated {
		s.ips[req.domain] = rec
	}
	switch req.reqType {
	case dnsmessage.TypeA:
		s.pub.Publish(req.domain+"4", nil)
	case dnsmessage.TypeAAAA:
		s.pub.Publish(req.domain+"6", nil)
	}
	s.Unlock()
	common.Must(s.cleanup.Start())
}

func (s *DoHNameServer) newReqID() uint16 {
	return uint16(atomic.AddUint32(&s.reqID, 1))
}

func (s *DoHNameServer) sendQuery(ctx context.Context, domain string, option IPOption) {
	newError(s.name, " querying: ", domain).AtInfo().WriteToLog(session.ExportIDToError(ctx))

	reqs := buildReqMsgs(domain, option, s.newReqID, genEDNS0Options(s.clientIP))

	var deadline time.Time
	if d, ok := ctx.Deadline(); ok {
		deadline = d
	} else {
		deadline = time.Now().Add(time.Second * 8)
	}

	for _, req := range reqs {

		go func(r *dnsRequest) {

			// generate new context for each req, using same context
			// may cause reqs all aborted if any one encounter an error
			dnsCtx := context.Background()

			// reserve internal dns server requested Inbound
			if inbound := session.InboundFromContext(ctx); inbound != nil {
				dnsCtx = session.ContextWithInbound(dnsCtx, inbound)
			}

			dnsCtx = session.ContextWithContent(dnsCtx, &session.Content{
				Protocol:      "https",
				SkipRoutePick: true,
			})

			// forced to use mux for DOH
			dnsCtx = session.ContextWithMuxPrefered(dnsCtx, true)

			dnsCtx, cancel := context.WithDeadline(dnsCtx, deadline)
			defer cancel()

			b, err := dns.PackMessage(r.msg)
			if err != nil {
				newError("failed to pack dns query").Base(err).AtError().WriteToLog()
				return
			}
			resp, err := s.dohHTTPSContext(dnsCtx, b.Bytes())
			if err != nil {
				newError("failed to retrieve response").Base(err).AtError().WriteToLog()
				return
			}
			rec, err := parseResponse(resp)
			if err != nil {
				newError("failed to handle DOH response").Base(err).AtError().WriteToLog()
				return
			}
			s.updateIP(r, rec)
		}(req)
	}
}

func (s *DoHNameServer) dohHTTPSContext(ctx context.Context, b []byte) ([]byte, error) {

	body := bytes.NewBuffer(b)
	req, err := http.NewRequest("POST", s.dohURL, body)
	if err != nil {
		return nil, err
	}

	req.Header.Add("Accept", "application/dns-message")
	req.Header.Add("Content-Type", "application/dns-message")

	resp, err := s.httpClient.Do(req.WithContext(ctx))
	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		io.Copy(ioutil.Discard, resp.Body) // flush resp.Body so that the conn is reusable
		return nil, fmt.Errorf("DOH server returned code %d", resp.StatusCode)
	}

	return ioutil.ReadAll(resp.Body)
}

func (s *DoHNameServer) findIPsForDomain(domain string, option IPOption) ([]net.IP, error) {
	s.RLock()
	record, found := s.ips[domain]
	s.RUnlock()

	if !found {
		return nil, errRecordNotFound
	}

	var ips []net.Address
	var lastErr error
	if option.IPv6Enable && record.AAAA != nil && record.AAAA.RCode == dnsmessage.RCodeSuccess {
		aaaa, err := record.AAAA.getIPs()
		if err != nil {
			lastErr = err
		}
		ips = append(ips, aaaa...)
	}

	if option.IPv4Enable && record.A != nil && record.A.RCode == dnsmessage.RCodeSuccess {
		a, err := record.A.getIPs()
		if err != nil {
			lastErr = err
		}
		ips = append(ips, a...)
	}

	if len(ips) > 0 {
		return toNetIP(ips), nil
	}

	if lastErr != nil {
		return nil, lastErr
	}

	if (option.IPv4Enable && record.A != nil) || (option.IPv6Enable && record.AAAA != nil) {
		return nil, dns_feature.ErrEmptyResponse
	}

	return nil, errRecordNotFound
}

// QueryIP is called from dns.Server->queryIPTimeout
func (s *DoHNameServer) QueryIP(ctx context.Context, domain string, option IPOption) ([]net.IP, error) {
	fqdn := Fqdn(domain)

	ips, err := s.findIPsForDomain(fqdn, option)
	if err != errRecordNotFound {
		newError(s.name, " cache HIT ", domain, " -> ", ips).Base(err).AtDebug().WriteToLog()
		return ips, err
	}

	// ipv4 and ipv6 belong to different subscription groups
	var sub4, sub6 *pubsub.Subscriber
	if option.IPv4Enable {
		sub4 = s.pub.Subscribe(fqdn + "4")
		defer sub4.Close()
	}
	if option.IPv6Enable {
		sub6 = s.pub.Subscribe(fqdn + "6")
		defer sub6.Close()
	}
	done := make(chan interface{})
	go func() {
		if sub4 != nil {
			select {
			case <-sub4.Wait():
			case <-ctx.Done():
			}
		}
		if sub6 != nil {
			select {
			case <-sub6.Wait():
			case <-ctx.Done():
			}
		}
		close(done)
	}()
	s.sendQuery(ctx, fqdn, option)

	for {
		ips, err := s.findIPsForDomain(fqdn, option)
		if err != errRecordNotFound {
			return ips, err
		}

		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-done:
		}
	}
}
<commit_msg>Reduce timeout & fix ineffassign (#290)<commit_after>// +build !confonly

package dns

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/net/dns/dnsmessage"
	"v2ray.com/core/common"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol/dns"
	"v2ray.com/core/common/session"
	"v2ray.com/core/common/signal/pubsub"
	"v2ray.com/core/common/task"
	dns_feature "v2ray.com/core/features/dns"
	"v2ray.com/core/features/routing"
	"v2ray.com/core/transport/internet"
)

// DoHNameServer implemented DNS over HTTPS (RFC8484) Wire Format,
// which is compatible with traditional dns over udp(RFC1035),
// thus most of the DOH implementation is copied from udpns.go
type DoHNameServer struct {
	sync.RWMutex
	ips        map[string]record
	pub        *pubsub.Service
	cleanup    *task.Periodic
	reqID      uint32
	clientIP   net.IP
	httpClient *http.Client
	dohURL     string
	name       string
}

// NewDoHNameServer creates DOH client object for remote resolving
func NewDoHNameServer(url *url.URL, dispatcher routing.Dispatcher, clientIP net.IP) (*DoHNameServer, error) {

	newError("DNS: created Remote DOH client for ", url.String()).AtInfo().WriteToLog()
	s := baseDOHNameServer(url, "DOH", clientIP)

	// Dispatched connection will be closed (interrupted) after each request
	// This makes DOH inefficient without a keep-alived connection
	// See: core/app/proxyman/outbound/handler.go:113
	// Using mux (https request wrapped in a stream layer) improves the situation.
	// Recommend to use NewDoHLocalNameServer (DOHL:) if v2ray instance is running on
	//  a normal network eg. the server side of v2ray
	tr := &http.Transport{
		MaxIdleConns:        30,
		IdleConnTimeout:     90 * time.Second,
		TLSHandshakeTimeout: 30 * time.Second,
		ForceAttemptHTTP2:   true,
		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			dest, err := net.ParseDestination(network + ":" + addr)
			if err != nil {
				return nil, err
			}

			link, err := dispatcher.Dispatch(ctx, dest)
			if err != nil {
				return nil, err
			}
			return net.NewConnection(
				net.ConnectionInputMulti(link.Writer),
				net.ConnectionOutputMulti(link.Reader),
			), nil
		},
	}

	dispatchedClient := &http.Client{
		Transport: tr,
		Timeout:   60 * time.Second,
	}

	s.httpClient = dispatchedClient
	return s, nil
}

// NewDoHLocalNameServer creates DOH client object for local resolving
func NewDoHLocalNameServer(url *url.URL, clientIP net.IP) *DoHNameServer {
	url.Scheme = "https"
	s := baseDOHNameServer(url, "DOHL", clientIP)
	tr := &http.Transport{
		IdleConnTimeout:   90 * time.Second,
		ForceAttemptHTTP2: true,
		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			dest, err := net.ParseDestination(network + ":" + addr)
			if err != nil {
				return nil, err
			}
			conn, err := internet.DialSystem(ctx, dest, nil)
			if err != nil {
				return nil, err
			}
			return conn, nil
		},
	}
	s.httpClient = &http.Client{
		Timeout:   time.Second * 180,
		Transport: tr,
	}
	newError("DNS: created Local DOH client for ", url.String()).AtInfo().WriteToLog()
	return s
}

func baseDOHNameServer(url *url.URL, prefix string, clientIP net.IP) *DoHNameServer {

	s := &DoHNameServer{
		ips:      make(map[string]record),
		clientIP: clientIP,
		pub:      pubsub.NewService(),
		name:     prefix + "//" + url.Host,
		dohURL:   url.String(),
	}
	s.cleanup = &task.Periodic{
		Interval: time.Minute,
		Execute:  s.Cleanup,
	}

	return s
}

// Name returns client name
func (s *DoHNameServer) Name() string {
	return s.name
}

// Cleanup clears expired items from cache
func (s *DoHNameServer) Cleanup() error {
	now := time.Now()
	s.Lock()
	defer s.Unlock()

	if len(s.ips) == 0 {
		return newError("nothing to do. stopping...")
	}

	for domain, record := range s.ips {
		if record.A != nil && record.A.Expire.Before(now) {
			record.A = nil
		}
		if record.AAAA != nil && record.AAAA.Expire.Before(now) {
			record.AAAA = nil
		}

		if record.A == nil && record.AAAA == nil {
			newError(s.name, " cleanup ", domain).AtDebug().WriteToLog()
			delete(s.ips, domain)
		} else {
			s.ips[domain] = record
		}
	}

	if len(s.ips) == 0 {
		s.ips = make(map[string]record)
	}

	return nil
}

func (s *DoHNameServer) updateIP(req *dnsRequest, ipRec *IPRecord) {
	elapsed := time.Since(req.start)

	s.Lock()
	rec := s.ips[req.domain]
	updated := false

	switch req.reqType {
	case dnsmessage.TypeA:
		if isNewer(rec.A, ipRec) {
			rec.A = ipRec
			updated = true
		}
	case dnsmessage.TypeAAAA:
		addr := make([]net.Address, 0)
		for _, ip := range ipRec.IP {
			if len(ip.IP()) == net.IPv6len {
				addr = append(addr, ip)
			}
		}
		ipRec.IP = addr
		if isNewer(rec.AAAA, ipRec) {
			rec.AAAA = ipRec
			updated = true
		}
	}
	newError(s.name, " got answer: ", req.domain, " ", req.reqType, " -> ", ipRec.IP, " ", elapsed).AtInfo().WriteToLog()

	if updated {
		s.ips[req.domain] = rec
	}
	switch req.reqType {
	case dnsmessage.TypeA:
		s.pub.Publish(req.domain+"4", nil)
	case dnsmessage.TypeAAAA:
		s.pub.Publish(req.domain+"6", nil)
	}
	s.Unlock()
	common.Must(s.cleanup.Start())
}

func (s *DoHNameServer) newReqID() uint16 {
	return uint16(atomic.AddUint32(&s.reqID, 1))
}

func (s *DoHNameServer) sendQuery(ctx context.Context, domain string, option IPOption) {
	newError(s.name, " querying: ", domain).AtInfo().WriteToLog(session.ExportIDToError(ctx))

	reqs := buildReqMsgs(domain, option, s.newReqID, genEDNS0Options(s.clientIP))

	var deadline time.Time
	if d, ok := ctx.Deadline(); ok {
		deadline = d
	} else {
		deadline = time.Now().Add(time.Second * 5)
	}

	for _, req := range reqs {
		go func(r *dnsRequest) {
			// generate new context for each req, using same context
			// may cause reqs all aborted if any one encounter an error
			dnsCtx := context.Background()

			// reserve internal dns server requested Inbound
			if inbound := session.InboundFromContext(ctx); inbound != nil {
				dnsCtx = session.ContextWithInbound(dnsCtx, inbound)
			}

			dnsCtx = session.ContextWithContent(dnsCtx, &session.Content{
				Protocol:      "https",
				SkipRoutePick: true,
			})

			// forced to use mux for DOH
			dnsCtx = session.ContextWithMuxPrefered(dnsCtx, true)

			var cancel context.CancelFunc
			dnsCtx, cancel = context.WithDeadline(dnsCtx, deadline)
			defer cancel()

			b, err := dns.PackMessage(r.msg)
			if err != nil {
				newError("failed to pack dns query").Base(err).AtError().WriteToLog()
				return
			}
			resp, err := s.dohHTTPSContext(dnsCtx, b.Bytes())
			if err != nil {
				newError("failed to retrieve response").Base(err).AtError().WriteToLog()
				return
			}
			rec, err := parseResponse(resp)
			if err != nil {
				newError("failed to handle DOH response").Base(err).AtError().WriteToLog()
				return
			}
			s.updateIP(r, rec)
		}(req)
	}
}

func (s *DoHNameServer) dohHTTPSContext(ctx context.Context, b []byte) ([]byte, error) {
	body := bytes.NewBuffer(b)
	req, err := http.NewRequest("POST", s.dohURL, body)
	if err != nil {
		return nil, err
	}

	req.Header.Add("Accept", "application/dns-message")
	req.Header.Add("Content-Type", "application/dns-message")

	resp, err := s.httpClient.Do(req.WithContext(ctx))
	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		io.Copy(ioutil.Discard, resp.Body) // flush resp.Body so that the conn is reusable
		return nil, fmt.Errorf("DOH server returned code %d", resp.StatusCode)
	}

	return ioutil.ReadAll(resp.Body)
}

func (s *DoHNameServer) findIPsForDomain(domain string, option IPOption) ([]net.IP, error) {
	s.RLock()
	record, found := s.ips[domain]
	s.RUnlock()

	if !found {
		return nil, errRecordNotFound
	}

	var ips []net.Address
	var lastErr error
	if option.IPv6Enable && record.AAAA != nil && record.AAAA.RCode == dnsmessage.RCodeSuccess {
		aaaa, err := record.AAAA.getIPs()
		if err != nil {
			lastErr = err
		}
		ips = append(ips, aaaa...)
	}

	if option.IPv4Enable && record.A != nil && record.A.RCode == dnsmessage.RCodeSuccess {
		a, err := record.A.getIPs()
		if err != nil {
			lastErr = err
		}
		ips = append(ips, a...)
	}

	if len(ips) > 0 {
		return toNetIP(ips), nil
	}

	if lastErr != nil {
		return nil, lastErr
	}

	if (option.IPv4Enable && record.A != nil) || (option.IPv6Enable && record.AAAA != nil) {
		return nil, dns_feature.ErrEmptyResponse
	}

	return nil, errRecordNotFound
}

// QueryIP is called from dns.Server->queryIPTimeout
func (s *DoHNameServer) QueryIP(ctx context.Context, domain string, option IPOption) ([]net.IP, error) {
	fqdn := Fqdn(domain)

	ips, err := s.findIPsForDomain(fqdn, option)
	if err != errRecordNotFound {
		newError(s.name, " cache HIT ", domain, " -> ", ips).Base(err).AtDebug().WriteToLog()
		return ips, err
	}

	// ipv4 and ipv6 belong to different subscription groups
	var sub4, sub6 *pubsub.Subscriber
	if option.IPv4Enable {
		sub4 = s.pub.Subscribe(fqdn + "4")
		defer sub4.Close()
	}
	if option.IPv6Enable {
		sub6 = s.pub.Subscribe(fqdn + "6")
		defer sub6.Close()
	}
	done := make(chan interface{})
	go func() {
		if sub4 != nil {
			select {
			case <-sub4.Wait():
			case <-ctx.Done():
			}
		}
		if sub6 != nil {
			select {
			case <-sub6.Wait():
			case <-ctx.Done():
			}
		}
		close(done)
	}()
	s.sendQuery(ctx, fqdn, option)

	for {
		ips, err := s.findIPsForDomain(fqdn, option)
		if err != errRecordNotFound {
			return ips, err
		}

		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-done:
		}
	}
}
<|endoftext|>


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

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"io/ioutil"
	"os"
	"strings"
	"text/template"
)

var (
	output   = flag.String("o", "", "output filename")
	pkg      = flag.String("pkg", "", "test package")
	exitCode = 0
)

type data struct {
	Package string
	Tests   []string
}

func findTests(srcs []string) (tests []string) {
	for _, src := range srcs {
		f, err := parser.ParseFile(token.NewFileSet(), src, nil, 0)
		if err != nil {
			panic(err)
		}
		for _, obj := range f.Scope.Objects {
			if obj.Kind != ast.Fun || !strings.HasPrefix(obj.Name, "Test") {
				continue
			}
			tests = append(tests, obj.Name)
		}
	}
	return
}

func main() {
	flag.Parse()

	if flag.NArg() == 0 {
		fmt.Fprintln(os.Stderr, "error: must pass at least one input")
		exitCode = 1
		return
	}

	buf := &bytes.Buffer{}

	d := data{
		Package: *pkg,
		Tests:   findTests(flag.Args()),
	}

	err := testMainTmpl.Execute(buf, d)
	if err != nil {
		panic(err)
	}

	err = ioutil.WriteFile(*output, buf.Bytes(), 0666)
	if err != nil {
		panic(err)
	}
}

var testMainTmpl = template.Must(template.New("testMain").Parse(`
package main

import (
	"testing"

	pkg "{{.Package}}"
)

var t = []testing.InternalTest{
{{range .Tests}}
	{"{{.}}", pkg.{{.}}},
{{end}}
}

func matchString(pat, str string) (bool, error) {
	return true, nil
}

func main() {
	testing.Main(matchString, t, nil, nil)
}
`))
<commit_msg>Fix -test.run test selection with gotestmain<commit_after>// Copyright 2015 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"io/ioutil"
	"os"
	"strings"
	"text/template"
)

var (
	output   = flag.String("o", "", "output filename")
	pkg      = flag.String("pkg", "", "test package")
	exitCode = 0
)

type data struct {
	Package string
	Tests   []string
}

func findTests(srcs []string) (tests []string) {
	for _, src := range srcs {
		f, err := parser.ParseFile(token.NewFileSet(), src, nil, 0)
		if err != nil {
			panic(err)
		}
		for _, obj := range f.Scope.Objects {
			if obj.Kind != ast.Fun || !strings.HasPrefix(obj.Name, "Test") {
				continue
			}
			tests = append(tests, obj.Name)
		}
	}
	return
}

func main() {
	flag.Parse()

	if flag.NArg() == 0 {
		fmt.Fprintln(os.Stderr, "error: must pass at least one input")
		exitCode = 1
		return
	}

	buf := &bytes.Buffer{}

	d := data{
		Package: *pkg,
		Tests:   findTests(flag.Args()),
	}

	err := testMainTmpl.Execute(buf, d)
	if err != nil {
		panic(err)
	}

	err = ioutil.WriteFile(*output, buf.Bytes(), 0666)
	if err != nil {
		panic(err)
	}
}

var testMainTmpl = template.Must(template.New("testMain").Parse(`
package main

import (
	"regexp"
	"testing"

	pkg "{{.Package}}"
)

var t = []testing.InternalTest{
{{range .Tests}}
	{"{{.}}", pkg.{{.}}},
{{end}}
}

var matchPat string
var matchRe *regexp.Regexp

func matchString(pat, str string) (result bool, err error) {
	if matchRe == nil || matchPat != pat {
		matchPat = pat
		matchRe, err = regexp.Compile(matchPat)
		if err != nil {
			return
		}
	}
	return matchRe.MatchString(str), nil
}

func main() {
	testing.Main(matchString, t, nil, nil)
}
`))
<|endoftext|>


<commit_before>package graphql

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"

	"github.com/samsarahq/thunder/batch"
	"github.com/samsarahq/thunder/diff"
	"github.com/samsarahq/thunder/reactive"
)

const (
	MaxSubscriptions = 200
	MinRerunInterval = 5 * time.Second
)

type JSONSocket interface {
	ReadJSON(value interface{}) error
	WriteJSON(value interface{}) error
	Close() error
}

type MakeCtxFunc func(context.Context) context.Context

type GraphqlLogger interface {
	StartExecution(ctx context.Context, tags map[string]string, initial bool)
	FinishExecution(ctx context.Context, tags map[string]string, delay time.Duration)
	Error(ctx context.Context, err error, tags map[string]string)
}

type conn struct {
	writeMu sync.Mutex
	socket  JSONSocket

	schema         *Schema
	mutationSchema *Schema
	ctx            context.Context
	makeCtx        MakeCtxFunc
	logger         GraphqlLogger
	middlewares    []MiddlewareFunc

	url string

	mutateMu sync.Mutex

	mu            sync.Mutex
	subscriptions map[string]*reactive.Rerunner
}

type inEnvelope struct {
	ID         string                 `json:"id"`
	Type       string                 `json:"type"`
	Message    json.RawMessage        `json:"message"`
	Extensions map[string]interface{} `json:"extensions,omitempty"`
}

type outEnvelope struct {
	ID       string                 `json:"id,omitempty"`
	Type     string                 `json:"type"`
	Message  interface{}            `json:"message,omitempty"`
	Metadata map[string]interface{} `json:"metadata,omitempty"`
}

type subscribeMessage struct {
	Query     string                 `json:"query"`
	Variables map[string]interface{} `json:"variables"`
}

type mutateMessage struct {
	Query     string                 `json:"query"`
	Variables map[string]interface{} `json:"variables"`
}

type SanitizedError interface {
	error
	SanitizedError() string
}

type SafeError struct {
	message string
}

type ClientError SafeError

func (e ClientError) Error() string {
	return e.message
}

func (e ClientError) SanitizedError() string {
	return e.message
}

func (e SafeError) Error() string {
	return e.message
}

func (e SafeError) SanitizedError() string {
	return e.message
}

func NewClientError(format string, a ...interface{}) error {
	return ClientError{message: fmt.Sprintf(format, a...)}
}

func NewSafeError(format string, a ...interface{}) error {
	return SafeError{message: fmt.Sprintf(format, a...)}
}

func sanitizeError(err error) string {
	if sanitized, ok := err.(SanitizedError); ok {
		return sanitized.SanitizedError()
	}
	return "Internal server error"
}

func isCloseError(err error) bool {
	_, ok := err.(*websocket.CloseError)
	return ok || err == websocket.ErrCloseSent
}

func (c *conn) writeOrClose(out outEnvelope) {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()

	if err := c.socket.WriteJSON(out); err != nil {
		if !isCloseError(err) {
			c.socket.Close()
			log.Printf("socket.WriteJSON: %s\n", err)
		}
	}
}

func mustMarshalJson(v interface{}) string {
	bytes, err := json.Marshal(v)
	if err != nil {
		panic(err)
	}
	return string(bytes)
}

func (c *conn) handleSubscribe(id string, subscribe *subscribeMessage, extensions map[string]interface{}) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if _, ok := c.subscriptions[id]; ok {
		return NewSafeError("duplicate subscription")
	}

	if len(c.subscriptions)+1 > MaxSubscriptions {
		return NewSafeError("too many subscriptions")
	}

	tags := map[string]string{"url": c.url, "query": subscribe.Query, "queryVariables": mustMarshalJson(subscribe.Variables), "id": id}

	query, err := Parse(subscribe.Query, subscribe.Variables)
	if query != nil {
		tags["queryType"] = query.Kind
		tags["queryName"] = query.Name
	}
	if err != nil {
		c.logger.Error(c.ctx, err, tags)
		return err
	}
	if err := PrepareQuery(c.schema.Query, query.SelectionSet); err != nil {
		c.logger.Error(c.ctx, err, tags)
		return err
	}

	var previous interface{}

	e := Executor{}

	initial := true
	c.subscriptions[id] = reactive.NewRerunner(c.ctx, func(ctx context.Context) (interface{}, error) {
		ctx = c.makeCtx(ctx)
		ctx = batch.WithBatching(ctx)

		start := time.Now()

		c.logger.StartExecution(ctx, tags, initial)

		var middlewares []MiddlewareFunc
		middlewares = append(middlewares, c.middlewares...)
		middlewares = append(middlewares, func(input *ComputationInput, next MiddlewareNextFunc) *ComputationOutput {
			output := next(input)
			output.Current, output.Error = e.Execute(input.Ctx, c.schema.Query, nil, input.ParsedQuery)
			return output
		})

		output := runMiddlewares(middlewares, &ComputationInput{
			Ctx:         ctx,
			Id:          id,
			ParsedQuery: query,
			Previous:    previous,
			Query:       subscribe.Query,
			Variables:   subscribe.Variables,
			Extensions:  extensions,
		})
		current, err := output.Current, output.Error

		c.logger.FinishExecution(ctx, tags, time.Since(start))

		if err != nil {
			if extractPathError(err) == context.Canceled {
				go c.closeSubscription(id)
				return nil, err
			}

			if !initial {
				// If this a re-computation, tell the Rerunner to retry the computation
				// without dumping the contents of the current computation cache.
				// Note that we are swallowing the propagation of the error in this case,
				// but we still log it.
				if _, ok := err.(SanitizedError); !ok {
					extraTags := map[string]string{"retry": "true"}
					for k, v := range tags {
						extraTags[k] = v
					}
					c.logger.Error(ctx, err, extraTags)
				}

				return nil, reactive.RetrySentinelError
			}

			c.writeOrClose(outEnvelope{
				ID:       id,
				Type:     "error",
				Message:  sanitizeError(err),
				Metadata: output.Metadata,
			})
			go c.closeSubscription(id)

			if _, ok := err.(SanitizedError); !ok {
				c.logger.Error(ctx, err, tags)
			}
			return nil, err
		}

		d := diff.Diff(previous, current)
		previous = current
		initial = false

		if initial || d != nil {
			c.writeOrClose(outEnvelope{
				ID:       id,
				Type:     "update",
				Message:  d,
				Metadata: output.Metadata,
			})
		}

		return nil, nil
	}, MinRerunInterval)

	return nil
}

func (c *conn) handleMutate(id string, mutate *mutateMessage, extensions map[string]interface{}) error {
	// TODO: deduplicate code
	c.mu.Lock()
	defer c.mu.Unlock()

	tags := map[string]string{"url": c.url, "query": mutate.Query, "queryVariables": mustMarshalJson(mutate.Variables), "id": id}

	query, err := Parse(mutate.Query, mutate.Variables)
	if query != nil {
		tags["queryType"] = query.Kind
		tags["queryName"] = query.Name
	}
	if err != nil {
		c.logger.Error(c.ctx, err, tags)
		return err
	}
	if err := PrepareQuery(c.mutationSchema.Mutation, query.SelectionSet); err != nil {
		c.logger.Error(c.ctx, err, tags)
		return err
	}

	e := Executor{}
	c.subscriptions[id] = reactive.NewRerunner(c.ctx, func(ctx context.Context) (interface{}, error) {
		// Serialize all mutates for a given connection.
		c.mutateMu.Lock()
		defer c.mutateMu.Unlock()

		ctx = c.makeCtx(ctx)
		ctx = batch.WithBatching(ctx)

		start := time.Now()
		c.logger.StartExecution(ctx, tags, true)

		var middlewares []MiddlewareFunc
		middlewares = append(middlewares, c.middlewares...)
		middlewares = append(middlewares, func(input *ComputationInput, next MiddlewareNextFunc) *ComputationOutput {
			output := next(input)
			output.Current, output.Error = e.Execute(input.Ctx, c.mutationSchema.Mutation, c.mutationSchema.Mutation, query)
			return output
		})

		output := runMiddlewares(middlewares, &ComputationInput{
			Ctx:         ctx,
			Id:          id,
			ParsedQuery: query,
			Previous:    nil,
			Query:       mutate.Query,
			Variables:   mutate.Variables,
			Extensions:  extensions,
		})
		current, err := output.Current, output.Error

		c.logger.FinishExecution(ctx, tags, time.Since(start))

		if err != nil {
			c.writeOrClose(outEnvelope{
				ID:       id,
				Type:     "error",
				Message:  sanitizeError(err),
				Metadata: output.Metadata,
			})

			go c.closeSubscription(id)

			if extractPathError(err) == context.Canceled {
				return nil, err
			}

			if _, ok := err.(SanitizedError); !ok {
				c.logger.Error(ctx, err, tags)
			}
			return nil, err
		}

		c.writeOrClose(outEnvelope{
			ID:       id,
			Type:     "result",
			Message:  diff.Diff(nil, current),
			Metadata: output.Metadata,
		})

		go c.rerunSubscriptionsImmediately()

		return nil, errors.New("stop")
	}, MinRerunInterval)

	return nil
}

func (c *conn) rerunSubscriptionsImmediately() {
	c.mu.Lock()
	defer c.mu.Unlock()

	for _, runner := range c.subscriptions {
		runner.RerunImmediately()
	}
}

func (c *conn) closeSubscription(id string) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if runner, ok := c.subscriptions[id]; ok {
		runner.Stop()
		delete(c.subscriptions, id)
	}
}

func (c *conn) closeSubscriptions() {
	c.mu.Lock()
	defer c.mu.Unlock()

	for id, runner := range c.subscriptions {
		runner.Stop()
		delete(c.subscriptions, id)
	}
}

func (c *conn) handle(e *inEnvelope) error {
	switch e.Type {
	case "subscribe":
		var subscribe subscribeMessage
		if err := json.Unmarshal(e.Message, &subscribe); err != nil {
			return err
		}
		return c.handleSubscribe(e.ID, &subscribe, e.Extensions)

	case "unsubscribe":
		c.closeSubscription(e.ID)
		return nil

	case "mutate":
		var mutate mutateMessage
		if err := json.Unmarshal(e.Message, &mutate); err != nil {
			return err
		}
		return c.handleMutate(e.ID, &mutate, e.Extensions)

	case "echo":
		c.writeOrClose(outEnvelope{
			ID:       e.ID,
			Type:     "echo",
			Message:  nil,
			Metadata: nil,
		})
		return nil

	case "url":
		var url string
		if err := json.Unmarshal(e.Message, &url); err != nil {
			return err
		}
		c.url = url
		return nil

	default:
		return NewSafeError("unknown message type")
	}
}

type simpleLogger struct {
}

func (s *simpleLogger) StartExecution(ctx context.Context, tags map[string]string, initial bool) {
}
func (s *simpleLogger) FinishExecution(ctx context.Context, tags map[string]string, delay time.Duration) {
}
func (s *simpleLogger) Error(ctx context.Context, err error, tags map[string]string) {
	log.Printf("error:%v\n%s", tags, err)
}

func Handler(schema *Schema) http.Handler {
	upgrader := &websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		socket, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Printf("upgrader.Upgrade: %v", err)
			return
		}
		defer socket.Close()

		makeCtx := func(ctx context.Context) context.Context {
			return ctx
		}

		ServeJSONSocket(r.Context(), socket, schema, makeCtx, &simpleLogger{})
	})
}

func (c *conn) Use(fn MiddlewareFunc) {
	c.middlewares = append(c.middlewares, fn)
}

func ServeJSONSocket(ctx context.Context, socket JSONSocket, schema *Schema, makeCtx MakeCtxFunc, logger GraphqlLogger) {
	conn := CreateJSONSocket(ctx, socket, schema, makeCtx, logger)
	conn.ServeJSONSocket()
}

func CreateJSONSocket(ctx context.Context, socket JSONSocket, schema *Schema, makeCtx MakeCtxFunc, logger GraphqlLogger) *conn {
	return &conn{
		socket: socket,
		ctx:    ctx,

		schema:         schema,
		mutationSchema: schema,
		makeCtx:        makeCtx,
		logger:         logger,

		subscriptions: make(map[string]*reactive.Rerunner),
	}
}

func CreateJSONSocketWithMutationSchema(ctx context.Context, socket JSONSocket, schema, mutationSchema *Schema, makeCtx MakeCtxFunc, logger GraphqlLogger) *conn {
	return &conn{
		socket: socket,
		ctx:    ctx,

		schema:         schema,
		mutationSchema: mutationSchema,
		makeCtx:        makeCtx,
		logger:         logger,

		subscriptions: make(map[string]*reactive.Rerunner),
	}
}

func (c *conn) ServeJSONSocket() {
	defer c.closeSubscriptions()

	for {
		var envelope inEnvelope
		if err := c.socket.ReadJSON(&envelope); err != nil {
			if !isCloseError(err) {
				log.Println("socket.ReadJSON:", err)
			}
			return
		}

		if err := c.handle(&envelope); err != nil {
			log.Println("c.handle:", err)
			c.writeOrClose(outEnvelope{
				ID:       envelope.ID,
				Type:     "error",
				Message:  sanitizeError(err),
				Metadata: nil,
			})
		}
	}
}
<commit_msg>graphql: directly pass envelope to handlers<commit_after>package graphql

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"

	"github.com/samsarahq/thunder/batch"
	"github.com/samsarahq/thunder/diff"
	"github.com/samsarahq/thunder/reactive"
)

const (
	MaxSubscriptions = 200
	MinRerunInterval = 5 * time.Second
)

type JSONSocket interface {
	ReadJSON(value interface{}) error
	WriteJSON(value interface{}) error
	Close() error
}

type MakeCtxFunc func(context.Context) context.Context

type GraphqlLogger interface {
	StartExecution(ctx context.Context, tags map[string]string, initial bool)
	FinishExecution(ctx context.Context, tags map[string]string, delay time.Duration)
	Error(ctx context.Context, err error, tags map[string]string)
}

type conn struct {
	writeMu sync.Mutex
	socket  JSONSocket

	schema         *Schema
	mutationSchema *Schema
	ctx            context.Context
	makeCtx        MakeCtxFunc
	logger         GraphqlLogger
	middlewares    []MiddlewareFunc

	url string

	mutateMu sync.Mutex

	mu            sync.Mutex
	subscriptions map[string]*reactive.Rerunner
}

type inEnvelope struct {
	ID         string                 `json:"id"`
	Type       string                 `json:"type"`
	Message    json.RawMessage        `json:"message"`
	Extensions map[string]interface{} `json:"extensions,omitempty"`
}

type outEnvelope struct {
	ID       string                 `json:"id,omitempty"`
	Type     string                 `json:"type"`
	Message  interface{}            `json:"message,omitempty"`
	Metadata map[string]interface{} `json:"metadata,omitempty"`
}

type subscribeMessage struct {
	Query     string                 `json:"query"`
	Variables map[string]interface{} `json:"variables"`
}

type mutateMessage struct {
	Query     string                 `json:"query"`
	Variables map[string]interface{} `json:"variables"`
}

type SanitizedError interface {
	error
	SanitizedError() string
}

type SafeError struct {
	message string
}

type ClientError SafeError

func (e ClientError) Error() string {
	return e.message
}

func (e ClientError) SanitizedError() string {
	return e.message
}

func (e SafeError) Error() string {
	return e.message
}

func (e SafeError) SanitizedError() string {
	return e.message
}

func NewClientError(format string, a ...interface{}) error {
	return ClientError{message: fmt.Sprintf(format, a...)}
}

func NewSafeError(format string, a ...interface{}) error {
	return SafeError{message: fmt.Sprintf(format, a...)}
}

func sanitizeError(err error) string {
	if sanitized, ok := err.(SanitizedError); ok {
		return sanitized.SanitizedError()
	}
	return "Internal server error"
}

func isCloseError(err error) bool {
	_, ok := err.(*websocket.CloseError)
	return ok || err == websocket.ErrCloseSent
}

func (c *conn) writeOrClose(out outEnvelope) {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()

	if err := c.socket.WriteJSON(out); err != nil {
		if !isCloseError(err) {
			c.socket.Close()
			log.Printf("socket.WriteJSON: %s\n", err)
		}
	}
}

func mustMarshalJson(v interface{}) string {
	bytes, err := json.Marshal(v)
	if err != nil {
		panic(err)
	}
	return string(bytes)
}

func (c *conn) handleSubscribe(in *inEnvelope) error {
	id := in.ID
	var subscribe subscribeMessage
	if err := json.Unmarshal(in.Message, &subscribe); err != nil {
		return err
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	if _, ok := c.subscriptions[id]; ok {
		return NewSafeError("duplicate subscription")
	}

	if len(c.subscriptions)+1 > MaxSubscriptions {
		return NewSafeError("too many subscriptions")
	}

	tags := map[string]string{"url": c.url, "query": subscribe.Query, "queryVariables": mustMarshalJson(subscribe.Variables), "id": id}

	query, err := Parse(subscribe.Query, subscribe.Variables)
	if query != nil {
		tags["queryType"] = query.Kind
		tags["queryName"] = query.Name
	}
	if err != nil {
		c.logger.Error(c.ctx, err, tags)
		return err
	}
	if err := PrepareQuery(c.schema.Query, query.SelectionSet); err != nil {
		c.logger.Error(c.ctx, err, tags)
		return err
	}

	var previous interface{}

	e := Executor{}

	initial := true
	c.subscriptions[id] = reactive.NewRerunner(c.ctx, func(ctx context.Context) (interface{}, error) {
		ctx = c.makeCtx(ctx)
		ctx = batch.WithBatching(ctx)

		start := time.Now()

		c.logger.StartExecution(ctx, tags, initial)

		var middlewares []MiddlewareFunc
		middlewares = append(middlewares, c.middlewares...)
		middlewares = append(middlewares, func(input *ComputationInput, next MiddlewareNextFunc) *ComputationOutput {
			output := next(input)
			output.Current, output.Error = e.Execute(input.Ctx, c.schema.Query, nil, input.ParsedQuery)
			return output
		})

		output := runMiddlewares(middlewares, &ComputationInput{
			Ctx:         ctx,
			Id:          id,
			ParsedQuery: query,
			Previous:    previous,
			Query:       subscribe.Query,
			Variables:   subscribe.Variables,
			Extensions:  in.Extensions,
		})
		current, err := output.Current, output.Error

		c.logger.FinishExecution(ctx, tags, time.Since(start))

		if err != nil {
			if extractPathError(err) == context.Canceled {
				go c.closeSubscription(id)
				return nil, err
			}

			if !initial {
				// If this a re-computation, tell the Rerunner to retry the computation
				// without dumping the contents of the current computation cache.
				// Note that we are swallowing the propagation of the error in this case,
				// but we still log it.
				if _, ok := err.(SanitizedError); !ok {
					extraTags := map[string]string{"retry": "true"}
					for k, v := range tags {
						extraTags[k] = v
					}
					c.logger.Error(ctx, err, extraTags)
				}

				return nil, reactive.RetrySentinelError
			}

			c.writeOrClose(outEnvelope{
				ID:       id,
				Type:     "error",
				Message:  sanitizeError(err),
				Metadata: output.Metadata,
			})
			go c.closeSubscription(id)

			if _, ok := err.(SanitizedError); !ok {
				c.logger.Error(ctx, err, tags)
			}
			return nil, err
		}

		d := diff.Diff(previous, current)
		previous = current
		initial = false

		if initial || d != nil {
			c.writeOrClose(outEnvelope{
				ID:       id,
				Type:     "update",
				Message:  d,
				Metadata: output.Metadata,
			})
		}

		return nil, nil
	}, MinRerunInterval)

	return nil
}

func (c *conn) handleMutate(in *inEnvelope) error {
	// TODO: deduplicate code
	id := in.ID
	var mutate mutateMessage
	if err := json.Unmarshal(in.Message, &mutate); err != nil {
		return err
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	tags := map[string]string{"url": c.url, "query": mutate.Query, "queryVariables": mustMarshalJson(mutate.Variables), "id": id}

	query, err := Parse(mutate.Query, mutate.Variables)
	if query != nil {
		tags["queryType"] = query.Kind
		tags["queryName"] = query.Name
	}
	if err != nil {
		c.logger.Error(c.ctx, err, tags)
		return err
	}
	if err := PrepareQuery(c.mutationSchema.Mutation, query.SelectionSet); err != nil {
		c.logger.Error(c.ctx, err, tags)
		return err
	}

	e := Executor{}
	c.subscriptions[id] = reactive.NewRerunner(c.ctx, func(ctx context.Context) (interface{}, error) {
		// Serialize all mutates for a given connection.
		c.mutateMu.Lock()
		defer c.mutateMu.Unlock()

		ctx = c.makeCtx(ctx)
		ctx = batch.WithBatching(ctx)

		start := time.Now()
		c.logger.StartExecution(ctx, tags, true)

		var middlewares []MiddlewareFunc
		middlewares = append(middlewares, c.middlewares...)
		middlewares = append(middlewares, func(input *ComputationInput, next MiddlewareNextFunc) *ComputationOutput {
			output := next(input)
			output.Current, output.Error = e.Execute(input.Ctx, c.mutationSchema.Mutation, c.mutationSchema.Mutation, query)
			return output
		})

		output := runMiddlewares(middlewares, &ComputationInput{
			Ctx:         ctx,
			Id:          id,
			ParsedQuery: query,
			Previous:    nil,
			Query:       mutate.Query,
			Variables:   mutate.Variables,
			Extensions:  in.Extensions,
		})
		current, err := output.Current, output.Error

		c.logger.FinishExecution(ctx, tags, time.Since(start))

		if err != nil {
			c.writeOrClose(outEnvelope{
				ID:       id,
				Type:     "error",
				Message:  sanitizeError(err),
				Metadata: output.Metadata,
			})

			go c.closeSubscription(id)

			if extractPathError(err) == context.Canceled {
				return nil, err
			}

			if _, ok := err.(SanitizedError); !ok {
				c.logger.Error(ctx, err, tags)
			}
			return nil, err
		}

		c.writeOrClose(outEnvelope{
			ID:       id,
			Type:     "result",
			Message:  diff.Diff(nil, current),
			Metadata: output.Metadata,
		})

		go c.rerunSubscriptionsImmediately()

		return nil, errors.New("stop")
	}, MinRerunInterval)

	return nil
}

func (c *conn) rerunSubscriptionsImmediately() {
	c.mu.Lock()
	defer c.mu.Unlock()

	for _, runner := range c.subscriptions {
		runner.RerunImmediately()
	}
}

func (c *conn) closeSubscription(id string) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if runner, ok := c.subscriptions[id]; ok {
		runner.Stop()
		delete(c.subscriptions, id)
	}
}

func (c *conn) closeSubscriptions() {
	c.mu.Lock()
	defer c.mu.Unlock()

	for id, runner := range c.subscriptions {
		runner.Stop()
		delete(c.subscriptions, id)
	}
}

func (c *conn) handle(e *inEnvelope) error {
	switch e.Type {
	case "subscribe":
		return c.handleSubscribe(e)

	case "unsubscribe":
		c.closeSubscription(e.ID)
		return nil

	case "mutate":
		return c.handleMutate(e)

	case "echo":
		c.writeOrClose(outEnvelope{
			ID:       e.ID,
			Type:     "echo",
			Message:  nil,
			Metadata: nil,
		})
		return nil

	case "url":
		var url string
		if err := json.Unmarshal(e.Message, &url); err != nil {
			return err
		}
		c.url = url
		return nil

	default:
		return NewSafeError("unknown message type")
	}
}

type simpleLogger struct {
}

func (s *simpleLogger) StartExecution(ctx context.Context, tags map[string]string, initial bool) {
}
func (s *simpleLogger) FinishExecution(ctx context.Context, tags map[string]string, delay time.Duration) {
}
func (s *simpleLogger) Error(ctx context.Context, err error, tags map[string]string) {
	log.Printf("error:%v\n%s", tags, err)
}

func Handler(schema *Schema) http.Handler {
	upgrader := &websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		socket, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Printf("upgrader.Upgrade: %v", err)
			return
		}
		defer socket.Close()

		makeCtx := func(ctx context.Context) context.Context {
			return ctx
		}

		ServeJSONSocket(r.Context(), socket, schema, makeCtx, &simpleLogger{})
	})
}

func (c *conn) Use(fn MiddlewareFunc) {
	c.middlewares = append(c.middlewares, fn)
}

func ServeJSONSocket(ctx context.Context, socket JSONSocket, schema *Schema, makeCtx MakeCtxFunc, logger GraphqlLogger) {
	conn := CreateJSONSocket(ctx, socket, schema, makeCtx, logger)
	conn.ServeJSONSocket()
}

func CreateJSONSocket(ctx context.Context, socket JSONSocket, schema *Schema, makeCtx MakeCtxFunc, logger GraphqlLogger) *conn {
	return &conn{
		socket: socket,
		ctx:    ctx,

		schema:         schema,
		mutationSchema: schema,
		makeCtx:        makeCtx,
		logger:         logger,

		subscriptions: make(map[string]*reactive.Rerunner),
	}
}

func CreateJSONSocketWithMutationSchema(ctx context.Context, socket JSONSocket, schema, mutationSchema *Schema, makeCtx MakeCtxFunc, logger GraphqlLogger) *conn {
	return &conn{
		socket: socket,
		ctx:    ctx,

		schema:         schema,
		mutationSchema: mutationSchema,
		makeCtx:        makeCtx,
		logger:         logger,

		subscriptions: make(map[string]*reactive.Rerunner),
	}
}

func (c *conn) ServeJSONSocket() {
	defer c.closeSubscriptions()

	for {
		var envelope inEnvelope
		if err := c.socket.ReadJSON(&envelope); err != nil {
			if !isCloseError(err) {
				log.Println("socket.ReadJSON:", err)
			}
			return
		}

		if err := c.handle(&envelope); err != nil {
			log.Println("c.handle:", err)
			c.writeOrClose(outEnvelope{
				ID:       envelope.ID,
				Type:     "error",
				Message:  sanitizeError(err),
				Metadata: nil,
			})
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"time"

	"github.com/robfig/cron/v3"
	"github.com/sirupsen/logrus"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
)

type reaper struct {
	clientSet *kubernetes.Clientset
	options   options
}

func newReaper() reaper {
	config, err := rest.InClusterConfig()
	if err != nil {
		logrus.WithError(err).Panic("error getting in cluster kubernetes config")
		panic(err)
	}
	clientSet, err := kubernetes.NewForConfig(config)
	if err != nil {
		logrus.WithError(err).Panic("unable to get client set for in cluster kubernetes config")
		panic(err)
	}
	if clientSet == nil {
		message := "kubernetes client set cannot be nil"
		logrus.Panic(message)
		panic(message)
	}
	options, err := loadOptions()
	if err != nil {
		logrus.WithError(err).Panic("error loading options")
		panic(err)
	}
	return reaper{
		clientSet: clientSet,
		options:   options,
	}
}

func (reaper reaper) getPods() *v1.PodList {
	coreClient := reaper.clientSet.CoreV1()
	pods := coreClient.Pods(reaper.options.namespace)
	listOptions := metav1.ListOptions{}
	if reaper.options.labelExclusion != nil || reaper.options.labelRequirement != nil {
		selector := labels.NewSelector()
		if reaper.options.labelExclusion != nil {
			selector = selector.Add(*reaper.options.labelExclusion)
		}
		if reaper.options.labelRequirement != nil {
			selector = selector.Add(*reaper.options.labelRequirement)
		}
		listOptions.LabelSelector = selector.String()
	}
	podList, err := pods.List(listOptions)
	if err != nil {
		logrus.WithError(err).Panic("unable to get pods from the cluster")
		panic(err)
	}
	filteredList := []v1.Pod{}
	if reaper.options.annotationRequirement != nil {
		for _, pod := range podList.Items {
			// convert the pod's annotations to an equivalent label selector
			selector := labels.Set(pod.Annotations)
			// include pod if its annotations match the selector
			if reaper.options.annotationRequirement.Matches(selector) {
				filteredList = append(filteredList, pod)
			}
		}
		podList.Items = filteredList
	}
	return podList
}

func (reaper reaper) reapPod(pod v1.Pod, reasons []string) {
	deleteOptions := &metav1.DeleteOptions{
		GracePeriodSeconds: reaper.options.gracePeriod,
	}
	if reaper.options.dryRun {
		logrus.WithFields(logrus.Fields{
			"pod":     pod.Name,
			"reasons": reasons,
		}).Info("pod would be reaped but pod-reaper is in dry-run mode")
	} else {
		logrus.WithFields(logrus.Fields{
			"pod":     pod.Name,
			"reasons": reasons,
		}).Info("reaping pod")
		err := reaper.clientSet.CoreV1().Pods(pod.Namespace).Delete(pod.Name, deleteOptions)
		if err != nil {
			// log the error, but continue on
			logrus.WithFields(logrus.Fields{
				"pod": pod.Name,
			}).WithError(err).Warn("unable to delete pod", err)
		}
	}
}

func (reaper reaper) scytheCycle() {
	logrus.Debug("starting reap cycle")
	pods := reaper.getPods()
	for _, pod := range pods.Items {
		shouldReap, reasons := reaper.options.rules.ShouldReap(pod)
		if shouldReap {
			reaper.reapPod(pod, reasons)
		}
	}
}

func cronWithOptionalSeconds() *cron.Cron {
	return cron.New(
		cron.WithParser(
			cron.NewParser(
				// include optional seconds
				cron.SecondOptional | cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor)))
}

func (reaper reaper) harvest() {
	runForever := reaper.options.runDuration == 0
	schedule := cronWithOptionalSeconds()
	_, err := schedule.AddFunc(reaper.options.schedule, func() {
		reaper.scytheCycle()
	})

	if err != nil {
		logrus.WithError(err).Panic("unable to create cron schedule: " + reaper.options.schedule)
		panic(err)
	}

	schedule.Start()

	if runForever {
		select {} // should only fail if no routine can make progress
	} else {
		time.Sleep(reaper.options.runDuration)
		schedule.Stop()
	}
}
<commit_msg>fix list scoping<commit_after>package main

import (
	"time"

	"github.com/robfig/cron/v3"
	"github.com/sirupsen/logrus"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
)

type reaper struct {
	clientSet *kubernetes.Clientset
	options   options
}

func newReaper() reaper {
	config, err := rest.InClusterConfig()
	if err != nil {
		logrus.WithError(err).Panic("error getting in cluster kubernetes config")
		panic(err)
	}
	clientSet, err := kubernetes.NewForConfig(config)
	if err != nil {
		logrus.WithError(err).Panic("unable to get client set for in cluster kubernetes config")
		panic(err)
	}
	if clientSet == nil {
		message := "kubernetes client set cannot be nil"
		logrus.Panic(message)
		panic(message)
	}
	options, err := loadOptions()
	if err != nil {
		logrus.WithError(err).Panic("error loading options")
		panic(err)
	}
	return reaper{
		clientSet: clientSet,
		options:   options,
	}
}

func (reaper reaper) getPods() *v1.PodList {
	coreClient := reaper.clientSet.CoreV1()
	pods := coreClient.Pods(reaper.options.namespace)
	listOptions := metav1.ListOptions{}
	if reaper.options.labelExclusion != nil || reaper.options.labelRequirement != nil {
		selector := labels.NewSelector()
		if reaper.options.labelExclusion != nil {
			selector = selector.Add(*reaper.options.labelExclusion)
		}
		if reaper.options.labelRequirement != nil {
			selector = selector.Add(*reaper.options.labelRequirement)
		}
		listOptions.LabelSelector = selector.String()
	}
	podList, err := pods.List(listOptions)
	if err != nil {
		logrus.WithError(err).Panic("unable to get pods from the cluster")
		panic(err)
	}
	if reaper.options.annotationRequirement != nil {
		filteredList := []v1.Pod{}
		for _, pod := range podList.Items {
			// convert the pod's annotations to an equivalent label selector
			selector := labels.Set(pod.Annotations)
			// include pod if its annotations match the selector
			if reaper.options.annotationRequirement.Matches(selector) {
				filteredList = append(filteredList, pod)
			}
		}
		podList.Items = filteredList
	}
	return podList
}

func (reaper reaper) reapPod(pod v1.Pod, reasons []string) {
	deleteOptions := &metav1.DeleteOptions{
		GracePeriodSeconds: reaper.options.gracePeriod,
	}
	if reaper.options.dryRun {
		logrus.WithFields(logrus.Fields{
			"pod":     pod.Name,
			"reasons": reasons,
		}).Info("pod would be reaped but pod-reaper is in dry-run mode")
	} else {
		logrus.WithFields(logrus.Fields{
			"pod":     pod.Name,
			"reasons": reasons,
		}).Info("reaping pod")
		err := reaper.clientSet.CoreV1().Pods(pod.Namespace).Delete(pod.Name, deleteOptions)
		if err != nil {
			// log the error, but continue on
			logrus.WithFields(logrus.Fields{
				"pod": pod.Name,
			}).WithError(err).Warn("unable to delete pod", err)
		}
	}
}

func (reaper reaper) scytheCycle() {
	logrus.Debug("starting reap cycle")
	pods := reaper.getPods()
	for _, pod := range pods.Items {
		shouldReap, reasons := reaper.options.rules.ShouldReap(pod)
		if shouldReap {
			reaper.reapPod(pod, reasons)
		}
	}
}

func cronWithOptionalSeconds() *cron.Cron {
	return cron.New(
		cron.WithParser(
			cron.NewParser(
				// include optional seconds
				cron.SecondOptional | cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor)))
}

func (reaper reaper) harvest() {
	runForever := reaper.options.runDuration == 0
	schedule := cronWithOptionalSeconds()
	_, err := schedule.AddFunc(reaper.options.schedule, func() {
		reaper.scytheCycle()
	})

	if err != nil {
		logrus.WithError(err).Panic("unable to create cron schedule: " + reaper.options.schedule)
		panic(err)
	}

	schedule.Start()

	if runForever {
		select {} // should only fail if no routine can make progress
	} else {
		time.Sleep(reaper.options.runDuration)
		schedule.Stop()
	}
}
<|endoftext|>