<commit_before>// The state package enables reading, observing, and changing
// the state stored in MongoDB of a whole environment
// managed by juju.
package mstate

import (
	"fmt"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"labix.org/v2/mgo/txn"
	"launchpad.net/juju-core/charm"
	"launchpad.net/juju-core/trivial"
	"launchpad.net/juju-core/version"
	"net/url"
)

const (
	environConfigKey = "environ"
)

// Tools describes a particular set of juju tools and where to find them.
type Tools struct {
	version.Binary
	URL string
}

type Life int8

const (
	Alive Life = iota
	Dying
	Dead
	nLife
)

var lifeStrings = [nLife]string{
	Alive: "alive",
	Dying: "dying",
	Dead:  "dead",
}

func (l Life) String() string {
	return lifeStrings[l]
}

// State represents the state of an environment
// managed by juju.
type State struct {
	db        *mgo.Database
	cfgnodes  *mgo.Collection
	charms    *mgo.Collection
	machines  *mgo.Collection
	relations *mgo.Collection
	services  *mgo.Collection
	txns      *mgo.Collection
	units     *mgo.Collection
	runner    *txn.Runner
}

// EnvironConfig returns the current configuration of the environment.
func (s *State) EnvironConfig() (*ConfigNode, error) {
	return readConfigNode(s, environConfigKey)
}

// AddMachine creates a new machine state.
func (s *State) AddMachine() (m *Machine, err error) {
	defer trivial.ErrorContextf(&err, "cannot add a new machine")
	id, err := s.sequence("machine")
	if err != nil {
		return nil, err
	}
	mdoc := machineDoc{
		Id:   id,
		Life: Alive,
	}
	op := []txn.Operation{{
		Collection: s.machines.Name,
		DocId:      id,
		Assert:     txn.DocMissing,
		Insert:     mdoc,
	}}
	err = s.runner.Run(op, "", nil)
	if err != nil {
		return nil, err
	}
	return newMachine(s, &mdoc), nil
}

// RemoveMachine removes the machine with the the given id.
func (s *State) RemoveMachine(id int) error {
	op := []txn.Operation{{
		Collection: s.machines.Name,
		DocId:      id,
		Assert:     bson.D{{"life", Alive}},
		Change:     bson.D{{"$set", bson.D{{"life", Dying}}}},
	}}
	err := s.runner.Run(op, "", nil)
	if err != nil {
		return fmt.Errorf("cannot remove machine %d: %v", id, err)
	}
	return nil
}

// AllMachines returns all machines in the environment.
func (s *State) AllMachines() (machines []*Machine, err error) {
	mdocs := []machineDoc{}
	sel := bson.D{{"life", Alive}}
	err = s.machines.Find(sel).Select(bson.D{{"_id", 1}}).All(&mdocs)
	if err != nil {
		return nil, fmt.Errorf("cannot get all machines: %v", err)
	}
	for _, doc := range mdocs {
		machines = append(machines, newMachine(s, &doc))
	}
	return
}

// Machine returns the machine with the given id.
func (s *State) Machine(id int) (*Machine, error) {
	mdoc := &machineDoc{}
	sel := bson.D{{"_id", id}, {"life", Alive}}
	err := s.machines.Find(sel).One(mdoc)
	if err != nil {
		return nil, fmt.Errorf("cannot get machine %d: %v", id, err)
	}
	return newMachine(s, mdoc), nil
}

// AddCharm adds the ch charm with curl to the state.  bundleUrl must be
// set to a URL where the bundle for ch may be downloaded from.
// On success the newly added charm state is returned.
func (s *State) AddCharm(ch charm.Charm, curl *charm.URL, bundleURL *url.URL, bundleSha256 string) (stch *Charm, err error) {
	cdoc := &charmDoc{
		URL:          curl,
		Meta:         ch.Meta(),
		Config:       ch.Config(),
		BundleURL:    bundleURL,
		BundleSha256: bundleSha256,
	}
	op := []txn.Operation{{
		Collection: s.charms.Name,
		DocId: curl,
		Assert: txn.DocMissing,
		Insert: cdoc,
	}}
	err = s.runner.Run(op, "", nil)
	if err != nil {
		return nil, fmt.Errorf("cannot add charm %q: %v", curl, err)
	}
	return newCharm(s, cdoc)
}

// Charm returns the charm with the given URL.
func (s *State) Charm(curl *charm.URL) (*Charm, error) {
	cdoc := &charmDoc{}
	err := s.charms.Find(bson.D{{"_id", curl}}).One(cdoc)
	if err != nil {
		return nil, fmt.Errorf("cannot get charm %q: %v", curl, err)
	}

	return newCharm(s, cdoc)
}

// AddService creates a new service state with the given unique name
// and the charm state.
func (s *State) AddService(name string, ch *Charm) (service *Service, err error) {
	sdoc := serviceDoc{
		Name:     name,
		CharmURL: ch.URL(),
		Life:     Alive,
	}
	op := []txn.Operation{{
		Collection: s.services.Name,
		DocId: name,
		Assert: txn.DocMissing,
		Insert: sdoc,
	}}
	err = s.runner.Run(op, "", nil)
	if err != nil {
		return nil, fmt.Errorf("cannot add service %q:", name, err)
	}
	return &Service{st: s, doc: sdoc}, nil
}

// RemoveService removes a service from the state. It will also remove all
// its units and break any of its existing relations.
func (s *State) RemoveService(svc *Service) (err error) {
	defer trivial.ErrorContextf(&err, "cannot remove service %q", svc)

	sel := bson.D{{"_id", svc.doc.Name}, {"life", Alive}}
	change := bson.D{{"$set", bson.D{{"life", Dying}}}}
	err = s.services.Update(sel, change)
	if err != nil {
		return err
	}

	sel = bson.D{{"service", svc.doc.Name}}
	change = bson.D{{"$set", bson.D{{"life", Dying}}}}
	_, err = s.units.UpdateAll(sel, change)
	return err
}

// Service returns a service state by name.
func (s *State) Service(name string) (service *Service, err error) {
	sdoc := serviceDoc{}
	sel := bson.D{{"_id", name}, {"life", Alive}}
	err = s.services.Find(sel).One(&sdoc)
	if err != nil {
		return nil, fmt.Errorf("cannot get service %q: %v", name, err)
	}
	return &Service{st: s, doc: sdoc}, nil
}

// AllServices returns all deployed services in the environment.
func (s *State) AllServices() (services []*Service, err error) {
	sdocs := []serviceDoc{}
	err = s.services.Find(bson.D{{"life", Alive}}).All(&sdocs)
	if err != nil {
		return nil, fmt.Errorf("cannot get all services")
	}
	for _, v := range sdocs {
		services = append(services, &Service{st: s, doc: v})
	}
	return services, nil
}

// AddRelation creates a new relation with the given endpoints.
func (s *State) AddRelation(endpoints ...RelationEndpoint) (r *Relation, err error) {
	defer trivial.ErrorContextf(&err, "cannot add relation %q", relationKey(endpoints))
	switch len(endpoints) {
	case 1:
		if endpoints[0].RelationRole != RolePeer {
			return nil, fmt.Errorf("single endpoint must be a peer relation")
		}
	case 2:
		if !endpoints[0].CanRelateTo(&endpoints[1]) {
			return nil, fmt.Errorf("endpoints do not relate")
		}
	default:
		return nil, fmt.Errorf("cannot relate %d endpoints", len(endpoints))
	}

	var scope charm.RelationScope
	for _, v := range endpoints {
		if v.RelationScope == charm.ScopeContainer {
			scope = charm.ScopeContainer
		}
		// BUG(aram): potential race in the time between getting the service
		// to validate the endpoint and actually writting the relation
		// into MongoDB; the service might have disappeared.
		_, err = s.Service(v.ServiceName)
		if err != nil {
			return nil, err
		}
	}
	if scope == charm.ScopeContainer {
		for i := range endpoints {
			endpoints[i].RelationScope = scope
		}
	}
	id, err := s.sequence("relation")
	if err != nil {
		return nil, err
	}
	doc := relationDoc{
		Id:        id,
		Key:       relationKey(endpoints),
		Endpoints: endpoints,
		Life:      Alive,
	}
	err = s.relations.Insert(doc)
	if err != nil {
		return nil, err
	}
	return newRelation(s, &doc), nil
}

// Relation returns the existing relation with the given endpoints.
func (s *State) Relation(endpoints ...RelationEndpoint) (r *Relation, err error) {
	defer trivial.ErrorContextf(&err, "cannot get relation %q", relationKey(endpoints))

	doc := relationDoc{}
	err = s.relations.Find(bson.D{{"key", relationKey(endpoints)}}).One(&doc)
	if err != nil {
		return nil, err
	}
	return newRelation(s, &doc), nil
}

// RemoveRelation removes the supplied relation.
func (s *State) RemoveRelation(r *Relation) (err error) {
	defer trivial.ErrorContextf(&err, "cannot remove relation %q", r.doc.Key)

	if r.doc.Life != Dead {
		panic(fmt.Errorf("relation %q is not dead", r))
	}
	sel := bson.D{
		{"_id", r.doc.Id},
		{"life", Dead},
	}
	err = s.relations.Remove(sel)
	if err != nil {
		return err
	}
	return nil
}

// Unit returns a unit by name.
func (s *State) Unit(name string) (*Unit, error) {
	doc := unitDoc{}
	err := s.units.FindId(name).One(&doc)
	if err != nil {
		return nil, fmt.Errorf("cannot get unit %q: %v", name, err)
	}
	return newUnit(s, &doc), nil
}
<commit_msg>mstate: use txn for RemoveService (but not for setting associated units to dying)<commit_after>// The state package enables reading, observing, and changing
// the state stored in MongoDB of a whole environment
// managed by juju.
package mstate

import (
	"fmt"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"labix.org/v2/mgo/txn"
	"launchpad.net/juju-core/charm"
	"launchpad.net/juju-core/trivial"
	"launchpad.net/juju-core/version"
	"net/url"
)

const (
	environConfigKey = "environ"
)

// Tools describes a particular set of juju tools and where to find them.
type Tools struct {
	version.Binary
	URL string
}

type Life int8

const (
	Alive Life = iota
	Dying
	Dead
	nLife
)

var lifeStrings = [nLife]string{
	Alive: "alive",
	Dying: "dying",
	Dead:  "dead",
}

func (l Life) String() string {
	return lifeStrings[l]
}

// State represents the state of an environment
// managed by juju.
type State struct {
	db        *mgo.Database
	cfgnodes  *mgo.Collection
	charms    *mgo.Collection
	machines  *mgo.Collection
	relations *mgo.Collection
	services  *mgo.Collection
	txns      *mgo.Collection
	units     *mgo.Collection
	runner    *txn.Runner
}

// EnvironConfig returns the current configuration of the environment.
func (s *State) EnvironConfig() (*ConfigNode, error) {
	return readConfigNode(s, environConfigKey)
}

// AddMachine creates a new machine state.
func (s *State) AddMachine() (m *Machine, err error) {
	defer trivial.ErrorContextf(&err, "cannot add a new machine")
	id, err := s.sequence("machine")
	if err != nil {
		return nil, err
	}
	mdoc := machineDoc{
		Id:   id,
		Life: Alive,
	}
	op := []txn.Operation{{
		Collection: s.machines.Name,
		DocId:      id,
		Assert:     txn.DocMissing,
		Insert:     mdoc,
	}}
	err = s.runner.Run(op, "", nil)
	if err != nil {
		return nil, err
	}
	return newMachine(s, &mdoc), nil
}

// RemoveMachine removes the machine with the the given id.
func (s *State) RemoveMachine(id int) error {
	op := []txn.Operation{{
		Collection: s.machines.Name,
		DocId:      id,
		Assert:     bson.D{{"life", Alive}},
		Change:     bson.D{{"$set", bson.D{{"life", Dying}}}},
	}}
	err := s.runner.Run(op, "", nil)
	if err != nil {
		return fmt.Errorf("cannot remove machine %d: %v", id, err)
	}
	return nil
}

// AllMachines returns all machines in the environment.
func (s *State) AllMachines() (machines []*Machine, err error) {
	mdocs := []machineDoc{}
	sel := bson.D{{"life", Alive}}
	err = s.machines.Find(sel).Select(bson.D{{"_id", 1}}).All(&mdocs)
	if err != nil {
		return nil, fmt.Errorf("cannot get all machines: %v", err)
	}
	for _, doc := range mdocs {
		machines = append(machines, newMachine(s, &doc))
	}
	return
}

// Machine returns the machine with the given id.
func (s *State) Machine(id int) (*Machine, error) {
	mdoc := &machineDoc{}
	sel := bson.D{{"_id", id}, {"life", Alive}}
	err := s.machines.Find(sel).One(mdoc)
	if err != nil {
		return nil, fmt.Errorf("cannot get machine %d: %v", id, err)
	}
	return newMachine(s, mdoc), nil
}

// AddCharm adds the ch charm with curl to the state.  bundleUrl must be
// set to a URL where the bundle for ch may be downloaded from.
// On success the newly added charm state is returned.
func (s *State) AddCharm(ch charm.Charm, curl *charm.URL, bundleURL *url.URL, bundleSha256 string) (stch *Charm, err error) {
	cdoc := &charmDoc{
		URL:          curl,
		Meta:         ch.Meta(),
		Config:       ch.Config(),
		BundleURL:    bundleURL,
		BundleSha256: bundleSha256,
	}
	op := []txn.Operation{{
		Collection: s.charms.Name,
		DocId:      curl,
		Assert:     txn.DocMissing,
		Insert:     cdoc,
	}}
	err = s.runner.Run(op, "", nil)
	if err != nil {
		return nil, fmt.Errorf("cannot add charm %q: %v", curl, err)
	}
	return newCharm(s, cdoc)
}

// Charm returns the charm with the given URL.
func (s *State) Charm(curl *charm.URL) (*Charm, error) {
	cdoc := &charmDoc{}
	err := s.charms.Find(bson.D{{"_id", curl}}).One(cdoc)
	if err != nil {
		return nil, fmt.Errorf("cannot get charm %q: %v", curl, err)
	}

	return newCharm(s, cdoc)
}

// AddService creates a new service state with the given unique name
// and the charm state.
func (s *State) AddService(name string, ch *Charm) (service *Service, err error) {
	sdoc := serviceDoc{
		Name:     name,
		CharmURL: ch.URL(),
		Life:     Alive,
	}
	op := []txn.Operation{{
		Collection: s.services.Name,
		DocId:      name,
		Assert:     txn.DocMissing,
		Insert:     sdoc,
	}}
	err = s.runner.Run(op, "", nil)
	if err != nil {
		return nil, fmt.Errorf("cannot add service %q:", name, err)
	}
	return &Service{st: s, doc: sdoc}, nil
}

// RemoveService removes a service from the state. It will also remove all
// its units and break any of its existing relations.
func (s *State) RemoveService(svc *Service) (err error) {
	defer trivial.ErrorContextf(&err, "cannot remove service %q", svc)

	op := []txn.Operation{{
		Collection: s.services.Name,
		DocId:      svc.doc.Name,
		Assert:     bson.D{{"_id", svc.doc.Name}, {"life", Alive}},
		Change:     bson.D{{"$set", bson.D{{"life", Dying}}}},
	}}
	err = s.runner.Run(op, "", nil)
	if err != nil {
		return err
	}

	// BUG(aram): this should be in the above transaction, but we
	// are missing UpdateAll for now.
	sel := bson.D{{"service", svc.doc.Name}}
	change := bson.D{{"$set", bson.D{{"life", Dying}}}}
	_, err = s.units.UpdateAll(sel, change)
	return err
}

// Service returns a service state by name.
func (s *State) Service(name string) (service *Service, err error) {
	sdoc := serviceDoc{}
	sel := bson.D{{"_id", name}, {"life", Alive}}
	err = s.services.Find(sel).One(&sdoc)
	if err != nil {
		return nil, fmt.Errorf("cannot get service %q: %v", name, err)
	}
	return &Service{st: s, doc: sdoc}, nil
}

// AllServices returns all deployed services in the environment.
func (s *State) AllServices() (services []*Service, err error) {
	sdocs := []serviceDoc{}
	err = s.services.Find(bson.D{{"life", Alive}}).All(&sdocs)
	if err != nil {
		return nil, fmt.Errorf("cannot get all services")
	}
	for _, v := range sdocs {
		services = append(services, &Service{st: s, doc: v})
	}
	return services, nil
}

// AddRelation creates a new relation with the given endpoints.
func (s *State) AddRelation(endpoints ...RelationEndpoint) (r *Relation, err error) {
	defer trivial.ErrorContextf(&err, "cannot add relation %q", relationKey(endpoints))
	switch len(endpoints) {
	case 1:
		if endpoints[0].RelationRole != RolePeer {
			return nil, fmt.Errorf("single endpoint must be a peer relation")
		}
	case 2:
		if !endpoints[0].CanRelateTo(&endpoints[1]) {
			return nil, fmt.Errorf("endpoints do not relate")
		}
	default:
		return nil, fmt.Errorf("cannot relate %d endpoints", len(endpoints))
	}

	var scope charm.RelationScope
	for _, v := range endpoints {
		if v.RelationScope == charm.ScopeContainer {
			scope = charm.ScopeContainer
		}
		// BUG(aram): potential race in the time between getting the service
		// to validate the endpoint and actually writting the relation
		// into MongoDB; the service might have disappeared.
		_, err = s.Service(v.ServiceName)
		if err != nil {
			return nil, err
		}
	}
	if scope == charm.ScopeContainer {
		for i := range endpoints {
			endpoints[i].RelationScope = scope
		}
	}
	id, err := s.sequence("relation")
	if err != nil {
		return nil, err
	}
	doc := relationDoc{
		Id:        id,
		Key:       relationKey(endpoints),
		Endpoints: endpoints,
		Life:      Alive,
	}
	err = s.relations.Insert(doc)
	if err != nil {
		return nil, err
	}
	return newRelation(s, &doc), nil
}

// Relation returns the existing relation with the given endpoints.
func (s *State) Relation(endpoints ...RelationEndpoint) (r *Relation, err error) {
	defer trivial.ErrorContextf(&err, "cannot get relation %q", relationKey(endpoints))

	doc := relationDoc{}
	err = s.relations.Find(bson.D{{"key", relationKey(endpoints)}}).One(&doc)
	if err != nil {
		return nil, err
	}
	return newRelation(s, &doc), nil
}

// RemoveRelation removes the supplied relation.
func (s *State) RemoveRelation(r *Relation) (err error) {
	defer trivial.ErrorContextf(&err, "cannot remove relation %q", r.doc.Key)

	if r.doc.Life != Dead {
		panic(fmt.Errorf("relation %q is not dead", r))
	}
	sel := bson.D{
		{"_id", r.doc.Id},
		{"life", Dead},
	}
	err = s.relations.Remove(sel)
	if err != nil {
		return err
	}
	return nil
}

// Unit returns a unit by name.
func (s *State) Unit(name string) (*Unit, error) {
	doc := unitDoc{}
	err := s.units.FindId(name).One(&doc)
	if err != nil {
		return nil, fmt.Errorf("cannot get unit %q: %v", name, err)
	}
	return newUnit(s, &doc), nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"encoding/binary"
	"log"
	"net"
	"os"
	"time"
	"github.com/miekg/dns"
	"github.com/miekg/pcap"
)

var (
	// We store the configuration of our local resolver in a global
	// variable for convenience.
	resolv_conf	*dns.ClientConfig

	// Use a global debug flag.
	debug		bool
)

// If we were passed name server addresses, parse them with this function.
func parse_root_server_addresses(addrs []string)(map[[4]byte]bool, map[[16]byte]bool) {
	if debug {
		fmt.Fprintf(os.Stderr, "parse_root_server_addresses()\n")
		fmt.Fprintf(os.Stderr, "addrs:%s\n", addrs)
	}
	root_addresses4 := make(map[[4]byte]bool)
	root_addresses6 := make(map[[16]byte]bool)
	for _, addr := range addrs {
		ip := net.ParseIP(addr)
		if ip.To4() == nil {
			// IPv6 address
			var aaaa [16]byte
			copy(aaaa[:], ip[0:16])
			root_addresses6[aaaa] = true
		} else {
			// IPv4 address
			var a [4]byte
			copy(a[:], ip.To4()[0:4])
			root_addresses4[a] = true
		}
	}
	return root_addresses4, root_addresses6
}

// We have a goroutine to act as a stub resolver, and use this
// structure to send the question in and get the results out.
type stub_resolve_info struct {
	ownername	string
	rtype		uint16
	answer		*dns.Msg
}

// A goroutine which performs stub lookups from a queue, writing
// the results to another queue.
func stub_resolve(questions <-chan stub_resolve_info,
                  results chan<- stub_resolve_info) {
	// make a client for our lookups
	dnsClient := new(dns.Client)
	dnsClient.Net = "tcp"
	// read each question on our channel
	for question := range questions {
		// build our answer
		var result stub_resolve_info = question
		result.answer = nil
		// make a DNS query based on our question
		query := new(dns.Msg)
		query.RecursionDesired = true
		query.SetQuestion(question.ownername, question.rtype)
		// check each resolver in turn
		for _, server := range resolv_conf.Servers {
			resolver := server + ":53"
			r , _, err := dnsClient.Exchange(query, resolver)
			// if we got an answer, use that and stop trying
			if (err == nil) && (r != nil) && (r.Rcode == dns.RcodeSuccess) {
				result.answer = r
				break
			}
		}
		// send back our answer (might be nil)
		results <- result
	}
}

// Lookup all of the IP addresses associated with the root name servers.
// Return two maps based on the results found, which have the keys of
// the binary values of the IPv4 and IPv6 addresses. (It's a bit clumsy,
// but it allows us to do quick and easy lookups of the addresses in the
// pcap later.)
func lookup_root_server_addresses() (map[[4]byte]bool, map[[16]byte]bool) {
	if debug {
		fmt.Fprintf(os.Stderr, "lookup_root_server_addresses()\n")
	}
	// look up the NS of the IANA root
	root_client := new(dns.Client)
	root_client.Net = "tcp"
	ns_query := new(dns.Msg)
	ns_query.SetQuestion(".", dns.TypeNS)
	// TODO: avoid hard-coding a particular root server here
	ns_response , _, err := root_client.Exchange(ns_query,
                                                     "k.root-servers.net:53")
	if err != nil {
		log.Fatal("Error looking up root name servers")
	}
	var root_servers []string
	for _, root_server := range ns_response.Answer {
		switch root_server.(type) {
		case *dns.NS:
			ns := root_server.(*dns.NS).Ns
			root_servers = append(root_servers, ns)
		}
	}
	// look up the addresses of the root servers
	questions := make(chan stub_resolve_info, 16)
	results := make(chan stub_resolve_info, len(root_servers)*2)
	for i := 0; i<8; i++ {
		go stub_resolve(questions, results)
	}
	for _, ns := range root_servers {
		info := new(stub_resolve_info)
		info.ownername = ns
		info.rtype = dns.TypeAAAA
		questions <- *info
		info = new(stub_resolve_info)
		info.ownername = ns
		info.rtype = dns.TypeA
		questions <- *info
	}
	root_addresses4 := make(map[[4]byte]bool)
	root_addresses6 := make(map[[16]byte]bool)
	for i := 0; i<len(root_servers)*2; i++ {
		response := <-results
		if response.answer == nil {
			log.Fatalf("Error looking up root server %s",
			           response.ownername)
		}
		for _, root_address := range response.answer.Answer {
			switch root_address.(type) {
			case *dns.AAAA:
				aaaa_s := root_address.(*dns.AAAA).AAAA.String()
				var aaaa [16]byte
				copy(aaaa[:], net.ParseIP(aaaa_s)[0:16])
				root_addresses6[aaaa] = true
			case *dns.A:
				a_s := root_address.(*dns.A).A.String()
				var a [4]byte
				copy(a[:], net.ParseIP(a_s).To4()[0:4])
				root_addresses4[a] = true
			}
		}
	}
	close(questions)
	return root_addresses4, root_addresses6
}

func ymmv_write(ip_family int, addr []byte, query dns.Msg,
		answer_time time.Time, answer dns.Msg) {
	// output magic value
	_, err := os.Stdout.Write([]byte("ymmv"))
	if err != nil {
		log.Fatal(err)
	}

	// output address family
	if ip_family == 4 {
		_, err = os.Stdout.Write([]byte("4"))
	} else if ip_family == 6 {
		_, err = os.Stdout.Write([]byte("6"))
	} else {
		log.Fatalf("Unknown ip_family %d\n", ip_family)
	}
	if err != nil {
		log.Fatal(err)
	}

	// output (U)DP or (T)CP
	_, err = os.Stdout.Write([]byte("u"))	// only support UDP for now...
	if err != nil {
		log.Fatal(err)
	}

	// output actual address
	_, err = os.Stdout.Write(addr)
	if err != nil {
		log.Fatal(err)
	}

	// output when the query happened (we don't know, so use 0)
	_, err = os.Stdout.Write([]byte{0, 0, 0, 0, 0, 0, 0, 0})

	// write the byte count of our query
	query_bytes, err := query.Pack()
	if err != nil {
		log.Fatal(err)
	}
	query_len := uint16(len(query_bytes))
	err = binary.Write(os.Stdout, binary.BigEndian, query_len)
	if err != nil {
		log.Fatal(err)
	}

	// write our query
	_, err = os.Stdout.Write(query_bytes)
	if err != nil {
		log.Fatal(err)
	}

	// output when the answer arrived
	seconds := uint32(answer_time.Unix())
	err = binary.Write(os.Stdout, binary.BigEndian, seconds)
	if err != nil {
		log.Fatal(err)
	}
	nanoseconds := uint32(answer_time.Nanosecond())
	err = binary.Write(os.Stdout, binary.BigEndian, nanoseconds)
	if err != nil {
		log.Fatal(err)
	}

	// write the byte count of our answer
	answer_bytes, err := answer.Pack()
	if err != nil {
		log.Fatal(err)
	}
	answer_len := uint16(len(answer_bytes))
	err = binary.Write(os.Stdout, binary.BigEndian, answer_len)
	if err != nil {
		log.Fatal(err)
	}

	// write our answer
	_, err = os.Stdout.Write(answer_bytes)
	if err != nil {
		log.Fatal(err)
	}

	// XXX: do we need to flush?
}

func parse_query(raw_answer []byte) (*dns.Msg, *dns.Msg, error) {
	// parse the query
	answer := new(dns.Msg)
	err := answer.Unpack(raw_answer)
	if err != nil {
		return nil, nil, err
	}
	answer.Id = 0
	// infer the answer and build that
	query := answer.Copy()
	query.Response = false
	query.Authoritative = false
	query.Truncated = false
	query.AuthenticatedData = true
	query.CheckingDisabled = false
	query.Rcode = 0
	query.Answer = nil
	query.Ns = nil
	old_extra := query.Extra
	query.Extra = nil
	// add our opt section back - probably not really
	// what we want, but what else can we do?
	if old_extra != nil {
		for _, extra := range old_extra {
			switch extra.(type) {
			case *dns.OPT:
				opt := extra.(*dns.OPT)
				query.Extra = []dns.RR{opt}
			}
		}
	}
	return query, answer, nil
}

// Look in the named file and find any packets that are from our root
// servers on port 53.
func pcap2ymmv(fname string,
               root_addresses4 map[[4]byte]bool,
               root_addresses6 map[[16]byte]bool) {
	// open our pcap file
	var file *os.File
	if fname == "-" {
		file = os.Stdin
	} else {
		named_file, err := os.Open(fname)
		if err != nil {
			log.Fatal(err)
		}
		file = named_file
	}
	pcap_file, err := pcap.NewReader(file)
	if err != nil {
		log.Fatal(err)
	}

	for {
		// reach each packet
		pkt := pcap_file.Next()
		if pkt == nil {
			break
		}
		pkt.Decode()

		// parse each header so we can see if we want this packet
		ip_family := 0
		var ip_addr []byte
		valid_udp := false
		for _, hdr := range pkt.Headers {
			switch hdr.(type) {
			case *pcap.Iphdr:
				// check that the packet comes from one
				// of the addresses that we are looking for
				iphdr := hdr.(*pcap.Iphdr)
				var addr [4]byte
				copy(addr[:], iphdr.SrcIp[0:4])
				_, found :=  root_addresses4[addr]
				if found {
					ip_family = 4
					ip_addr = make([]byte, 4)
					copy(ip_addr[:], addr[0:4])
				}
			case *pcap.Ip6hdr:
				iphdr := hdr.(*pcap.Ip6hdr)
				var addr [16]byte
				copy(addr[:], iphdr.SrcIp[0:16])
				_, found :=  root_addresses6[addr]
				if found {
					ip_family = 6
					ip_addr = make([]byte, 16)
					copy(ip_addr[:], addr[0:16])
				}
			case *pcap.Udphdr:
				udphdr := hdr.(*pcap.Udphdr)
				if udphdr.SrcPort == 53 {
					valid_udp = true
				}
			}
		}

		// if we got a valid IP and UDP packet, process it
		if (ip_family != 0) && valid_udp {
			// parse the payload as the DNS message
			query, answer, err := parse_query(pkt.Payload)
			if err != nil {
				fmt.Fprintf(os.Stderr, "error unpacking DNS message: %s\n", err)
			} else {
				ymmv_write(ip_family, ip_addr,
					   *query, pkt.Time, *answer)
			}
		}
	}
	file.Close()
}

// Main function.
func main() {
	// turn on debugging if desired
	if (len(os.Args) > 1) && (os.Args[1] == "-d") {
		debug = true
		os.Args = append(os.Args[:1], os.Args[2:]...)
	}

	// initialize our stub resolver
	var ( err error )
	resolv_conf, err = dns.ClientConfigFromFile("/etc/resolv.conf")
	if err != nil {
		log.Fatal(err)
	}

	// get root server addresses
	var root_addresses4 map[[4]byte]bool
	var root_addresses6 map[[16]byte]bool
	if len(os.Args) > 1 {
		root_addresses4, root_addresses6 = parse_root_server_addresses(os.Args[1:])
	} else {
		root_addresses4, root_addresses6 = lookup_root_server_addresses()
	}

	// process stdin as a pcap file
	pcap2ymmv("-", root_addresses4, root_addresses6)
}
<commit_msg>Add some debugging. "go fmt"<commit_after>package main

import (
	"encoding/binary"
	"fmt"
	"github.com/miekg/dns"
	"github.com/miekg/pcap"
	"log"
	"net"
	"os"
	"time"
)

var (
	// We store the configuration of our local resolver in a global
	// variable for convenience.
	resolv_conf *dns.ClientConfig

	// Use a global debug flag.
	debug bool
)

// If we were passed name server addresses, parse them with this function.
func parse_root_server_addresses(addrs []string) (map[[4]byte]bool, map[[16]byte]bool) {
	if debug {
		fmt.Fprintf(os.Stderr, "parse_root_server_addresses()\n")
		fmt.Fprintf(os.Stderr, "addrs:%s\n", addrs)
	}
	root_addresses4 := make(map[[4]byte]bool)
	root_addresses6 := make(map[[16]byte]bool)
	for _, addr := range addrs {
		ip := net.ParseIP(addr)
		if ip.To4() == nil {
			// IPv6 address
			var aaaa [16]byte
			copy(aaaa[:], ip[0:16])
			root_addresses6[aaaa] = true
		} else {
			// IPv4 address
			var a [4]byte
			copy(a[:], ip.To4()[0:4])
			root_addresses4[a] = true
		}
	}
	return root_addresses4, root_addresses6
}

// We have a goroutine to act as a stub resolver, and use this
// structure to send the question in and get the results out.
type stub_resolve_info struct {
	ownername string
	rtype     uint16
	answer    *dns.Msg
}

// A goroutine which performs stub lookups from a queue, writing
// the results to another queue.
func stub_resolve(questions <-chan stub_resolve_info,
	results chan<- stub_resolve_info) {
	// make a client for our lookups
	dnsClient := new(dns.Client)
	dnsClient.Net = "tcp"
	// read each question on our channel
	for question := range questions {
		// build our answer
		var result stub_resolve_info = question
		result.answer = nil
		// make a DNS query based on our question
		query := new(dns.Msg)
		query.RecursionDesired = true
		query.SetQuestion(question.ownername, question.rtype)
		// check each resolver in turn
		for _, server := range resolv_conf.Servers {
			resolver := server + ":53"
			r, _, err := dnsClient.Exchange(query, resolver)
			// if we got an answer, use that and stop trying
			if (err == nil) && (r != nil) && (r.Rcode == dns.RcodeSuccess) {
				result.answer = r
				break
			}
		}
		// send back our answer (might be nil)
		results <- result
	}
}

// Lookup all of the IP addresses associated with the root name servers.
// Return two maps based on the results found, which have the keys of
// the binary values of the IPv4 and IPv6 addresses. (It's a bit clumsy,
// but it allows us to do quick and easy lookups of the addresses in the
// pcap later.)
func lookup_root_server_addresses() (map[[4]byte]bool, map[[16]byte]bool) {
	if debug {
		fmt.Fprintf(os.Stderr, "lookup_root_server_addresses()\n")
	}
	// look up the NS of the IANA root
	root_client := new(dns.Client)
	root_client.Net = "tcp"
	ns_query := new(dns.Msg)
	ns_query.SetQuestion(".", dns.TypeNS)
	// TODO: avoid hard-coding a particular root server here
	ns_response, _, err := root_client.Exchange(ns_query,
		"k.root-servers.net:53")
	if err != nil {
		log.Fatal("Error looking up root name servers")
	}
	var root_servers []string
	for _, root_server := range ns_response.Answer {
		switch root_server.(type) {
		case *dns.NS:
			ns := root_server.(*dns.NS).Ns
			root_servers = append(root_servers, ns)
		}
	}
	// look up the addresses of the root servers
	questions := make(chan stub_resolve_info, 16)
	results := make(chan stub_resolve_info, len(root_servers)*2)
	for i := 0; i < 8; i++ {
		go stub_resolve(questions, results)
	}
	for _, ns := range root_servers {
		info := new(stub_resolve_info)
		info.ownername = ns
		info.rtype = dns.TypeAAAA
		questions <- *info
		info = new(stub_resolve_info)
		info.ownername = ns
		info.rtype = dns.TypeA
		questions <- *info
	}
	root_addresses4 := make(map[[4]byte]bool)
	root_addresses6 := make(map[[16]byte]bool)
	for i := 0; i < len(root_servers)*2; i++ {
		response := <-results
		if response.answer == nil {
			log.Fatalf("Error looking up root server %s",
				response.ownername)
		}
		for _, root_address := range response.answer.Answer {
			switch root_address.(type) {
			case *dns.AAAA:
				aaaa_s := root_address.(*dns.AAAA).AAAA.String()
				if debug {
					fmt.Fprintf(os.Stderr, "server: %s\n", aaaa_s)
				}
				var aaaa [16]byte
				copy(aaaa[:], net.ParseIP(aaaa_s)[0:16])
				root_addresses6[aaaa] = true
			case *dns.A:
				a_s := root_address.(*dns.A).A.String()
				if debug {
					fmt.Fprintf(os.Stderr, "server: %s\n", a_s)
				}
				var a [4]byte
				copy(a[:], net.ParseIP(a_s).To4()[0:4])
				root_addresses4[a] = true
			}
		}
	}
	close(questions)
	return root_addresses4, root_addresses6
}

func ymmv_write(ip_family int, addr []byte, query dns.Msg,
	answer_time time.Time, answer dns.Msg) {
	// output magic value
	_, err := os.Stdout.Write([]byte("ymmv"))
	if err != nil {
		log.Fatal(err)
	}

	// output address family
	if ip_family == 4 {
		_, err = os.Stdout.Write([]byte("4"))
	} else if ip_family == 6 {
		_, err = os.Stdout.Write([]byte("6"))
	} else {
		log.Fatalf("Unknown ip_family %d\n", ip_family)
	}
	if err != nil {
		log.Fatal(err)
	}

	// output (U)DP or (T)CP
	_, err = os.Stdout.Write([]byte("u")) // only support UDP for now...
	if err != nil {
		log.Fatal(err)
	}

	// output actual address
	_, err = os.Stdout.Write(addr)
	if err != nil {
		log.Fatal(err)
	}

	// output when the query happened (we don't know, so use 0)
	_, err = os.Stdout.Write([]byte{0, 0, 0, 0, 0, 0, 0, 0})

	// write the byte count of our query
	query_bytes, err := query.Pack()
	if err != nil {
		log.Fatal(err)
	}
	query_len := uint16(len(query_bytes))
	err = binary.Write(os.Stdout, binary.BigEndian, query_len)
	if err != nil {
		log.Fatal(err)
	}

	// write our query
	_, err = os.Stdout.Write(query_bytes)
	if err != nil {
		log.Fatal(err)
	}

	// output when the answer arrived
	seconds := uint32(answer_time.Unix())
	err = binary.Write(os.Stdout, binary.BigEndian, seconds)
	if err != nil {
		log.Fatal(err)
	}
	nanoseconds := uint32(answer_time.Nanosecond())
	err = binary.Write(os.Stdout, binary.BigEndian, nanoseconds)
	if err != nil {
		log.Fatal(err)
	}

	// write the byte count of our answer
	answer_bytes, err := answer.Pack()
	if err != nil {
		log.Fatal(err)
	}
	answer_len := uint16(len(answer_bytes))
	err = binary.Write(os.Stdout, binary.BigEndian, answer_len)
	if err != nil {
		log.Fatal(err)
	}

	// write our answer
	_, err = os.Stdout.Write(answer_bytes)
	if err != nil {
		log.Fatal(err)
	}

	// XXX: do we need to flush?
}

func parse_query(raw_answer []byte) (*dns.Msg, *dns.Msg, error) {
	// parse the query
	answer := new(dns.Msg)
	err := answer.Unpack(raw_answer)
	if err != nil {
		return nil, nil, err
	}
	answer.Id = 0
	// infer the answer and build that
	query := answer.Copy()
	query.Response = false
	query.Authoritative = false
	query.Truncated = false
	query.AuthenticatedData = true
	query.CheckingDisabled = false
	query.Rcode = 0
	query.Answer = nil
	query.Ns = nil
	old_extra := query.Extra
	query.Extra = nil
	// add our opt section back - probably not really
	// what we want, but what else can we do?
	if old_extra != nil {
		for _, extra := range old_extra {
			switch extra.(type) {
			case *dns.OPT:
				opt := extra.(*dns.OPT)
				query.Extra = []dns.RR{opt}
			}
		}
	}
	return query, answer, nil
}

// Look in the named file and find any packets that are from our root
// servers on port 53.
func pcap2ymmv(fname string,
	root_addresses4 map[[4]byte]bool,
	root_addresses6 map[[16]byte]bool) {
	// open our pcap file
	var file *os.File
	if fname == "-" {
		file = os.Stdin
	} else {
		named_file, err := os.Open(fname)
		if err != nil {
			log.Fatal(err)
		}
		file = named_file
	}
	pcap_file, err := pcap.NewReader(file)
	if err != nil {
		log.Fatal(err)
	}

	for {
		// reach each packet
		pkt := pcap_file.Next()
		if pkt == nil {
			break
		}
		pkt.Decode()
		if debug {
			fmt.Fprintf(os.Stderr, "pcap2ymmv read packet\n")
		}

		// parse each header so we can see if we want this packet
		ip_family := 0
		var ip_addr []byte
		valid_udp := false
		for _, hdr := range pkt.Headers {
			switch hdr.(type) {
			case *pcap.Iphdr:
				// check that the packet comes from one
				// of the addresses that we are looking for
				iphdr := hdr.(*pcap.Iphdr)
				if debug {
					fmt.Fprintf(os.Stderr, "pcap2ymmv source IP %s\n", net.IP(iphdr.SrcIp).String())
				}
				var addr [4]byte
				copy(addr[:], iphdr.SrcIp[0:4])
				_, found := root_addresses4[addr]
				if found {
					if debug {
						fmt.Fprintf(os.Stderr, "pcap2ymmv address match\n")
					}
					ip_family = 4
					ip_addr = make([]byte, 4)
					copy(ip_addr[:], addr[0:4])
				}
			case *pcap.Ip6hdr:
				iphdr := hdr.(*pcap.Ip6hdr)
				if debug {
					fmt.Fprintf(os.Stderr, "pcap2ymmv source IP %s\n", net.IP(iphdr.SrcIp).String())
				}
				var addr [16]byte
				copy(addr[:], iphdr.SrcIp[0:16])
				_, found := root_addresses6[addr]
				if found {
					if debug {
						fmt.Fprintf(os.Stderr, "pcap2ymmv address match\n")
					}
					ip_family = 6
					ip_addr = make([]byte, 16)
					copy(ip_addr[:], addr[0:16])
				}
			case *pcap.Udphdr:
				udphdr := hdr.(*pcap.Udphdr)
				if udphdr.SrcPort == 53 {
					valid_udp = true
				}
			}
		}

		// if we got a valid IP and UDP packet, process it
		if (ip_family != 0) && valid_udp {
			if debug {
				fmt.Fprintf(os.Stderr, "pcap2ymmv matched packet\n")
			}
			// parse the payload as the DNS message
			query, answer, err := parse_query(pkt.Payload)
			if err != nil {
				fmt.Fprintf(os.Stderr, "error unpacking DNS message: %s\n", err)
			} else {
				ymmv_write(ip_family, ip_addr,
					*query, pkt.Time, *answer)
				os.Stdout.Sync()
			}
		}
	}
	file.Close()
}

// Main function.
func main() {
	// turn on debugging if desired
	if (len(os.Args) > 1) && (os.Args[1] == "-d") {
		debug = true
		os.Args = append(os.Args[:1], os.Args[2:]...)
	}

	// initialize our stub resolver
	var (
		err error
	)
	resolv_conf, err = dns.ClientConfigFromFile("/etc/resolv.conf")
	if err != nil {
		log.Fatal(err)
	}

	// get root server addresses
	var root_addresses4 map[[4]byte]bool
	var root_addresses6 map[[16]byte]bool
	if len(os.Args) > 1 {
		root_addresses4, root_addresses6 = parse_root_server_addresses(os.Args[1:])
	} else {
		root_addresses4, root_addresses6 = lookup_root_server_addresses()
	}

	// process stdin as a pcap file
	pcap2ymmv("-", root_addresses4, root_addresses6)
}
<|endoftext|>


<commit_before>package getter

import (
	"path/filepath"
	"testing"
)

func TestTarGzipDecompressor(t *testing.T) {
	cases := []TestDecompressCase{
		{
			"empty.tar.gz",
			false,
			true,
			nil,
			"",
		},

		{
			"single.tar.gz",
			false,
			false,
			nil,
			"d3b07384d113edec49eaa6238ad5ff00",
		},

		{
			"single.tar.gz",
			true,
			false,
			[]string{"file"},
			"",
		},

		{
			"multiple.tar.gz",
			true,
			false,
			[]string{"file1", "file2"},
			"",
		},

		{
			"multiple.tar.gz",
			false,
			true,
			nil,
			"",
		},

		{
			"multiple_dir.tar.gz",
			true,
			false,
			[]string{"dir/", "dir/test2", "test1"},
			"",
		},
	}

	for i, tc := range cases {
		cases[i].Input = filepath.Join("./test-fixtures", "decompress-tgz", tc.Input)
	}

	TestDecompressor(t, new(TarGzipDecompressor), cases)
}
<commit_msg>tgz/windows: Fix the TestTarGzipDecompressor test<commit_after>package getter

import (
	"path/filepath"
	"runtime"
	"testing"
)

func TestTarGzipDecompressor(t *testing.T) {

	multiplePaths := []string{"dir/", "dir/test2", "test1"}
	if runtime.GOOS == "windows" {
		multiplePaths = []string{"dir/", "dir\\test2", "test1"}
	}

	cases := []TestDecompressCase{
		{
			"empty.tar.gz",
			false,
			true,
			nil,
			"",
		},

		{
			"single.tar.gz",
			false,
			false,
			nil,
			"d3b07384d113edec49eaa6238ad5ff00",
		},

		{
			"single.tar.gz",
			true,
			false,
			[]string{"file"},
			"",
		},

		{
			"multiple.tar.gz",
			true,
			false,
			[]string{"file1", "file2"},
			"",
		},

		{
			"multiple.tar.gz",
			false,
			true,
			nil,
			"",
		},

		{
			"multiple_dir.tar.gz",
			true,
			false,
			multiplePaths,
			"",
		},
	}

	for i, tc := range cases {
		cases[i].Input = filepath.Join("./test-fixtures", "decompress-tgz", tc.Input)
	}

	TestDecompressor(t, new(TarGzipDecompressor), cases)
}
<|endoftext|>


<commit_before>package stream_test

import (
	"context"
	"io"
	"net"
	"sync"
	"testing"
	"time"

	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/sec"
	"github.com/libp2p/go-libp2p-core/transport"
	st "github.com/libp2p/go-libp2p-transport-upgrader"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr/net"

	"github.com/stretchr/testify/require"
)

func init() {
	transport.AcceptTimeout = 1 * time.Hour
}

type MuxAdapter struct {
	tpt sec.SecureTransport
}

func (mux *MuxAdapter) SecureInbound(ctx context.Context, insecure net.Conn) (sec.SecureConn, bool, error) {
	sconn, err := mux.tpt.SecureInbound(ctx, insecure)
	return sconn, true, err
}

func (mux *MuxAdapter) SecureOutbound(ctx context.Context, insecure net.Conn, p peer.ID) (sec.SecureConn, bool, error) {
	sconn, err := mux.tpt.SecureOutbound(ctx, insecure, p)
	return sconn, false, err
}

func createListener(t *testing.T, upgrader *st.Upgrader) transport.Listener {
	t.Helper()
	addr, err := ma.NewMultiaddr("/ip4/127.0.0.1/tcp/0")
	require.NoError(t, err)
	ln, err := manet.Listen(addr)
	require.NoError(t, err)
	return upgrader.UpgradeListener(nil, ln)
}

func TestAcceptSingleConn(t *testing.T) {
	require := require.New(t)

	id, upgrader := createUpgrader(t)
	ln := createListener(t, upgrader)
	defer ln.Close()

	cconn, err := dial(t, upgrader, ln.Multiaddr(), id)
	require.NoError(err)

	sconn, err := ln.Accept()
	require.NoError(err)

	testConn(t, cconn, sconn)
}

func TestAcceptMultipleConns(t *testing.T) {
	require := require.New(t)

	id, upgrader := createUpgrader(t)
	ln := createListener(t, upgrader)
	defer ln.Close()

	var toClose []io.Closer
	defer func() {
		for _, c := range toClose {
			_ = c.Close()
		}
	}()

	for i := 0; i < 10; i++ {
		cconn, err := dial(t, upgrader, ln.Multiaddr(), id)
		require.NoError(err)
		toClose = append(toClose, cconn)

		sconn, err := ln.Accept()
		require.NoError(err)
		toClose = append(toClose, sconn)

		testConn(t, cconn, sconn)
	}
}

func TestConnectionsClosedIfNotAccepted(t *testing.T) {
	require := require.New(t)

	const timeout = 200 * time.Millisecond
	transport.AcceptTimeout = timeout
	defer func() { transport.AcceptTimeout = 1 * time.Hour }()

	id, upgrader := createUpgrader(t)
	ln := createListener(t, upgrader)
	defer ln.Close()

	conn, err := dial(t, upgrader, ln.Multiaddr(), id)
	require.NoError(err)

	errCh := make(chan error)
	go func() {
		defer conn.Close()
		str, err := conn.OpenStream(context.Background())
		if err != nil {
			errCh <- err
			return
		}
		// start a Read. It will block until the connection is closed
		_, _ = str.Read([]byte{0})
		errCh <- nil
	}()

	time.Sleep(timeout / 2)
	select {
	case err := <-errCh:
		t.Fatalf("connection closed earlier than expected. expected nothing on channel, got: %v", err)
	default:
	}

	time.Sleep(timeout)
	require.Nil(<-errCh)
}

func TestFailedUpgradeOnListen(t *testing.T) {
	require := require.New(t)

	id, upgrader := createUpgrader(t)
	upgrader.Muxer = &errorMuxer{}
	ln := createListener(t, upgrader)
	defer ln.Close()

	errCh := make(chan error)
	go func() {
		_, err := ln.Accept()
		errCh <- err
	}()

	_, err := dial(t, upgrader, ln.Multiaddr(), id)
	require.Error(err)

	// close the listener.
	ln.Close()
	require.Error(<-errCh)
}

func TestListenerClose(t *testing.T) {
	require := require.New(t)

	_, upgrader := createUpgrader(t)
	ln := createListener(t, upgrader)

	errCh := make(chan error)
	go func() {
		_, err := ln.Accept()
		errCh <- err
	}()

	select {
	case err := <-errCh:
		t.Fatalf("connection closed earlier than expected. expected nothing on channel, got: %v", err)
	case <-time.After(200 * time.Millisecond):
		// nothing in 200ms.
	}

	// unblocks Accept when it is closed.
	err := ln.Close()
	require.NoError(err)
	err = <-errCh
	require.Error(err)
	require.Contains(err.Error(), "use of closed network connection")

	// doesn't accept new connections when it is closed
	_, err = dial(t, upgrader, ln.Multiaddr(), peer.ID("1"))
	require.Error(err)
}

func TestListenerCloseClosesQueued(t *testing.T) {
	require := require.New(t)

	id, upgrader := createUpgrader(t)
	ln := createListener(t, upgrader)

	var conns []transport.CapableConn
	for i := 0; i < 10; i++ {
		conn, err := dial(t, upgrader, ln.Multiaddr(), id)
		require.NoError(err)
		conns = append(conns, conn)
	}

	// wait for all the dials to happen.
	time.Sleep(500 * time.Millisecond)

	// all the connections are opened.
	for _, c := range conns {
		require.False(c.IsClosed())
	}

	// expect that all the connections will be closed.
	err := ln.Close()
	require.NoError(err)

	// all the connections are closed.
	require.Eventually(func() bool {
		for _, c := range conns {
			if !c.IsClosed() {
				return false
			}
		}
		return true
	}, 3*time.Second, 100*time.Millisecond)

	for _, c := range conns {
		_ = c.Close()
	}
}

func TestConcurrentAccept(t *testing.T) {
	var num = 3 * st.AcceptQueueLength

	id, upgrader := createUpgrader(t)
	blockingMuxer := newBlockingMuxer()
	upgrader.Muxer = blockingMuxer

	ln := createListener(t, upgrader)
	defer ln.Close()

	accepted := make(chan transport.CapableConn, num)
	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				return
			}
			_ = conn.Close()
			accepted <- conn
		}
	}()

	// start num dials, which all block while setting up the muxer
	errCh := make(chan error, num)
	var wg sync.WaitGroup
	for i := 0; i < num; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()

			conn, err := dial(t, upgrader, ln.Multiaddr(), id)
			if err != nil {
				errCh <- err
				return
			}
			defer conn.Close()

			_, err = conn.AcceptStream() // wait for conn to be accepted.
			errCh <- err
		}()
	}

	time.Sleep(200 * time.Millisecond)
	// the dials are still blocked, so we shouldn't have any connection available yet
	require.Empty(t, accepted)
	blockingMuxer.Unblock() // make all dials succeed
	require.Eventually(t, func() bool { return len(accepted) == num }, 3*time.Second, 100*time.Millisecond)
	wg.Wait()
}

func TestAcceptQueueBacklogged(t *testing.T) {
	require := require.New(t)

	id, upgrader := createUpgrader(t)
	ln := createListener(t, upgrader)
	defer ln.Close()

	// setup AcceptQueueLength connections, but don't accept any of them
	errCh := make(chan error, st.AcceptQueueLength+1)
	doDial := func() {
		conn, err := dial(t, upgrader, ln.Multiaddr(), id)
		errCh <- err
		if conn != nil {
			_ = conn.Close()
		}
	}

	for i := 0; i < st.AcceptQueueLength; i++ {
		go doDial()
	}

	require.Eventually(func() bool { return len(errCh) == st.AcceptQueueLength }, 2*time.Second, 100*time.Millisecond)

	// dial a new connection. This connection should not complete setup, since the queue is full
	go doDial()

	time.Sleep(500 * time.Millisecond)
	require.Len(errCh, st.AcceptQueueLength)

	// accept a single connection. Now the new connection should be set up, and fill the queue again
	conn, err := ln.Accept()
	require.NoError(err)
	_ = conn.Close()

	require.Eventually(func() bool { return len(errCh) == st.AcceptQueueLength+1 }, 2*time.Second, 100*time.Millisecond)
}

func TestListenerConnectionGater(t *testing.T) {
	require := require.New(t)

	testGater := &testGater{}
	id, upgrader := createUpgrader(t)
	upgrader.ConnGater = testGater

	ln := createListener(t, upgrader)
	defer ln.Close()

	// no gating.
	conn, err := dial(t, upgrader, ln.Multiaddr(), id)
	require.NoError(err)
	require.False(conn.IsClosed())
	_ = conn.Close()

	// rejecting after handshake.
	testGater.BlockSecured(true)
	testGater.BlockAccept(false)
	conn, err = dial(t, upgrader, ln.Multiaddr(), peer.ID("invalid"))
	require.Error(err)
	require.Nil(conn)

	// rejecting on accept will trigger first.
	testGater.BlockSecured(true)
	testGater.BlockAccept(true)
	conn, err = dial(t, upgrader, ln.Multiaddr(), peer.ID("invalid"))
	require.Error(err)
	require.Nil(conn)

	// rejecting only on acceptance.
	testGater.BlockSecured(false)
	testGater.BlockAccept(true)
	conn, err = dial(t, upgrader, ln.Multiaddr(), peer.ID("invalid"))
	require.Error(err)
	require.Nil(conn)

	// back to normal
	testGater.BlockSecured(false)
	testGater.BlockAccept(false)
	conn, err = dial(t, upgrader, ln.Multiaddr(), id)
	require.NoError(err)
	require.False(conn.IsClosed())
	_ = conn.Close()
}
<commit_msg>increase timeout in TestConnectionsClosedIfNotAccepted on CI<commit_after>package stream_test

import (
	"context"
	"io"
	"net"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/sec"
	"github.com/libp2p/go-libp2p-core/transport"
	st "github.com/libp2p/go-libp2p-transport-upgrader"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr/net"

	"github.com/stretchr/testify/require"
)

func init() {
	transport.AcceptTimeout = 1 * time.Hour
}

type MuxAdapter struct {
	tpt sec.SecureTransport
}

func (mux *MuxAdapter) SecureInbound(ctx context.Context, insecure net.Conn) (sec.SecureConn, bool, error) {
	sconn, err := mux.tpt.SecureInbound(ctx, insecure)
	return sconn, true, err
}

func (mux *MuxAdapter) SecureOutbound(ctx context.Context, insecure net.Conn, p peer.ID) (sec.SecureConn, bool, error) {
	sconn, err := mux.tpt.SecureOutbound(ctx, insecure, p)
	return sconn, false, err
}

func createListener(t *testing.T, upgrader *st.Upgrader) transport.Listener {
	t.Helper()
	addr, err := ma.NewMultiaddr("/ip4/127.0.0.1/tcp/0")
	require.NoError(t, err)
	ln, err := manet.Listen(addr)
	require.NoError(t, err)
	return upgrader.UpgradeListener(nil, ln)
}

func TestAcceptSingleConn(t *testing.T) {
	require := require.New(t)

	id, upgrader := createUpgrader(t)
	ln := createListener(t, upgrader)
	defer ln.Close()

	cconn, err := dial(t, upgrader, ln.Multiaddr(), id)
	require.NoError(err)

	sconn, err := ln.Accept()
	require.NoError(err)

	testConn(t, cconn, sconn)
}

func TestAcceptMultipleConns(t *testing.T) {
	require := require.New(t)

	id, upgrader := createUpgrader(t)
	ln := createListener(t, upgrader)
	defer ln.Close()

	var toClose []io.Closer
	defer func() {
		for _, c := range toClose {
			_ = c.Close()
		}
	}()

	for i := 0; i < 10; i++ {
		cconn, err := dial(t, upgrader, ln.Multiaddr(), id)
		require.NoError(err)
		toClose = append(toClose, cconn)

		sconn, err := ln.Accept()
		require.NoError(err)
		toClose = append(toClose, sconn)

		testConn(t, cconn, sconn)
	}
}

func TestConnectionsClosedIfNotAccepted(t *testing.T) {
	require := require.New(t)

	var timeout = 100 * time.Millisecond
	if os.Getenv("CI") != "" {
		timeout = 500 * time.Millisecond
	}
	origAcceptTimeout := transport.AcceptTimeout
	transport.AcceptTimeout = timeout
	t.Cleanup(func() { transport.AcceptTimeout = origAcceptTimeout })

	id, upgrader := createUpgrader(t)
	ln := createListener(t, upgrader)
	defer ln.Close()

	conn, err := dial(t, upgrader, ln.Multiaddr(), id)
	require.NoError(err)

	errCh := make(chan error)
	go func() {
		defer conn.Close()
		str, err := conn.OpenStream(context.Background())
		if err != nil {
			errCh <- err
			return
		}
		// start a Read. It will block until the connection is closed
		_, _ = str.Read([]byte{0})
		errCh <- nil
	}()

	time.Sleep(timeout / 2)
	select {
	case err := <-errCh:
		t.Fatalf("connection closed earlier than expected. expected nothing on channel, got: %v", err)
	default:
	}

	time.Sleep(timeout)
	require.Nil(<-errCh)
}

func TestFailedUpgradeOnListen(t *testing.T) {
	require := require.New(t)

	id, upgrader := createUpgrader(t)
	upgrader.Muxer = &errorMuxer{}
	ln := createListener(t, upgrader)
	defer ln.Close()

	errCh := make(chan error)
	go func() {
		_, err := ln.Accept()
		errCh <- err
	}()

	_, err := dial(t, upgrader, ln.Multiaddr(), id)
	require.Error(err)

	// close the listener.
	ln.Close()
	require.Error(<-errCh)
}

func TestListenerClose(t *testing.T) {
	require := require.New(t)

	_, upgrader := createUpgrader(t)
	ln := createListener(t, upgrader)

	errCh := make(chan error)
	go func() {
		_, err := ln.Accept()
		errCh <- err
	}()

	select {
	case err := <-errCh:
		t.Fatalf("connection closed earlier than expected. expected nothing on channel, got: %v", err)
	case <-time.After(200 * time.Millisecond):
		// nothing in 200ms.
	}

	// unblocks Accept when it is closed.
	err := ln.Close()
	require.NoError(err)
	err = <-errCh
	require.Error(err)
	require.Contains(err.Error(), "use of closed network connection")

	// doesn't accept new connections when it is closed
	_, err = dial(t, upgrader, ln.Multiaddr(), peer.ID("1"))
	require.Error(err)
}

func TestListenerCloseClosesQueued(t *testing.T) {
	require := require.New(t)

	id, upgrader := createUpgrader(t)
	ln := createListener(t, upgrader)

	var conns []transport.CapableConn
	for i := 0; i < 10; i++ {
		conn, err := dial(t, upgrader, ln.Multiaddr(), id)
		require.NoError(err)
		conns = append(conns, conn)
	}

	// wait for all the dials to happen.
	time.Sleep(500 * time.Millisecond)

	// all the connections are opened.
	for _, c := range conns {
		require.False(c.IsClosed())
	}

	// expect that all the connections will be closed.
	err := ln.Close()
	require.NoError(err)

	// all the connections are closed.
	require.Eventually(func() bool {
		for _, c := range conns {
			if !c.IsClosed() {
				return false
			}
		}
		return true
	}, 3*time.Second, 100*time.Millisecond)

	for _, c := range conns {
		_ = c.Close()
	}
}

func TestConcurrentAccept(t *testing.T) {
	var num = 3 * st.AcceptQueueLength

	id, upgrader := createUpgrader(t)
	blockingMuxer := newBlockingMuxer()
	upgrader.Muxer = blockingMuxer

	ln := createListener(t, upgrader)
	defer ln.Close()

	accepted := make(chan transport.CapableConn, num)
	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				return
			}
			_ = conn.Close()
			accepted <- conn
		}
	}()

	// start num dials, which all block while setting up the muxer
	errCh := make(chan error, num)
	var wg sync.WaitGroup
	for i := 0; i < num; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()

			conn, err := dial(t, upgrader, ln.Multiaddr(), id)
			if err != nil {
				errCh <- err
				return
			}
			defer conn.Close()

			_, err = conn.AcceptStream() // wait for conn to be accepted.
			errCh <- err
		}()
	}

	time.Sleep(200 * time.Millisecond)
	// the dials are still blocked, so we shouldn't have any connection available yet
	require.Empty(t, accepted)
	blockingMuxer.Unblock() // make all dials succeed
	require.Eventually(t, func() bool { return len(accepted) == num }, 3*time.Second, 100*time.Millisecond)
	wg.Wait()
}

func TestAcceptQueueBacklogged(t *testing.T) {
	require := require.New(t)

	id, upgrader := createUpgrader(t)
	ln := createListener(t, upgrader)
	defer ln.Close()

	// setup AcceptQueueLength connections, but don't accept any of them
	errCh := make(chan error, st.AcceptQueueLength+1)
	doDial := func() {
		conn, err := dial(t, upgrader, ln.Multiaddr(), id)
		errCh <- err
		if conn != nil {
			_ = conn.Close()
		}
	}

	for i := 0; i < st.AcceptQueueLength; i++ {
		go doDial()
	}

	require.Eventually(func() bool { return len(errCh) == st.AcceptQueueLength }, 2*time.Second, 100*time.Millisecond)

	// dial a new connection. This connection should not complete setup, since the queue is full
	go doDial()

	time.Sleep(500 * time.Millisecond)
	require.Len(errCh, st.AcceptQueueLength)

	// accept a single connection. Now the new connection should be set up, and fill the queue again
	conn, err := ln.Accept()
	require.NoError(err)
	_ = conn.Close()

	require.Eventually(func() bool { return len(errCh) == st.AcceptQueueLength+1 }, 2*time.Second, 100*time.Millisecond)
}

func TestListenerConnectionGater(t *testing.T) {
	require := require.New(t)

	testGater := &testGater{}
	id, upgrader := createUpgrader(t)
	upgrader.ConnGater = testGater

	ln := createListener(t, upgrader)
	defer ln.Close()

	// no gating.
	conn, err := dial(t, upgrader, ln.Multiaddr(), id)
	require.NoError(err)
	require.False(conn.IsClosed())
	_ = conn.Close()

	// rejecting after handshake.
	testGater.BlockSecured(true)
	testGater.BlockAccept(false)
	conn, err = dial(t, upgrader, ln.Multiaddr(), peer.ID("invalid"))
	require.Error(err)
	require.Nil(conn)

	// rejecting on accept will trigger first.
	testGater.BlockSecured(true)
	testGater.BlockAccept(true)
	conn, err = dial(t, upgrader, ln.Multiaddr(), peer.ID("invalid"))
	require.Error(err)
	require.Nil(conn)

	// rejecting only on acceptance.
	testGater.BlockSecured(false)
	testGater.BlockAccept(true)
	conn, err = dial(t, upgrader, ln.Multiaddr(), peer.ID("invalid"))
	require.Error(err)
	require.Nil(conn)

	// back to normal
	testGater.BlockSecured(false)
	testGater.BlockAccept(false)
	conn, err = dial(t, upgrader, ln.Multiaddr(), id)
	require.NoError(err)
	require.False(conn.IsClosed())
	_ = conn.Close()
}
<|endoftext|>


<commit_before>/*
Copyright 2016, Cossack Labs Limited

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 base

import (
	"bytes"
	"encoding/binary"

	"errors"
	"github.com/cossacklabs/acra/keystore"
	"github.com/cossacklabs/acra/utils"
	"github.com/cossacklabs/themis/gothemis/cell"
	"github.com/cossacklabs/themis/gothemis/keys"
	"github.com/cossacklabs/themis/gothemis/message"
)

// GetDataLengthFromAcraStruct unpack data length value from AcraStruct
func GetDataLengthFromAcraStruct(data []byte) int {
	dataLengthBlock := data[GetMinAcraStructLength()-DataLengthSize : GetMinAcraStructLength()]
	return int(binary.LittleEndian.Uint64(dataLengthBlock))
}

// GetMinAcraStructLength returns minimal length of AcraStruct
// because in golang we can't declare byte array as constant we need to calculate length of TagBegin in runtime
// or hardcode as constant and maintain len(TagBegin) == CONST_VALUE
func GetMinAcraStructLength() int {
	return len(TagBegin) + KeyBlockLength + DataLengthSize
}

// Errors show incorrect AcraStruct length
var (
	ErrIncorrectAcraStructTagBegin   = errors.New("AcraStruct has incorrect TagBegin")
	ErrIncorrectAcraStructLength     = errors.New("AcraStruct has incorrect length")
	ErrIncorrectAcraStructDataLength = errors.New("AcraStruct has incorrect data length value")
)

// ValidateAcraStructLength check that data has minimal length for AcraStruct and data block equal to data length in AcraStruct
func ValidateAcraStructLength(data []byte) error {
	baseLength := GetMinAcraStructLength()
	if len(data) < baseLength {
		return ErrIncorrectAcraStructLength
	}
	if !bytes.Equal(data[:len(TagBegin)], TagBegin) {
		return ErrIncorrectAcraStructTagBegin
	}
	dataLength := GetDataLengthFromAcraStruct(data)
	if dataLength != len(data[GetMinAcraStructLength():]) {
		return ErrIncorrectAcraStructDataLength
	}
	return nil
}

// DecryptAcrastruct returns plaintext data from AcraStruct, decrypting it using Themis SecureCell in Seal mode,
// using zone as context and privateKey as decryption key.
// Returns error if decryption failed.
func DecryptAcrastruct(data []byte, privateKey *keys.PrivateKey, zone []byte) ([]byte, error) {
	if err := ValidateAcraStructLength(data); err != nil {
		return nil, err
	}
	innerData := data[len(TagBegin):]
	pubkey := &keys.PublicKey{Value: innerData[:PublicKeyLength]}
	smessage := message.New(privateKey, pubkey)
	symmetricKey, err := smessage.Unwrap(innerData[PublicKeyLength:KeyBlockLength])
	if err != nil {
		return []byte{}, err
	}
	//
	var length uint64
	// convert from little endian
	err = binary.Read(bytes.NewReader(innerData[KeyBlockLength:KeyBlockLength+DataLengthSize]), binary.LittleEndian, &length)
	if err != nil {
		return []byte{}, err
	}
	scell := cell.New(symmetricKey, cell.CELL_MODE_SEAL)
	decrypted, err := scell.Unprotect(innerData[KeyBlockLength+DataLengthSize:], nil, zone)
	// fill zero symmetric_key
	utils.FillSlice(byte(0), symmetricKey)
	if err != nil {
		return []byte{}, err
	}
	return decrypted, nil
}

// CheckPoisonRecord checks if AcraStruct could be decrypted using Poison Record private key.
// Returns true if AcraStruct is poison record, returns false otherwise.
// Returns error if Poison record key is not found.
func CheckPoisonRecord(data []byte, keystorage keystore.KeyStore) (bool, error) {
	poisonKeypair, err := keystorage.GetPoisonKeyPair()
	if err != nil {
		// we can't check on poisoning
		return true, err
	}
	_, err = DecryptAcrastruct(data, poisonKeypair.Private, nil)
	utils.FillSlice(byte(0), poisonKeypair.Private.Value)
	if err == nil {
		// decryption success so it was encrypted with private key for poison records
		return true, nil
	}
	return false, nil
}
<commit_msg>Multi-key decryption utility<commit_after>/*
Copyright 2016, Cossack Labs Limited

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 base

import (
	"bytes"
	"encoding/binary"

	"errors"
	"github.com/cossacklabs/acra/keystore"
	"github.com/cossacklabs/acra/utils"
	"github.com/cossacklabs/themis/gothemis/cell"
	"github.com/cossacklabs/themis/gothemis/keys"
	"github.com/cossacklabs/themis/gothemis/message"
)

// GetDataLengthFromAcraStruct unpack data length value from AcraStruct
func GetDataLengthFromAcraStruct(data []byte) int {
	dataLengthBlock := data[GetMinAcraStructLength()-DataLengthSize : GetMinAcraStructLength()]
	return int(binary.LittleEndian.Uint64(dataLengthBlock))
}

// GetMinAcraStructLength returns minimal length of AcraStruct
// because in golang we can't declare byte array as constant we need to calculate length of TagBegin in runtime
// or hardcode as constant and maintain len(TagBegin) == CONST_VALUE
func GetMinAcraStructLength() int {
	return len(TagBegin) + KeyBlockLength + DataLengthSize
}

// Errors show incorrect AcraStruct length
var (
	ErrIncorrectAcraStructTagBegin   = errors.New("AcraStruct has incorrect TagBegin")
	ErrIncorrectAcraStructLength     = errors.New("AcraStruct has incorrect length")
	ErrIncorrectAcraStructDataLength = errors.New("AcraStruct has incorrect data length value")
)

// ValidateAcraStructLength check that data has minimal length for AcraStruct and data block equal to data length in AcraStruct
func ValidateAcraStructLength(data []byte) error {
	baseLength := GetMinAcraStructLength()
	if len(data) < baseLength {
		return ErrIncorrectAcraStructLength
	}
	if !bytes.Equal(data[:len(TagBegin)], TagBegin) {
		return ErrIncorrectAcraStructTagBegin
	}
	dataLength := GetDataLengthFromAcraStruct(data)
	if dataLength != len(data[GetMinAcraStructLength():]) {
		return ErrIncorrectAcraStructDataLength
	}
	return nil
}

// DecryptAcrastruct returns plaintext data from AcraStruct, decrypting it using Themis SecureCell in Seal mode,
// using zone as context and privateKey as decryption key.
// Returns error if decryption failed.
func DecryptAcrastruct(data []byte, privateKey *keys.PrivateKey, zone []byte) ([]byte, error) {
	if err := ValidateAcraStructLength(data); err != nil {
		return nil, err
	}
	innerData := data[len(TagBegin):]
	pubkey := &keys.PublicKey{Value: innerData[:PublicKeyLength]}
	smessage := message.New(privateKey, pubkey)
	symmetricKey, err := smessage.Unwrap(innerData[PublicKeyLength:KeyBlockLength])
	if err != nil {
		return []byte{}, err
	}
	//
	var length uint64
	// convert from little endian
	err = binary.Read(bytes.NewReader(innerData[KeyBlockLength:KeyBlockLength+DataLengthSize]), binary.LittleEndian, &length)
	if err != nil {
		return []byte{}, err
	}
	scell := cell.New(symmetricKey, cell.CELL_MODE_SEAL)
	decrypted, err := scell.Unprotect(innerData[KeyBlockLength+DataLengthSize:], nil, zone)
	// fill zero symmetric_key
	utils.FillSlice(byte(0), symmetricKey)
	if err != nil {
		return []byte{}, err
	}
	return decrypted, nil
}

// DecryptRotatedAcrastruct tries decrypting an AcraStruct with a set of rotated keys.
// It either returns decrypted data if one of the keys succeeds, or an error if none is good.
func DecryptRotatedAcrastruct(data []byte, privateKeys []*keys.PrivateKey, zone []byte) ([]byte, error) {
	var err error
	for _, privateKey := range privateKeys {
		data, err := DecryptAcrastruct(data, privateKey, zone)
		if err == nil {
			return data, err
		}
	}
	return nil, err
}

// CheckPoisonRecord checks if AcraStruct could be decrypted using Poison Record private key.
// Returns true if AcraStruct is poison record, returns false otherwise.
// Returns error if Poison record key is not found.
func CheckPoisonRecord(data []byte, keystorage keystore.KeyStore) (bool, error) {
	poisonKeypair, err := keystorage.GetPoisonKeyPair()
	if err != nil {
		// we can't check on poisoning
		return true, err
	}
	_, err = DecryptAcrastruct(data, poisonKeypair.Private, nil)
	utils.FillSlice(byte(0), poisonKeypair.Private.Value)
	if err == nil {
		// decryption success so it was encrypted with private key for poison records
		return true, nil
	}
	return false, nil
}
<|endoftext|>


<commit_before>// Copyright 2021 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package main

import (
	"bytes"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/pkg/rpctype"
	"github.com/google/syzkaller/prog"
)

func TestFinalizeCallSet(t *testing.T) {
	target, err := prog.GetTarget("test", "64")
	if err != nil {
		t.Fatalf("failed to initialise test target: %v", err)
	}

	vrf := Verifier{
		target: target,
		reasons: map[*prog.Syscall]string{
			target.SyscallMap["test$res0"]:  "foo",
			target.SyscallMap["minimize$0"]: "bar",
		},
		calls: map[*prog.Syscall]bool{
			target.SyscallMap["minimize$0"]: true,
			target.SyscallMap["test$res0"]:  true,
			target.SyscallMap["disabled1"]:  true,
		},
		reportReasons: true,
	}

	out := bytes.Buffer{}
	vrf.finalizeCallSet(&out)
	wantLines := []string{
		"The following calls have been disabled:\n",
		"\ttest$res0: foo\n",
		"\tminimize$0: bar\n",
	}
	output := out.String()
	for _, line := range wantLines {
		if !strings.Contains(output, line) {
			t.Errorf("finalizeCallSet: %q missing in reported output", line)
		}
	}

	wantCalls, gotCalls := map[*prog.Syscall]bool{
		target.SyscallMap["disabled1"]: true,
	}, vrf.calls
	if diff := cmp.Diff(wantCalls, gotCalls); diff != "" {
		t.Errorf("srv.calls mismatch (-want +got):\n%s", diff)
	}
}

func TestUpdateUnsupported(t *testing.T) {
	target, err := prog.GetTarget("test", "64")
	if err != nil {
		t.Fatalf("failed to initialise test target: %v", err)
	}

	tests := []struct {
		name           string
		vrfPools       map[int]*poolInfo
		wantPools      map[int]*poolInfo
		wantCalls      map[*prog.Syscall]bool
		wantNotChecked int
		nilCT          bool
	}{
		{
			name:           "choice table not generated",
			vrfPools:       map[int]*poolInfo{0: {}, 1: {}},
			wantPools:      map[int]*poolInfo{0: {checked: true}, 1: {}},
			wantNotChecked: 1,
			wantCalls: map[*prog.Syscall]bool{
				target.SyscallMap["minimize$0"]:     true,
				target.SyscallMap["breaks_returns"]: true,
				target.SyscallMap["test$res0"]:      true,
				target.SyscallMap["test$union0"]:    true,
			},
			nilCT: true,
		},
		{
			name:           "choice table generated",
			vrfPools:       map[int]*poolInfo{0: {}},
			wantPools:      map[int]*poolInfo{0: {checked: true}},
			wantNotChecked: 0,
			wantCalls: map[*prog.Syscall]bool{
				target.SyscallMap["minimize$0"]:     true,
				target.SyscallMap["breaks_returns"]: true,
			},
			nilCT: false,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			vrf := Verifier{
				target:        target,
				pools:         test.vrfPools,
				reasons:       make(map[*prog.Syscall]string),
				reportReasons: true,
				calls: map[*prog.Syscall]bool{
					target.SyscallMap["minimize$0"]:     true,
					target.SyscallMap["breaks_returns"]: true,
					target.SyscallMap["test$res0"]:      true,
					target.SyscallMap["test$union0"]:    true,
				},
				stats: MakeStats(),
			}
			vrf.Init()

			a := &rpctype.UpdateUnsupportedArgs{
				Pool: 0,
				UnsupportedCalls: []rpctype.SyscallReason{
					{ID: 142, Reason: "foo"},
					{ID: 2, Reason: "bar"},
					{ID: 156, Reason: "tar"},
				}}
			if err := vrf.srv.UpdateUnsupported(a, nil); err != nil {
				t.Fatalf("srv.UpdateUnsupported failed: %v", err)
			}

			if diff := cmp.Diff(test.wantPools, vrf.pools, cmp.AllowUnexported(poolInfo{})); diff != "" {
				t.Errorf("srv.pools mismatch (-want +got):\n%s", diff)
			}

			wantReasons := map[*prog.Syscall]string{
				target.SyscallMap["test$res0"]:   "foo",
				target.SyscallMap["test$union0"]: "tar",
			}
			if diff := cmp.Diff(wantReasons, vrf.reasons); diff != "" {
				t.Errorf("srv.reasons mismatch (-want +got):\n%s", diff)
			}

			if diff := cmp.Diff(test.wantCalls, vrf.calls); diff != "" {
				t.Errorf("srv.calls mismatch (-want +got):\n%s", diff)
			}

			if want, got := test.wantNotChecked, vrf.srv.notChecked; want != got {
				t.Errorf("srv.notChecked: got %d want %d", got, want)
			}

			if want, got := test.nilCT, vrf.choiceTable == nil; want != got {
				t.Errorf("vrf.choiceTable == nil: want nil, got: %v", vrf.choiceTable)
			}
		})
	}
}

func TestUpdateUnsupportedNotCalledTwice(t *testing.T) {
	vrf := Verifier{
		pools: map[int]*poolInfo{
			0: {checked: false},
			1: {checked: false},
		},
	}
	srv, err := startRPCServer(&vrf)
	if err != nil {
		t.Fatalf("failed to initialise RPC server: %v", err)
	}
	a := &rpctype.UpdateUnsupportedArgs{Pool: 0}

	if err := srv.UpdateUnsupported(a, nil); err != nil {
		t.Fatalf("srv.UpdateUnsupported failed: %v", err)
	}
	if want, got := 1, srv.notChecked; want != got {
		t.Errorf("srv.notChecked: got %d want %d", got, want)
	}

	if err := srv.UpdateUnsupported(a, nil); err != nil {
		t.Fatalf("srv.UpdateUnsupported failed: %v", err)
	}
	if want, got := 1, srv.notChecked; want != got {
		t.Fatalf("srv.UpdateUnsupported called twice")
	}

	wantPools := map[int]*poolInfo{
		0: {checked: true},
		1: {checked: false},
	}
	if diff := cmp.Diff(wantPools, vrf.pools, cmp.AllowUnexported(poolInfo{})); diff != "" {
		t.Errorf("srv.pools mismatch (-want +got):\n%s", diff)
	}
}

func TestSaveDiffResults(t *testing.T) {
	tests := []struct {
		name      string
		res       []*ExecResult
		prog      string
		wantExist bool
		wantStats *Stats
	}{
		{
			name: "report written",
			res: []*ExecResult{
				makeExecResult(0, []int{1, 3, 2}),
				makeExecResult(1, []int{1, 3, 5}),
			},
			wantExist: true,
			wantStats: &Stats{
				TotalMismatches: 1,
				Calls: map[string]*CallStats{
					"breaks_returns": makeCallStats("breaks_returns", 1, 0, map[ReturnState]bool{}),
					"test$res0":      makeCallStats("test$res0", 1, 1, map[ReturnState]bool{{Errno: 2}: true, {Errno: 5}: true}),
					"minimize$0":     makeCallStats("minimize$0", 1, 0, map[ReturnState]bool{}),
				},
			},
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			prog := getTestProgram(t)
			vrf := Verifier{
				resultsdir: makeTestResultDirectory(t),
				stats:      emptyTestStats(),
			}
			resultFile := filepath.Join(vrf.resultsdir, "result-0")

			vrf.AddCallsExecutionStat(test.res, prog)
			vrf.SaveDiffResults(test.res, prog)

			if diff := cmp.Diff(test.wantStats, vrf.stats); diff != "" {
				t.Errorf("vrf.stats mismatch (-want +got):\n%s", diff)
			}

			if got, want := osutil.IsExist(resultFile), test.wantExist; got != want {
				t.Errorf("osutil.IsExist report file: got %v want %v", got, want)
			}
			os.Remove(filepath.Join(vrf.resultsdir, "result-0"))
		})
	}
}

func TestCreateReport(t *testing.T) {
	rr := ResultReport{
		Prog: "breaks_returns()\n" +
			"minimize$0(0x1, 0x1)\n" +
			"test$res0()\n",
		Reports: []*CallReport{
			{Call: "breaks_returns", States: map[int]ReturnState{
				0: returnState(1, 1),
				1: returnState(1, 1),
				2: returnState(1, 1)}},
			{Call: "minimize$0", States: map[int]ReturnState{
				0: returnState(3, 3),
				1: returnState(3, 3),
				2: returnState(3, 3)}},
			{Call: "test$res0", States: map[int]ReturnState{
				0: returnState(2, 7),
				1: returnState(5, 3),
				2: returnState(22, 1)},
				Mismatch: true},
		},
	}
	got := string(createReport(&rr, 3))
	want := "ERRNO mismatches found for program:\n\n" +
		"[=] breaks_returns()\n" +
		"\t↳ Pool: 0, Flags: 1, Errno: 1 (operation not permitted)\n" +
		"\t↳ Pool: 1, Flags: 1, Errno: 1 (operation not permitted)\n" +
		"\t↳ Pool: 2, Flags: 1, Errno: 1 (operation not permitted)\n\n" +
		"[=] minimize$0(0x1, 0x1)\n" +
		"\t↳ Pool: 0, Flags: 3, Errno: 3 (no such process)\n" +
		"\t↳ Pool: 1, Flags: 3, Errno: 3 (no such process)\n" +
		"\t↳ Pool: 2, Flags: 3, Errno: 3 (no such process)\n\n" +
		"[!] test$res0()\n" +
		"\t↳ Pool: 0, Flags: 7, Errno: 2 (no such file or directory)\n" +
		"\t↳ Pool: 1, Flags: 3, Errno: 5 (input/output error)\n" +
		"\t↳ Pool: 2, Flags: 1, Errno: 22 (invalid argument)\n\n"
	if diff := cmp.Diff(got, want); diff != "" {
		t.Errorf("createReport: (-want +got):\n%s", diff)
	}
}
<commit_msg>syz-verifier: fix syscalls API mapping in test (#3037)<commit_after>// Copyright 2021 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package main

import (
	"bytes"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/pkg/rpctype"
	"github.com/google/syzkaller/prog"
)

func TestFinalizeCallSet(t *testing.T) {
	target, err := prog.GetTarget("test", "64")
	if err != nil {
		t.Fatalf("failed to initialise test target: %v", err)
	}

	vrf := Verifier{
		target: target,
		reasons: map[*prog.Syscall]string{
			target.SyscallMap["test$res0"]:  "foo",
			target.SyscallMap["minimize$0"]: "bar",
		},
		calls: map[*prog.Syscall]bool{
			target.SyscallMap["minimize$0"]: true,
			target.SyscallMap["test$res0"]:  true,
			target.SyscallMap["disabled1"]:  true,
		},
		reportReasons: true,
	}

	out := bytes.Buffer{}
	vrf.finalizeCallSet(&out)
	wantLines := []string{
		"The following calls have been disabled:\n",
		"\ttest$res0: foo\n",
		"\tminimize$0: bar\n",
	}
	output := out.String()
	for _, line := range wantLines {
		if !strings.Contains(output, line) {
			t.Errorf("finalizeCallSet: %q missing in reported output", line)
		}
	}

	wantCalls, gotCalls := map[*prog.Syscall]bool{
		target.SyscallMap["disabled1"]: true,
	}, vrf.calls
	if diff := cmp.Diff(wantCalls, gotCalls); diff != "" {
		t.Errorf("srv.calls mismatch (-want +got):\n%s", diff)
	}
}

func TestUpdateUnsupported(t *testing.T) {
	target, err := prog.GetTarget("test", "64")
	if err != nil {
		t.Fatalf("failed to initialise test target: %v", err)
	}

	tests := []struct {
		name           string
		vrfPools       map[int]*poolInfo
		wantPools      map[int]*poolInfo
		wantCalls      map[*prog.Syscall]bool
		wantNotChecked int
		nilCT          bool
	}{
		{
			name:           "choice table not generated",
			vrfPools:       map[int]*poolInfo{0: {}, 1: {}},
			wantPools:      map[int]*poolInfo{0: {checked: true}, 1: {}},
			wantNotChecked: 1,
			wantCalls: map[*prog.Syscall]bool{
				target.SyscallMap["minimize$0"]:     true,
				target.SyscallMap["breaks_returns"]: true,
				target.SyscallMap["test$res0"]:      true,
				target.SyscallMap["test$union0"]:    true,
			},
			nilCT: true,
		},
		{
			name:           "choice table generated",
			vrfPools:       map[int]*poolInfo{0: {}},
			wantPools:      map[int]*poolInfo{0: {checked: true}},
			wantNotChecked: 0,
			wantCalls: map[*prog.Syscall]bool{
				target.SyscallMap["minimize$0"]:     true,
				target.SyscallMap["breaks_returns"]: true,
			},
			nilCT: false,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			vrf := Verifier{
				target:        target,
				pools:         test.vrfPools,
				reasons:       make(map[*prog.Syscall]string),
				reportReasons: true,
				calls: map[*prog.Syscall]bool{
					target.SyscallMap["minimize$0"]:     true,
					target.SyscallMap["breaks_returns"]: true,
					target.SyscallMap["test$res0"]:      true,
					target.SyscallMap["test$union0"]:    true,
				},
				stats: MakeStats(),
			}
			vrf.Init()

			a := &rpctype.UpdateUnsupportedArgs{
				Pool: 0,
				UnsupportedCalls: []rpctype.SyscallReason{
					{ID: target.SyscallMap["test$res0"].ID, Reason: "foo"},
					{ID: 2, Reason: "bar"},
					{ID: target.SyscallMap["test$union0"].ID, Reason: "tar"},
				}}
			if err := vrf.srv.UpdateUnsupported(a, nil); err != nil {
				t.Fatalf("srv.UpdateUnsupported failed: %v", err)
			}

			if diff := cmp.Diff(test.wantPools, vrf.pools, cmp.AllowUnexported(poolInfo{})); diff != "" {
				t.Errorf("srv.pools mismatch (-want +got):\n%s", diff)
			}

			wantReasons := map[*prog.Syscall]string{
				target.SyscallMap["test$res0"]:   "foo",
				target.SyscallMap["test$union0"]: "tar",
			}
			if diff := cmp.Diff(wantReasons, vrf.reasons); diff != "" {
				t.Errorf("srv.reasons mismatch (-want +got):\n%s", diff)
			}

			if diff := cmp.Diff(test.wantCalls, vrf.calls); diff != "" {
				t.Errorf("srv.calls mismatch (-want +got):\n%s", diff)
			}

			if want, got := test.wantNotChecked, vrf.srv.notChecked; want != got {
				t.Errorf("srv.notChecked: got %d want %d", got, want)
			}

			if want, got := test.nilCT, vrf.choiceTable == nil; want != got {
				t.Errorf("vrf.choiceTable == nil: want nil, got: %v", vrf.choiceTable)
			}
		})
	}
}

func TestUpdateUnsupportedNotCalledTwice(t *testing.T) {
	vrf := Verifier{
		pools: map[int]*poolInfo{
			0: {checked: false},
			1: {checked: false},
		},
	}
	srv, err := startRPCServer(&vrf)
	if err != nil {
		t.Fatalf("failed to initialise RPC server: %v", err)
	}
	a := &rpctype.UpdateUnsupportedArgs{Pool: 0}

	if err := srv.UpdateUnsupported(a, nil); err != nil {
		t.Fatalf("srv.UpdateUnsupported failed: %v", err)
	}
	if want, got := 1, srv.notChecked; want != got {
		t.Errorf("srv.notChecked: got %d want %d", got, want)
	}

	if err := srv.UpdateUnsupported(a, nil); err != nil {
		t.Fatalf("srv.UpdateUnsupported failed: %v", err)
	}
	if want, got := 1, srv.notChecked; want != got {
		t.Fatalf("srv.UpdateUnsupported called twice")
	}

	wantPools := map[int]*poolInfo{
		0: {checked: true},
		1: {checked: false},
	}
	if diff := cmp.Diff(wantPools, vrf.pools, cmp.AllowUnexported(poolInfo{})); diff != "" {
		t.Errorf("srv.pools mismatch (-want +got):\n%s", diff)
	}
}

func TestSaveDiffResults(t *testing.T) {
	tests := []struct {
		name      string
		res       []*ExecResult
		prog      string
		wantExist bool
		wantStats *Stats
	}{
		{
			name: "report written",
			res: []*ExecResult{
				makeExecResult(0, []int{1, 3, 2}),
				makeExecResult(1, []int{1, 3, 5}),
			},
			wantExist: true,
			wantStats: &Stats{
				TotalMismatches: 1,
				Calls: map[string]*CallStats{
					"breaks_returns": makeCallStats("breaks_returns", 1, 0, map[ReturnState]bool{}),
					"test$res0":      makeCallStats("test$res0", 1, 1, map[ReturnState]bool{{Errno: 2}: true, {Errno: 5}: true}),
					"minimize$0":     makeCallStats("minimize$0", 1, 0, map[ReturnState]bool{}),
				},
			},
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			prog := getTestProgram(t)
			vrf := Verifier{
				resultsdir: makeTestResultDirectory(t),
				stats:      emptyTestStats(),
			}
			resultFile := filepath.Join(vrf.resultsdir, "result-0")

			vrf.AddCallsExecutionStat(test.res, prog)
			vrf.SaveDiffResults(test.res, prog)

			if diff := cmp.Diff(test.wantStats, vrf.stats); diff != "" {
				t.Errorf("vrf.stats mismatch (-want +got):\n%s", diff)
			}

			if got, want := osutil.IsExist(resultFile), test.wantExist; got != want {
				t.Errorf("osutil.IsExist report file: got %v want %v", got, want)
			}
			os.Remove(filepath.Join(vrf.resultsdir, "result-0"))
		})
	}
}

func TestCreateReport(t *testing.T) {
	rr := ResultReport{
		Prog: "breaks_returns()\n" +
			"minimize$0(0x1, 0x1)\n" +
			"test$res0()\n",
		Reports: []*CallReport{
			{Call: "breaks_returns", States: map[int]ReturnState{
				0: returnState(1, 1),
				1: returnState(1, 1),
				2: returnState(1, 1)}},
			{Call: "minimize$0", States: map[int]ReturnState{
				0: returnState(3, 3),
				1: returnState(3, 3),
				2: returnState(3, 3)}},
			{Call: "test$res0", States: map[int]ReturnState{
				0: returnState(2, 7),
				1: returnState(5, 3),
				2: returnState(22, 1)},
				Mismatch: true},
		},
	}
	got := string(createReport(&rr, 3))
	want := "ERRNO mismatches found for program:\n\n" +
		"[=] breaks_returns()\n" +
		"\t↳ Pool: 0, Flags: 1, Errno: 1 (operation not permitted)\n" +
		"\t↳ Pool: 1, Flags: 1, Errno: 1 (operation not permitted)\n" +
		"\t↳ Pool: 2, Flags: 1, Errno: 1 (operation not permitted)\n\n" +
		"[=] minimize$0(0x1, 0x1)\n" +
		"\t↳ Pool: 0, Flags: 3, Errno: 3 (no such process)\n" +
		"\t↳ Pool: 1, Flags: 3, Errno: 3 (no such process)\n" +
		"\t↳ Pool: 2, Flags: 3, Errno: 3 (no such process)\n\n" +
		"[!] test$res0()\n" +
		"\t↳ Pool: 0, Flags: 7, Errno: 2 (no such file or directory)\n" +
		"\t↳ Pool: 1, Flags: 3, Errno: 5 (input/output error)\n" +
		"\t↳ Pool: 2, Flags: 1, Errno: 22 (invalid argument)\n\n"
	if diff := cmp.Diff(got, want); diff != "" {
		t.Errorf("createReport: (-want +got):\n%s", diff)
	}
}
<|endoftext|>


<commit_before>package clickhouse

import (
	"bufio"
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
	"fmt"
	"net"
	"reflect"
	"regexp"
	"sync"
	"time"

	"github.com/kshvakov/clickhouse/lib/binary"
	"github.com/kshvakov/clickhouse/lib/column"
	"github.com/kshvakov/clickhouse/lib/data"
	"github.com/kshvakov/clickhouse/lib/protocol"
	"github.com/kshvakov/clickhouse/lib/types"
)

type (
	Date     = types.Date
	DateTime = types.DateTime
	UUID     = types.UUID
)

var (
	ErrInsertInNotBatchMode = errors.New("insert statement supported only in the batch mode (use begin/commit)")
	ErrLimitDataRequestInTx = errors.New("data request has already been prepared in transaction")
)

var (
	splitInsertRe = regexp.MustCompile(`(?i)\sVALUES\s*\(`)
)

type logger func(format string, v ...interface{})

type clickhouse struct {
	sync.Mutex
	data.ServerInfo
	data.ClientInfo
	logf          logger
	conn          *connect
	block         *data.Block
	buffer        *bufio.Writer
	decoder       *binary.Decoder
	encoder       *binary.Encoder
	compress      bool
	blockSize     int
	inTransaction bool
}

func (ch *clickhouse) Prepare(query string) (driver.Stmt, error) {
	return ch.prepareContext(context.Background(), query)
}

func (ch *clickhouse) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	return ch.prepareContext(ctx, query)
}

func (ch *clickhouse) prepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	ch.logf("[prepare] %s", query)
	switch {
	case ch.conn.closed:
		return nil, driver.ErrBadConn
	case ch.block != nil:
		return nil, ErrLimitDataRequestInTx
	case isInsert(query):
		if !ch.inTransaction {
			return nil, ErrInsertInNotBatchMode
		}
		return ch.insert(query)
	}
	return &stmt{
		ch:       ch,
		query:    query,
		numInput: numInput(query),
	}, nil
}

func (ch *clickhouse) insert(query string) (_ driver.Stmt, err error) {
	if err := ch.sendQuery(splitInsertRe.Split(query, -1)[0] + " VALUES "); err != nil {
		return nil, err
	}
	if ch.block, err = ch.readMeta(); err != nil {
		return nil, err
	}
	return &stmt{
		ch:       ch,
		isInsert: true,
	}, nil
}

func (ch *clickhouse) Begin() (driver.Tx, error) {
	return ch.beginTx(context.Background(), txOptions{})
}

func (ch *clickhouse) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
	return ch.beginTx(ctx, txOptions{
		Isolation: int(opts.Isolation),
		ReadOnly:  opts.ReadOnly,
	})
}

type txOptions struct {
	Isolation int
	ReadOnly  bool
}

func (ch *clickhouse) beginTx(ctx context.Context, opts txOptions) (*clickhouse, error) {
	ch.logf("[begin] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
	switch {
	case ch.inTransaction:
		return nil, sql.ErrTxDone
	case ch.conn.closed:
		return nil, driver.ErrBadConn
	}
	if finish := ch.watchCancel(ctx); finish != nil {
		defer finish()
	}
	ch.block = nil
	ch.inTransaction = true
	return ch, nil
}

func (ch *clickhouse) Commit() error {
	ch.logf("[commit] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
	defer func() {
		if ch.block != nil {
			ch.block.Reset()
			ch.block = nil
		}
		ch.inTransaction = false
	}()
	switch {
	case !ch.inTransaction:
		return sql.ErrTxDone
	case ch.conn.closed:
		return driver.ErrBadConn
	}
	if ch.block != nil {
		if err := ch.writeBlock(ch.block); err != nil {
			return err
		}
		// Send empty block as marker of end of data.
		if err := ch.writeBlock(&data.Block{}); err != nil {
			return err
		}
		if err := ch.encoder.Flush(); err != nil {
			return err
		}
		return ch.process()
	}
	return nil
}

func (ch *clickhouse) Rollback() error {
	ch.logf("[rollback] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
	if !ch.inTransaction {
		return sql.ErrTxDone
	}
	if ch.block != nil {
		ch.block.Reset()
	}
	ch.block = nil
	ch.buffer = nil
	ch.inTransaction = false
	return ch.conn.Close()
}

func (ch *clickhouse) CheckNamedValue(nv *driver.NamedValue) error {
	switch nv.Value.(type) {
	case column.IP, column.UUID:
		return nil
	case nil, []byte, int8, int16, int32, int64, uint8, uint16, uint32, uint64, float32, float64, string, time.Time:
		return nil
	}
	switch v := nv.Value.(type) {
	case
		[]int, []int8, []int16, []int32, []int64,
		[]uint, []uint8, []uint16, []uint32, []uint64,
		[]float32, []float64,
		[]string:
		return nil
	case net.IP:
		return nil
	case driver.Valuer:
		value, err := v.Value()
		if err != nil {
			return err
		}
		nv.Value = value
	default:
		switch value := reflect.ValueOf(nv.Value); value.Kind() {
		case reflect.Slice:
			return nil
		case reflect.Bool:
			nv.Value = uint8(0)
			if value.Bool() {
				nv.Value = uint8(1)
			}
		case reflect.Int8:
			nv.Value = int8(value.Int())
		case reflect.Int16:
			nv.Value = int16(value.Int())
		case reflect.Int32:
			nv.Value = int32(value.Int())
		case reflect.Int64:
			nv.Value = value.Int()
		case reflect.Uint8:
			nv.Value = uint8(value.Uint())
		case reflect.Uint16:
			nv.Value = uint16(value.Uint())
		case reflect.Uint32:
			nv.Value = uint32(value.Uint())
		case reflect.Uint64:
			nv.Value = uint64(value.Uint())
		case reflect.Float32:
			nv.Value = float32(value.Float())
		case reflect.Float64:
			nv.Value = float64(value.Float())
		case reflect.String:
			nv.Value = value.String()
		}
	}
	return nil
}

func (ch *clickhouse) Close() error {
	ch.block = nil
	return ch.conn.Close()
}

func (ch *clickhouse) process() error {
	packet, err := ch.decoder.Uvarint()
	if err != nil {
		return err
	}
	for {
		switch packet {
		case protocol.ServerPong:
			ch.logf("[process] <- pong")
			return nil
		case protocol.ServerException:
			ch.logf("[process] <- exception")
			return ch.exception()
		case protocol.ServerProgress:
			progress, err := ch.progress()
			if err != nil {
				return err
			}
			ch.logf("[process] <- progress: rows=%d, bytes=%d, total rows=%d",
				progress.rows,
				progress.bytes,
				progress.totalRows,
			)
		case protocol.ServerProfileInfo:
			profileInfo, err := ch.profileInfo()
			if err != nil {
				return err
			}
			ch.logf("[process] <- profiling: rows=%d, bytes=%d, blocks=%d", profileInfo.rows, profileInfo.bytes, profileInfo.blocks)
		case protocol.ServerData:
			block, err := ch.readBlock()
			if err != nil {
				return err
			}
			ch.logf("[process] <- data: packet=%d, columns=%d, rows=%d", packet, block.NumColumns, block.NumRows)
		case protocol.ServerEndOfStream:
			ch.logf("[process] <- end of stream")
			return nil
		default:
			ch.conn.Close()
			return fmt.Errorf("[process] unexpected packet [%d] from server", packet)
		}
		if packet, err = ch.decoder.Uvarint(); err != nil {
			return err
		}
	}
}

func (ch *clickhouse) cancel() error {
	ch.logf("[cancel request]")
	if err := ch.encoder.Uvarint(protocol.ClientCancel); err != nil {
		return err
	}
	return ch.conn.Close()
}

func (ch *clickhouse) watchCancel(ctx context.Context) func() {
	if done := ctx.Done(); done != nil {
		finished := make(chan struct{})
		go func() {
			select {
			case <-done:
				ch.cancel()
				finished <- struct{}{}
				ch.logf("[cancel] <- done")
			case <-finished:
				ch.logf("[cancel] <- finished")
			}
		}()
		return func() {
			select {
			case <-finished:
			case finished <- struct{}{}:
			}
		}
	}
	return func() {}
}
<commit_msg>Fix cancel not closing if write error<commit_after>package clickhouse

import (
	"bufio"
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
	"fmt"
	"net"
	"reflect"
	"regexp"
	"sync"
	"time"

	"github.com/kshvakov/clickhouse/lib/binary"
	"github.com/kshvakov/clickhouse/lib/column"
	"github.com/kshvakov/clickhouse/lib/data"
	"github.com/kshvakov/clickhouse/lib/protocol"
	"github.com/kshvakov/clickhouse/lib/types"
)

type (
	Date     = types.Date
	DateTime = types.DateTime
	UUID     = types.UUID
)

var (
	ErrInsertInNotBatchMode = errors.New("insert statement supported only in the batch mode (use begin/commit)")
	ErrLimitDataRequestInTx = errors.New("data request has already been prepared in transaction")
)

var (
	splitInsertRe = regexp.MustCompile(`(?i)\sVALUES\s*\(`)
)

type logger func(format string, v ...interface{})

type clickhouse struct {
	sync.Mutex
	data.ServerInfo
	data.ClientInfo
	logf          logger
	conn          *connect
	block         *data.Block
	buffer        *bufio.Writer
	decoder       *binary.Decoder
	encoder       *binary.Encoder
	compress      bool
	blockSize     int
	inTransaction bool
}

func (ch *clickhouse) Prepare(query string) (driver.Stmt, error) {
	return ch.prepareContext(context.Background(), query)
}

func (ch *clickhouse) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	return ch.prepareContext(ctx, query)
}

func (ch *clickhouse) prepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	ch.logf("[prepare] %s", query)
	switch {
	case ch.conn.closed:
		return nil, driver.ErrBadConn
	case ch.block != nil:
		return nil, ErrLimitDataRequestInTx
	case isInsert(query):
		if !ch.inTransaction {
			return nil, ErrInsertInNotBatchMode
		}
		return ch.insert(query)
	}
	return &stmt{
		ch:       ch,
		query:    query,
		numInput: numInput(query),
	}, nil
}

func (ch *clickhouse) insert(query string) (_ driver.Stmt, err error) {
	if err := ch.sendQuery(splitInsertRe.Split(query, -1)[0] + " VALUES "); err != nil {
		return nil, err
	}
	if ch.block, err = ch.readMeta(); err != nil {
		return nil, err
	}
	return &stmt{
		ch:       ch,
		isInsert: true,
	}, nil
}

func (ch *clickhouse) Begin() (driver.Tx, error) {
	return ch.beginTx(context.Background(), txOptions{})
}

func (ch *clickhouse) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
	return ch.beginTx(ctx, txOptions{
		Isolation: int(opts.Isolation),
		ReadOnly:  opts.ReadOnly,
	})
}

type txOptions struct {
	Isolation int
	ReadOnly  bool
}

func (ch *clickhouse) beginTx(ctx context.Context, opts txOptions) (*clickhouse, error) {
	ch.logf("[begin] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
	switch {
	case ch.inTransaction:
		return nil, sql.ErrTxDone
	case ch.conn.closed:
		return nil, driver.ErrBadConn
	}
	if finish := ch.watchCancel(ctx); finish != nil {
		defer finish()
	}
	ch.block = nil
	ch.inTransaction = true
	return ch, nil
}

func (ch *clickhouse) Commit() error {
	ch.logf("[commit] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
	defer func() {
		if ch.block != nil {
			ch.block.Reset()
			ch.block = nil
		}
		ch.inTransaction = false
	}()
	switch {
	case !ch.inTransaction:
		return sql.ErrTxDone
	case ch.conn.closed:
		return driver.ErrBadConn
	}
	if ch.block != nil {
		if err := ch.writeBlock(ch.block); err != nil {
			return err
		}
		// Send empty block as marker of end of data.
		if err := ch.writeBlock(&data.Block{}); err != nil {
			return err
		}
		if err := ch.encoder.Flush(); err != nil {
			return err
		}
		return ch.process()
	}
	return nil
}

func (ch *clickhouse) Rollback() error {
	ch.logf("[rollback] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
	if !ch.inTransaction {
		return sql.ErrTxDone
	}
	if ch.block != nil {
		ch.block.Reset()
	}
	ch.block = nil
	ch.buffer = nil
	ch.inTransaction = false
	return ch.conn.Close()
}

func (ch *clickhouse) CheckNamedValue(nv *driver.NamedValue) error {
	switch nv.Value.(type) {
	case column.IP, column.UUID:
		return nil
	case nil, []byte, int8, int16, int32, int64, uint8, uint16, uint32, uint64, float32, float64, string, time.Time:
		return nil
	}
	switch v := nv.Value.(type) {
	case
		[]int, []int8, []int16, []int32, []int64,
		[]uint, []uint8, []uint16, []uint32, []uint64,
		[]float32, []float64,
		[]string:
		return nil
	case net.IP:
		return nil
	case driver.Valuer:
		value, err := v.Value()
		if err != nil {
			return err
		}
		nv.Value = value
	default:
		switch value := reflect.ValueOf(nv.Value); value.Kind() {
		case reflect.Slice:
			return nil
		case reflect.Bool:
			nv.Value = uint8(0)
			if value.Bool() {
				nv.Value = uint8(1)
			}
		case reflect.Int8:
			nv.Value = int8(value.Int())
		case reflect.Int16:
			nv.Value = int16(value.Int())
		case reflect.Int32:
			nv.Value = int32(value.Int())
		case reflect.Int64:
			nv.Value = value.Int()
		case reflect.Uint8:
			nv.Value = uint8(value.Uint())
		case reflect.Uint16:
			nv.Value = uint16(value.Uint())
		case reflect.Uint32:
			nv.Value = uint32(value.Uint())
		case reflect.Uint64:
			nv.Value = uint64(value.Uint())
		case reflect.Float32:
			nv.Value = float32(value.Float())
		case reflect.Float64:
			nv.Value = float64(value.Float())
		case reflect.String:
			nv.Value = value.String()
		}
	}
	return nil
}

func (ch *clickhouse) Close() error {
	ch.block = nil
	return ch.conn.Close()
}

func (ch *clickhouse) process() error {
	packet, err := ch.decoder.Uvarint()
	if err != nil {
		return err
	}
	for {
		switch packet {
		case protocol.ServerPong:
			ch.logf("[process] <- pong")
			return nil
		case protocol.ServerException:
			ch.logf("[process] <- exception")
			return ch.exception()
		case protocol.ServerProgress:
			progress, err := ch.progress()
			if err != nil {
				return err
			}
			ch.logf("[process] <- progress: rows=%d, bytes=%d, total rows=%d",
				progress.rows,
				progress.bytes,
				progress.totalRows,
			)
		case protocol.ServerProfileInfo:
			profileInfo, err := ch.profileInfo()
			if err != nil {
				return err
			}
			ch.logf("[process] <- profiling: rows=%d, bytes=%d, blocks=%d", profileInfo.rows, profileInfo.bytes, profileInfo.blocks)
		case protocol.ServerData:
			block, err := ch.readBlock()
			if err != nil {
				return err
			}
			ch.logf("[process] <- data: packet=%d, columns=%d, rows=%d", packet, block.NumColumns, block.NumRows)
		case protocol.ServerEndOfStream:
			ch.logf("[process] <- end of stream")
			return nil
		default:
			ch.conn.Close()
			return fmt.Errorf("[process] unexpected packet [%d] from server", packet)
		}
		if packet, err = ch.decoder.Uvarint(); err != nil {
			return err
		}
	}
}

func (ch *clickhouse) cancel() error {
	ch.logf("[cancel request]")
	// even if we fail to write the cancel, we still need to close
	err := ch.encoder.Uvarint(protocol.ClientCancel)
	if err == nil {
		err = ch.encoder.Flush()
	}
	// return the close error if there was one, otherwise return the write error
	if cerr := ch.conn.Close(); cerr != nil {
		return cerr
	}
	return err
}

func (ch *clickhouse) watchCancel(ctx context.Context) func() {
	if done := ctx.Done(); done != nil {
		finished := make(chan struct{})
		go func() {
			select {
			case <-done:
				ch.cancel()
				finished <- struct{}{}
				ch.logf("[cancel] <- done")
			case <-finished:
				ch.logf("[cancel] <- finished")
			}
		}()
		return func() {
			select {
			case <-finished:
			case finished <- struct{}{}:
			}
		}
	}
	return func() {}
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"log"
	"strconv"

	"github.com/spf13/cobra"
)

var importModelHelp = `
model [clusterId] [modelName]
Import a model from an H2O cluster into steam
Examples:

	$ steam import model 42 model3
`

func importModel(c *context) *cobra.Command {
	cmd := newCmd(c, importModelHelp, func(c *context, args []string) {
		if len(args) != 2 {
			log.Fatalln("Incorrect number of arguments. See 'steam help import model'.")
		}

		clusterId, err := strconv.ParseInt(args[0], 10, 64)
		if err != nil {
			log.Fatalf("Incorrect value for clusterId: %s: %v", args[0], err)
		}
		modelName := args[1]

		if _, err := c.remote.GetModelFromCluster(clusterId, modelName); err != nil {
			log.Fatalln(err)
		}

		fmt.Println("Retireved model:", modelName)
	})

	return cmd
}
<commit_msg>STEAM-154 Rename GetModelFromCluster -> ImportModelFromCluster<commit_after>package cli

import (
	"fmt"
	"log"
	"strconv"

	"github.com/spf13/cobra"
)

var importModelHelp = `
model [clusterId] [modelName]
Import a model from an H2O cluster into steam
Examples:

	$ steam import model 42 model3
`

func importModel(c *context) *cobra.Command {
	cmd := newCmd(c, importModelHelp, func(c *context, args []string) {
		if len(args) != 2 {
			log.Fatalln("Incorrect number of arguments. See 'steam help import model'.")
		}

		clusterId, err := strconv.ParseInt(args[0], 10, 64)
		if err != nil {
			log.Fatalf("Incorrect value for clusterId: %s: %v", args[0], err)
		}
		modelName := args[1]

		if _, err := c.remote.ImportModelFromCluster(clusterId, modelName); err != nil {
			log.Fatalln(err)
		}

		fmt.Println("Retireved model:", modelName)
	})

	return cmd
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/mdreizin/smartling/model"
	"gopkg.in/urfave/cli.v1"
)

var listCommand = cli.Command{
	Name: "list",
	Aliases: []string{
		"ls",
	},
	Usage: "Shows a list of local translations",
	Before: func(c *cli.Context) error {
		return invokeActions([]action{
			ensureMetadataAction,
			injectContainerAction,
			injectProjectConfigAction,
		}, c)
	},
	Action: func(c *cli.Context) error {
		projectConfig := c.App.Metadata[projectConfigMetadataKey].(*model.ProjectConfig)

		for _, resource := range projectConfig.Resources {
			for _, v := range resource.Files() {
				fmt.Println(v)
			}
		}

		return nil
	},
	After: func(c *cli.Context) error {
		return invokeActions([]action{
			persistAuthTokenAction,
		}, c)
	},
}
<commit_msg>feat(listCommand): uses custom `log`<commit_after>package main

import (
	"fmt"
	"github.com/mdreizin/smartling/model"
	"gopkg.in/urfave/cli.v1"
)

var listCommand = cli.Command{
	Name: "list",
	Aliases: []string{
		"ls",
	},
	Usage: "Shows a list of local translations",
	Before: func(c *cli.Context) error {
		return invokeActions([]action{
			ensureMetadataAction,
			injectContainerAction,
			injectProjectConfigAction,
		}, c)
	},
	Action: func(c *cli.Context) error {
		projectConfig := c.App.Metadata[projectConfigMetadataKey].(*model.ProjectConfig)
		i := 0

		for _, resource := range projectConfig.Resources {
			logInfo(fmt.Sprintf("Using {PathGlob=%v PathExclude=%v}", resource.PathGlob, resource.PathExclude))

			for _, v := range resource.Files() {
				logInfo(v)
				i++
			}
		}

		logInfo(fmt.Sprintf("%d files", i))

		return nil
	},
	After: func(c *cli.Context) error {
		return invokeActions([]action{
			persistAuthTokenAction,
		}, c)
	},
}
<|endoftext|>