<commit_before>// +skip_license_check

/*
This file contains portions of code directly taken from the 'xenolf/lego' project.
A copy of the license for this code can be found in the file named LICENSE in
this directory.
*/

package util

import (
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/golang/glog"
	"github.com/miekg/dns"
)

type preCheckDNSFunc func(fqdn, value string, nameservers []string) (bool, error)

var (
	// PreCheckDNS checks DNS propagation before notifying ACME that
	// the DNS challenge is ready.
	PreCheckDNS preCheckDNSFunc = checkDNSPropagation

	fqdnToZoneLock sync.RWMutex
	fqdnToZone     = map[string]string{}
)

const defaultResolvConf = "/etc/resolv.conf"

var defaultNameservers = []string{
	"8.8.8.8:53",
	"8.8.4.4:53",
}

var RecursiveNameservers = getNameservers(defaultResolvConf, defaultNameservers)

// DNSTimeout is used to override the default DNS timeout of 10 seconds.
var DNSTimeout = 10 * time.Second

// getNameservers attempts to get systems nameservers before falling back to the defaults
func getNameservers(path string, defaults []string) []string {
	config, err := dns.ClientConfigFromFile(path)
	if err != nil || len(config.Servers) == 0 {
		return defaults
	}

	systemNameservers := []string{}
	for _, server := range config.Servers {
		// ensure all servers have a port number
		if _, _, err := net.SplitHostPort(server); err != nil {
			systemNameservers = append(systemNameservers, net.JoinHostPort(server, "53"))
		} else {
			systemNameservers = append(systemNameservers, server)
		}
	}
	return systemNameservers
}

// Update FQDN with CNAME if any
func updateDomainWithCName(r *dns.Msg, fqdn string) string {
	for _, rr := range r.Answer {
		if cn, ok := rr.(*dns.CNAME); ok {
			if cn.Hdr.Name == fqdn {
				glog.Infof("Updating FQDN: %s with it's CNAME: %s", fqdn, cn.Target)
				fqdn = cn.Target
				break
			}
		}
	}

	return fqdn
}

// checkDNSPropagation checks if the expected TXT record has been propagated to all authoritative nameservers.
func checkDNSPropagation(fqdn, value string, nameservers []string) (bool, error) {
	// Initial attempt to resolve at the recursive NS
	r, err := dnsQuery(fqdn, dns.TypeTXT, nameservers, true)
	if err != nil {
		return false, err
	}
	if r.Rcode == dns.RcodeSuccess {
		fqdn = updateDomainWithCName(r, fqdn)
	}

	authoritativeNss, err := lookupNameservers(fqdn, nameservers)
	if err != nil {
		return false, err
	}

	return checkAuthoritativeNss(fqdn, value, authoritativeNss)
}

// checkAuthoritativeNss queries each of the given nameservers for the expected TXT record.
func checkAuthoritativeNss(fqdn, value string, nameservers []string) (bool, error) {
	for _, ns := range nameservers {
		r, err := dnsQuery(fqdn, dns.TypeTXT, []string{net.JoinHostPort(ns, "53")}, false)
		if err != nil {
			return false, err
		}

		// NXDomain response is not really an error, just waiting for propagation to happen
		if !(r.Rcode == dns.RcodeSuccess || r.Rcode == dns.RcodeNameError) {
			return false, fmt.Errorf("NS %s returned %s for %s", ns, dns.RcodeToString[r.Rcode], fqdn)
		}

		glog.V(6).Infof("Looking up TXT records for %q", fqdn)
		var found bool
		for _, rr := range r.Answer {
			if txt, ok := rr.(*dns.TXT); ok {
				if strings.Join(txt.Txt, "") == value {
					found = true
					break
				}
			}
		}

		if !found {
			return false, nil
		}
	}

	return true, nil
}

// dnsQuery will query a nameserver, iterating through the supplied servers as it retries
// The nameserver should include a port, to facilitate testing where we talk to a mock dns server.
func dnsQuery(fqdn string, rtype uint16, nameservers []string, recursive bool) (in *dns.Msg, err error) {
	m := new(dns.Msg)
	m.SetQuestion(fqdn, rtype)
	m.SetEdns0(4096, false)

	if !recursive {
		m.RecursionDesired = false
	}

	// Will retry the request based on the number of servers (n+1)
	for i := 1; i <= len(nameservers)+1; i++ {
		ns := nameservers[i%len(nameservers)]
		udp := &dns.Client{Net: "udp", Timeout: DNSTimeout}
		in, _, err = udp.Exchange(m, ns)

		if err == dns.ErrTruncated {
			tcp := &dns.Client{Net: "tcp", Timeout: DNSTimeout}
			// If the TCP request succeeds, the err will reset to nil
			in, _, err = tcp.Exchange(m, ns)
		}

		if err == nil {
			break
		}
	}
	return
}

// lookupNameservers returns the authoritative nameservers for the given fqdn.
func lookupNameservers(fqdn string, nameservers []string) ([]string, error) {
	var authoritativeNss []string

	zone, err := FindZoneByFqdn(fqdn, nameservers)
	if err != nil {
		return nil, fmt.Errorf("Could not determine the zone: %v", err)
	}

	r, err := dnsQuery(zone, dns.TypeNS, nameservers, true)
	if err != nil {
		return nil, err
	}

	for _, rr := range r.Answer {
		if ns, ok := rr.(*dns.NS); ok {
			authoritativeNss = append(authoritativeNss, strings.ToLower(ns.Ns))
		}
	}

	if len(authoritativeNss) > 0 {
		return authoritativeNss, nil
	}
	return nil, fmt.Errorf("Could not determine authoritative nameservers")
}

// FindZoneByFqdn determines the zone apex for the given fqdn by recursing up the
// domain labels until the nameserver returns a SOA record in the answer section.
func FindZoneByFqdn(fqdn string, nameservers []string) (string, error) {
	fqdnToZoneLock.RLock()
	// Do we have it cached?
	if zone, ok := fqdnToZone[fqdn]; ok {
		fqdnToZoneLock.RUnlock()
		return zone, nil
	}
	fqdnToZoneLock.RUnlock()

	labelIndexes := dns.Split(fqdn)
	for _, index := range labelIndexes {
		domain := fqdn[index:]

		in, err := dnsQuery(domain, dns.TypeSOA, nameservers, true)
		if err != nil {
			return "", err
		}

		// Any response code other than NOERROR and NXDOMAIN is treated as error
		if in.Rcode != dns.RcodeNameError && in.Rcode != dns.RcodeSuccess {
			return "", fmt.Errorf("Unexpected response code '%s' for %s",
				dns.RcodeToString[in.Rcode], domain)
		}

		// Check if we got a SOA RR in the answer section
		if in.Rcode == dns.RcodeSuccess {

			// CNAME records cannot/should not exist at the root of a zone.
			// So we skip a domain when a CNAME is found.
			if dnsMsgContainsCNAME(in) {
				continue
			}

			for _, ans := range in.Answer {
				if soa, ok := ans.(*dns.SOA); ok {
					fqdnToZoneLock.Lock()
					defer fqdnToZoneLock.Unlock()

					zone := soa.Hdr.Name
					fqdnToZone[fqdn] = zone
					return zone, nil
				}
			}
		}
	}

	return "", fmt.Errorf("Could not find the start of authority")
}

// dnsMsgContainsCNAME checks for a CNAME answer in msg
func dnsMsgContainsCNAME(msg *dns.Msg) bool {
	for _, ans := range msg.Answer {
		if _, ok := ans.(*dns.CNAME); ok {
			return true
		}
	}
	return false
}

// ToFqdn converts the name into a fqdn appending a trailing dot.
func ToFqdn(name string) string {
	n := len(name)
	if n == 0 || name[n-1] == '.' {
		return name
	}
	return name + "."
}

// UnFqdn converts the fqdn into a name removing the trailing dot.
func UnFqdn(name string) string {
	n := len(name)
	if n != 0 && name[n-1] == '.' {
		return name[:n-1]
	}
	return name
}

// WaitFor polls the given function 'f', once every 'interval', up to 'timeout'.
func WaitFor(timeout, interval time.Duration, f func() (bool, error)) error {
	var lastErr string
	timeup := time.After(timeout)
	for {
		select {
		case <-timeup:
			return fmt.Errorf("Time limit exceeded. Last error: %s", lastErr)
		default:
		}

		stop, err := f()
		if stop {
			return nil
		}
		if err != nil {
			lastErr = err.Error()
		}

		time.Sleep(interval)
	}
}
<commit_msg>Retry dns queries with TCP if UDP has an i/o timeout<commit_after>// +skip_license_check

/*
This file contains portions of code directly taken from the 'xenolf/lego' project.
A copy of the license for this code can be found in the file named LICENSE in
this directory.
*/

package util

import (
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/golang/glog"
	"github.com/miekg/dns"
)

type preCheckDNSFunc func(fqdn, value string, nameservers []string) (bool, error)

var (
	// PreCheckDNS checks DNS propagation before notifying ACME that
	// the DNS challenge is ready.
	PreCheckDNS preCheckDNSFunc = checkDNSPropagation

	fqdnToZoneLock sync.RWMutex
	fqdnToZone     = map[string]string{}
)

const defaultResolvConf = "/etc/resolv.conf"

var defaultNameservers = []string{
	"8.8.8.8:53",
	"8.8.4.4:53",
}

var RecursiveNameservers = getNameservers(defaultResolvConf, defaultNameservers)

// DNSTimeout is used to override the default DNS timeout of 10 seconds.
var DNSTimeout = 10 * time.Second

// getNameservers attempts to get systems nameservers before falling back to the defaults
func getNameservers(path string, defaults []string) []string {
	config, err := dns.ClientConfigFromFile(path)
	if err != nil || len(config.Servers) == 0 {
		return defaults
	}

	systemNameservers := []string{}
	for _, server := range config.Servers {
		// ensure all servers have a port number
		if _, _, err := net.SplitHostPort(server); err != nil {
			systemNameservers = append(systemNameservers, net.JoinHostPort(server, "53"))
		} else {
			systemNameservers = append(systemNameservers, server)
		}
	}
	return systemNameservers
}

// Update FQDN with CNAME if any
func updateDomainWithCName(r *dns.Msg, fqdn string) string {
	for _, rr := range r.Answer {
		if cn, ok := rr.(*dns.CNAME); ok {
			if cn.Hdr.Name == fqdn {
				glog.Infof("Updating FQDN: %s with it's CNAME: %s", fqdn, cn.Target)
				fqdn = cn.Target
				break
			}
		}
	}

	return fqdn
}

// checkDNSPropagation checks if the expected TXT record has been propagated to all authoritative nameservers.
func checkDNSPropagation(fqdn, value string, nameservers []string) (bool, error) {
	// Initial attempt to resolve at the recursive NS
	r, err := dnsQuery(fqdn, dns.TypeTXT, nameservers, true)
	if err != nil {
		return false, err
	}
	if r.Rcode == dns.RcodeSuccess {
		fqdn = updateDomainWithCName(r, fqdn)
	}

	authoritativeNss, err := lookupNameservers(fqdn, nameservers)
	if err != nil {
		return false, err
	}

	return checkAuthoritativeNss(fqdn, value, authoritativeNss)
}

// checkAuthoritativeNss queries each of the given nameservers for the expected TXT record.
func checkAuthoritativeNss(fqdn, value string, nameservers []string) (bool, error) {
	for _, ns := range nameservers {
		r, err := dnsQuery(fqdn, dns.TypeTXT, []string{net.JoinHostPort(ns, "53")}, false)
		if err != nil {
			return false, err
		}

		// NXDomain response is not really an error, just waiting for propagation to happen
		if !(r.Rcode == dns.RcodeSuccess || r.Rcode == dns.RcodeNameError) {
			return false, fmt.Errorf("NS %s returned %s for %s", ns, dns.RcodeToString[r.Rcode], fqdn)
		}

		glog.V(6).Infof("Looking up TXT records for %q", fqdn)
		var found bool
		for _, rr := range r.Answer {
			if txt, ok := rr.(*dns.TXT); ok {
				if strings.Join(txt.Txt, "") == value {
					found = true
					break
				}
			}
		}

		if !found {
			return false, nil
		}
	}

	return true, nil
}

// dnsQuery will query a nameserver, iterating through the supplied servers as it retries
// The nameserver should include a port, to facilitate testing where we talk to a mock dns server.
func dnsQuery(fqdn string, rtype uint16, nameservers []string, recursive bool) (in *dns.Msg, err error) {
	m := new(dns.Msg)
	m.SetQuestion(fqdn, rtype)
	m.SetEdns0(4096, false)

	if !recursive {
		m.RecursionDesired = false
	}

	// Will retry the request based on the number of servers (n+1)
	for i := 1; i <= len(nameservers)+1; i++ {
		ns := nameservers[i%len(nameservers)]
		udp := &dns.Client{Net: "udp", Timeout: DNSTimeout}
		in, _, err = udp.Exchange(m, ns)

		if err == dns.ErrTruncated ||
			(err != nil && strings.HasPrefix(err.Error(), "read udp") && strings.HasSuffix(err.Error(), "i/o timeout")) {
			glog.V(6).Infof("UDP dns lookup failed, retrying with TCP: %v", err)
			tcp := &dns.Client{Net: "tcp", Timeout: DNSTimeout}
			// If the TCP request succeeds, the err will reset to nil
			in, _, err = tcp.Exchange(m, ns)
		}

		if err == nil {
			break
		}
	}
	return
}

// lookupNameservers returns the authoritative nameservers for the given fqdn.
func lookupNameservers(fqdn string, nameservers []string) ([]string, error) {
	var authoritativeNss []string

	zone, err := FindZoneByFqdn(fqdn, nameservers)
	if err != nil {
		return nil, fmt.Errorf("Could not determine the zone: %v", err)
	}

	r, err := dnsQuery(zone, dns.TypeNS, nameservers, true)
	if err != nil {
		return nil, err
	}

	for _, rr := range r.Answer {
		if ns, ok := rr.(*dns.NS); ok {
			authoritativeNss = append(authoritativeNss, strings.ToLower(ns.Ns))
		}
	}

	if len(authoritativeNss) > 0 {
		return authoritativeNss, nil
	}
	return nil, fmt.Errorf("Could not determine authoritative nameservers")
}

// FindZoneByFqdn determines the zone apex for the given fqdn by recursing up the
// domain labels until the nameserver returns a SOA record in the answer section.
func FindZoneByFqdn(fqdn string, nameservers []string) (string, error) {
	fqdnToZoneLock.RLock()
	// Do we have it cached?
	if zone, ok := fqdnToZone[fqdn]; ok {
		fqdnToZoneLock.RUnlock()
		return zone, nil
	}
	fqdnToZoneLock.RUnlock()

	labelIndexes := dns.Split(fqdn)
	for _, index := range labelIndexes {
		domain := fqdn[index:]

		in, err := dnsQuery(domain, dns.TypeSOA, nameservers, true)
		if err != nil {
			return "", err
		}

		// Any response code other than NOERROR and NXDOMAIN is treated as error
		if in.Rcode != dns.RcodeNameError && in.Rcode != dns.RcodeSuccess {
			return "", fmt.Errorf("Unexpected response code '%s' for %s",
				dns.RcodeToString[in.Rcode], domain)
		}

		// Check if we got a SOA RR in the answer section
		if in.Rcode == dns.RcodeSuccess {

			// CNAME records cannot/should not exist at the root of a zone.
			// So we skip a domain when a CNAME is found.
			if dnsMsgContainsCNAME(in) {
				continue
			}

			for _, ans := range in.Answer {
				if soa, ok := ans.(*dns.SOA); ok {
					fqdnToZoneLock.Lock()
					defer fqdnToZoneLock.Unlock()

					zone := soa.Hdr.Name
					fqdnToZone[fqdn] = zone
					return zone, nil
				}
			}
		}
	}

	return "", fmt.Errorf("Could not find the start of authority")
}

// dnsMsgContainsCNAME checks for a CNAME answer in msg
func dnsMsgContainsCNAME(msg *dns.Msg) bool {
	for _, ans := range msg.Answer {
		if _, ok := ans.(*dns.CNAME); ok {
			return true
		}
	}
	return false
}

// ToFqdn converts the name into a fqdn appending a trailing dot.
func ToFqdn(name string) string {
	n := len(name)
	if n == 0 || name[n-1] == '.' {
		return name
	}
	return name + "."
}

// UnFqdn converts the fqdn into a name removing the trailing dot.
func UnFqdn(name string) string {
	n := len(name)
	if n != 0 && name[n-1] == '.' {
		return name[:n-1]
	}
	return name
}

// WaitFor polls the given function 'f', once every 'interval', up to 'timeout'.
func WaitFor(timeout, interval time.Duration, f func() (bool, error)) error {
	var lastErr string
	timeup := time.After(timeout)
	for {
		select {
		case <-timeup:
			return fmt.Errorf("Time limit exceeded. Last error: %s", lastErr)
		default:
		}

		stop, err := f()
		if stop {
			return nil
		}
		if err != nil {
			lastErr = err.Error()
		}

		time.Sleep(interval)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The cert-manager Authors.

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 dns

import (
	"errors"
	"testing"

	cmacme "github.com/jetstack/cert-manager/pkg/apis/acme/v1"
	v1 "github.com/jetstack/cert-manager/pkg/apis/certmanager/v1"
	"github.com/jetstack/cert-manager/pkg/controller/test"
	"github.com/jetstack/cert-manager/pkg/issuer/acme/dns/acmedns"
	"github.com/jetstack/cert-manager/pkg/issuer/acme/dns/azuredns"
	"github.com/jetstack/cert-manager/pkg/issuer/acme/dns/clouddns"
	"github.com/jetstack/cert-manager/pkg/issuer/acme/dns/cloudflare"
	"github.com/jetstack/cert-manager/pkg/issuer/acme/dns/digitalocean"
	"github.com/jetstack/cert-manager/pkg/issuer/acme/dns/route53"
	"github.com/jetstack/cert-manager/pkg/issuer/acme/dns/util"
	"github.com/jetstack/cert-manager/test/unit/gen"
)

const (
	defaultTestIssuerName = "test-issuer"
)

type solverFixture struct {
	// The Solver under test
	Solver *Solver
	*test.Builder

	// Issuer to be passed to functions on the Solver (a default will be used if nil)
	Issuer v1.GenericIssuer
	// Challenge resource to use during tests
	Challenge *cmacme.Challenge

	dnsProviders *fakeDNSProviders

	// PreFn will run before the test is run, but after the fixture has been initialised.
	// This is useful if you want to load the clientset with some resources *after* the
	// fixture has been created.
	PreFn func(*testing.T, *solverFixture)
	// CheckFn should performs checks to ensure the output of the test is as expected.
	// Optional additional values may be provided, which represent the output of the
	// function under test.
	CheckFn func(*testing.T, *solverFixture, ...interface{})
	// Err should be true if an error is expected from the function under test
	Err bool

	// testResources is used to store references to resources used or created during
	// the test.
	testResources map[string]interface{}
}

func (s *solverFixture) Setup(t *testing.T) {
	if s.Issuer == nil {
		s.Issuer = gen.Issuer(defaultTestIssuerName, gen.SetIssuerACME(cmacme.ACMEIssuer{}))
	}
	if s.testResources == nil {
		s.testResources = map[string]interface{}{}
	}
	if s.Builder == nil {
		s.Builder = &test.Builder{}
	}
	if s.Builder.T == nil {
		s.Builder.T = t
	}
	if s.dnsProviders == nil {
		s.dnsProviders = newFakeDNSProviders()
	}
	s.Solver = buildFakeSolver(s.Builder, s.dnsProviders.constructors)
	if s.PreFn != nil {
		s.PreFn(t, s)
		s.Builder.Sync()
	}
}

func (s *solverFixture) Finish(t *testing.T, args ...interface{}) {
	defer s.Builder.Stop()
	// resync listers before running checks
	s.Builder.Sync()
	// run custom checks
	if s.CheckFn != nil {
		s.CheckFn(t, s, args...)
	}
}

func buildFakeSolver(b *test.Builder, dnsProviders dnsProviderConstructors) *Solver {
	b.Init()
	s := &Solver{
		Context:                 b.Context,
		secretLister:            b.Context.KubeSharedInformerFactory.Core().V1().Secrets().Lister(),
		dnsProviderConstructors: dnsProviders,
	}
	b.Start()
	return s
}

type fakeDNSProviderCall struct {
	name string
	args []interface{}
}

type fakeDNSProviders struct {
	constructors dnsProviderConstructors
	calls        []fakeDNSProviderCall
}

func (f *fakeDNSProviders) call(name string, args ...interface{}) {
	f.calls = append(f.calls, fakeDNSProviderCall{name: name, args: args})
}

func newFakeDNSProviders() *fakeDNSProviders {
	f := &fakeDNSProviders{
		calls: []fakeDNSProviderCall{},
	}
	f.constructors = dnsProviderConstructors{
		cloudDNS: func(project string, serviceAccount []byte, dns01Nameservers []string, ambient bool, hostedZoneName string) (*clouddns.DNSProvider, error) {
			f.call("clouddns", project, serviceAccount, util.RecursiveNameservers, ambient, hostedZoneName)
			return nil, nil
		},
		cloudFlare: func(email, apikey, apiToken string, dns01Nameservers []string) (*cloudflare.DNSProvider, error) {
			f.call("cloudflare", email, apikey, apiToken, util.RecursiveNameservers)
			if email == "" || (apikey == "" && apiToken == "") {
				return nil, errors.New("invalid email or apikey or apitoken")
			}
			return nil, nil
		},
		route53: func(accessKey, secretKey, hostedZoneID, region, role string, ambient bool, dns01Nameservers []string) (*route53.DNSProvider, error) {
			f.call("route53", accessKey, secretKey, hostedZoneID, region, role, ambient, util.RecursiveNameservers)
			return nil, nil
		},
		azureDNS: func(environment, clientID, clientSecret, subscriptionID, tenentID, resourceGroupName, hostedZoneName string, dns01Nameservers []string, ambient bool) (*azuredns.DNSProvider, error) {
			f.call("azuredns", clientID, clientSecret, subscriptionID, tenentID, resourceGroupName, hostedZoneName, util.RecursiveNameservers, ambient)
			return nil, nil
		},
		acmeDNS: func(host string, accountJson []byte, dns01Nameservers []string) (*acmedns.DNSProvider, error) {
			f.call("acmedns", host, accountJson, dns01Nameservers)
			return nil, nil
		},
		digitalOcean: func(token string, dns01Nameservers []string) (*digitalocean.DNSProvider, error) {
			f.call("digitalocean", token, util.RecursiveNameservers)
			return nil, nil
		},
	}
	return f
}
<commit_msg>spelling: tenant<commit_after>/*
Copyright 2020 The cert-manager Authors.

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 dns

import (
	"errors"
	"testing"

	cmacme "github.com/jetstack/cert-manager/pkg/apis/acme/v1"
	v1 "github.com/jetstack/cert-manager/pkg/apis/certmanager/v1"
	"github.com/jetstack/cert-manager/pkg/controller/test"
	"github.com/jetstack/cert-manager/pkg/issuer/acme/dns/acmedns"
	"github.com/jetstack/cert-manager/pkg/issuer/acme/dns/azuredns"
	"github.com/jetstack/cert-manager/pkg/issuer/acme/dns/clouddns"
	"github.com/jetstack/cert-manager/pkg/issuer/acme/dns/cloudflare"
	"github.com/jetstack/cert-manager/pkg/issuer/acme/dns/digitalocean"
	"github.com/jetstack/cert-manager/pkg/issuer/acme/dns/route53"
	"github.com/jetstack/cert-manager/pkg/issuer/acme/dns/util"
	"github.com/jetstack/cert-manager/test/unit/gen"
)

const (
	defaultTestIssuerName = "test-issuer"
)

type solverFixture struct {
	// The Solver under test
	Solver *Solver
	*test.Builder

	// Issuer to be passed to functions on the Solver (a default will be used if nil)
	Issuer v1.GenericIssuer
	// Challenge resource to use during tests
	Challenge *cmacme.Challenge

	dnsProviders *fakeDNSProviders

	// PreFn will run before the test is run, but after the fixture has been initialised.
	// This is useful if you want to load the clientset with some resources *after* the
	// fixture has been created.
	PreFn func(*testing.T, *solverFixture)
	// CheckFn should performs checks to ensure the output of the test is as expected.
	// Optional additional values may be provided, which represent the output of the
	// function under test.
	CheckFn func(*testing.T, *solverFixture, ...interface{})
	// Err should be true if an error is expected from the function under test
	Err bool

	// testResources is used to store references to resources used or created during
	// the test.
	testResources map[string]interface{}
}

func (s *solverFixture) Setup(t *testing.T) {
	if s.Issuer == nil {
		s.Issuer = gen.Issuer(defaultTestIssuerName, gen.SetIssuerACME(cmacme.ACMEIssuer{}))
	}
	if s.testResources == nil {
		s.testResources = map[string]interface{}{}
	}
	if s.Builder == nil {
		s.Builder = &test.Builder{}
	}
	if s.Builder.T == nil {
		s.Builder.T = t
	}
	if s.dnsProviders == nil {
		s.dnsProviders = newFakeDNSProviders()
	}
	s.Solver = buildFakeSolver(s.Builder, s.dnsProviders.constructors)
	if s.PreFn != nil {
		s.PreFn(t, s)
		s.Builder.Sync()
	}
}

func (s *solverFixture) Finish(t *testing.T, args ...interface{}) {
	defer s.Builder.Stop()
	// resync listers before running checks
	s.Builder.Sync()
	// run custom checks
	if s.CheckFn != nil {
		s.CheckFn(t, s, args...)
	}
}

func buildFakeSolver(b *test.Builder, dnsProviders dnsProviderConstructors) *Solver {
	b.Init()
	s := &Solver{
		Context:                 b.Context,
		secretLister:            b.Context.KubeSharedInformerFactory.Core().V1().Secrets().Lister(),
		dnsProviderConstructors: dnsProviders,
	}
	b.Start()
	return s
}

type fakeDNSProviderCall struct {
	name string
	args []interface{}
}

type fakeDNSProviders struct {
	constructors dnsProviderConstructors
	calls        []fakeDNSProviderCall
}

func (f *fakeDNSProviders) call(name string, args ...interface{}) {
	f.calls = append(f.calls, fakeDNSProviderCall{name: name, args: args})
}

func newFakeDNSProviders() *fakeDNSProviders {
	f := &fakeDNSProviders{
		calls: []fakeDNSProviderCall{},
	}
	f.constructors = dnsProviderConstructors{
		cloudDNS: func(project string, serviceAccount []byte, dns01Nameservers []string, ambient bool, hostedZoneName string) (*clouddns.DNSProvider, error) {
			f.call("clouddns", project, serviceAccount, util.RecursiveNameservers, ambient, hostedZoneName)
			return nil, nil
		},
		cloudFlare: func(email, apikey, apiToken string, dns01Nameservers []string) (*cloudflare.DNSProvider, error) {
			f.call("cloudflare", email, apikey, apiToken, util.RecursiveNameservers)
			if email == "" || (apikey == "" && apiToken == "") {
				return nil, errors.New("invalid email or apikey or apitoken")
			}
			return nil, nil
		},
		route53: func(accessKey, secretKey, hostedZoneID, region, role string, ambient bool, dns01Nameservers []string) (*route53.DNSProvider, error) {
			f.call("route53", accessKey, secretKey, hostedZoneID, region, role, ambient, util.RecursiveNameservers)
			return nil, nil
		},
		azureDNS: func(environment, clientID, clientSecret, subscriptionID, tenantID, resourceGroupName, hostedZoneName string, dns01Nameservers []string, ambient bool) (*azuredns.DNSProvider, error) {
			f.call("azuredns", clientID, clientSecret, subscriptionID, tenantID, resourceGroupName, hostedZoneName, util.RecursiveNameservers, ambient)
			return nil, nil
		},
		acmeDNS: func(host string, accountJson []byte, dns01Nameservers []string) (*acmedns.DNSProvider, error) {
			f.call("acmedns", host, accountJson, dns01Nameservers)
			return nil, nil
		},
		digitalOcean: func(token string, dns01Nameservers []string) (*digitalocean.DNSProvider, error) {
			f.call("digitalocean", token, util.RecursiveNameservers)
			return nil, nil
		},
	}
	return f
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Kubernetes Authors.

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 profile holds the definition of a scheduling Profile.
package profile

import (
	"errors"
	"fmt"
	"github.com/google/go-cmp/cmp"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/kubernetes/scheme"
	"k8s.io/client-go/tools/events"
	"k8s.io/kubernetes/pkg/scheduler/apis/config"
	"k8s.io/kubernetes/pkg/scheduler/framework"
	frameworkruntime "k8s.io/kubernetes/pkg/scheduler/framework/runtime"
)

// RecorderFactory builds an EventRecorder for a given scheduler name.
type RecorderFactory func(string) events.EventRecorder

// newProfile builds a Profile for the given configuration.
func newProfile(cfg config.KubeSchedulerProfile, r frameworkruntime.Registry, recorderFact RecorderFactory,
	opts ...frameworkruntime.Option) (framework.Framework, error) {
	recorder := recorderFact(cfg.SchedulerName)
	opts = append(opts, frameworkruntime.WithEventRecorder(recorder))
	fwk, err := frameworkruntime.NewFramework(r, &cfg, opts...)
	if err != nil {
		return nil, err
	}
	return fwk, nil
}

// Map holds frameworks indexed by scheduler name.
type Map map[string]framework.Framework

// NewMap builds the frameworks given by the configuration, indexed by name.
func NewMap(cfgs []config.KubeSchedulerProfile, r frameworkruntime.Registry, recorderFact RecorderFactory,
	opts ...frameworkruntime.Option) (Map, error) {
	m := make(Map)
	v := cfgValidator{m: m}

	for _, cfg := range cfgs {
		p, err := newProfile(cfg, r, recorderFact, opts...)
		if err != nil {
			return nil, fmt.Errorf("creating profile for scheduler name %s: %v", cfg.SchedulerName, err)
		}
		if err := v.validate(cfg, p); err != nil {
			return nil, err
		}
		m[cfg.SchedulerName] = p
	}
	return m, nil
}

// HandlesSchedulerName returns whether a profile handles the given scheduler name.
func (m Map) HandlesSchedulerName(name string) bool {
	_, ok := m[name]
	return ok
}

// NewRecorderFactory returns a RecorderFactory for the broadcaster.
func NewRecorderFactory(b events.EventBroadcaster) RecorderFactory {
	return func(name string) events.EventRecorder {
		return b.NewRecorder(scheme.Scheme, name)
	}
}

type cfgValidator struct {
	m             Map
	queueSort     string
	queueSortArgs runtime.Object
}

func (v *cfgValidator) validate(cfg config.KubeSchedulerProfile, f framework.Framework) error {
	if len(f.ProfileName()) == 0 {
		return errors.New("scheduler name is needed")
	}
	if cfg.Plugins == nil {
		return fmt.Errorf("plugins required for profile with scheduler name %q", f.ProfileName())
	}
	if v.m[f.ProfileName()] != nil {
		return fmt.Errorf("duplicate profile with scheduler name %q", f.ProfileName())
	}

	queueSort := f.ListPlugins().QueueSort.Enabled[0].Name
	var queueSortArgs runtime.Object
	for _, plCfg := range cfg.PluginConfig {
		if plCfg.Name == queueSort {
			queueSortArgs = plCfg.Args
		}
	}
	if len(v.queueSort) == 0 {
		v.queueSort = queueSort
		v.queueSortArgs = queueSortArgs
		return nil
	}
	if v.queueSort != queueSort {
		return fmt.Errorf("different queue sort plugins for profile %q: %q, first: %q", cfg.SchedulerName, queueSort, v.queueSort)
	}
	if !cmp.Equal(v.queueSortArgs, queueSortArgs) {
		return fmt.Errorf("different queue sort plugin args for profile %q", cfg.SchedulerName)
	}
	return nil
}
<commit_msg>cleanup: return frameworkruntime.NewFramework directly<commit_after>/*
Copyright 2020 The Kubernetes Authors.

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 profile holds the definition of a scheduling Profile.
package profile

import (
	"errors"
	"fmt"

	"github.com/google/go-cmp/cmp"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/kubernetes/scheme"
	"k8s.io/client-go/tools/events"
	"k8s.io/kubernetes/pkg/scheduler/apis/config"
	"k8s.io/kubernetes/pkg/scheduler/framework"
	frameworkruntime "k8s.io/kubernetes/pkg/scheduler/framework/runtime"
)

// RecorderFactory builds an EventRecorder for a given scheduler name.
type RecorderFactory func(string) events.EventRecorder

// newProfile builds a Profile for the given configuration.
func newProfile(cfg config.KubeSchedulerProfile, r frameworkruntime.Registry, recorderFact RecorderFactory,
	opts ...frameworkruntime.Option) (framework.Framework, error) {
	recorder := recorderFact(cfg.SchedulerName)
	opts = append(opts, frameworkruntime.WithEventRecorder(recorder))
	return frameworkruntime.NewFramework(r, &cfg, opts...)
}

// Map holds frameworks indexed by scheduler name.
type Map map[string]framework.Framework

// NewMap builds the frameworks given by the configuration, indexed by name.
func NewMap(cfgs []config.KubeSchedulerProfile, r frameworkruntime.Registry, recorderFact RecorderFactory,
	opts ...frameworkruntime.Option) (Map, error) {
	m := make(Map)
	v := cfgValidator{m: m}

	for _, cfg := range cfgs {
		p, err := newProfile(cfg, r, recorderFact, opts...)
		if err != nil {
			return nil, fmt.Errorf("creating profile for scheduler name %s: %v", cfg.SchedulerName, err)
		}
		if err := v.validate(cfg, p); err != nil {
			return nil, err
		}
		m[cfg.SchedulerName] = p
	}
	return m, nil
}

// HandlesSchedulerName returns whether a profile handles the given scheduler name.
func (m Map) HandlesSchedulerName(name string) bool {
	_, ok := m[name]
	return ok
}

// NewRecorderFactory returns a RecorderFactory for the broadcaster.
func NewRecorderFactory(b events.EventBroadcaster) RecorderFactory {
	return func(name string) events.EventRecorder {
		return b.NewRecorder(scheme.Scheme, name)
	}
}

type cfgValidator struct {
	m             Map
	queueSort     string
	queueSortArgs runtime.Object
}

func (v *cfgValidator) validate(cfg config.KubeSchedulerProfile, f framework.Framework) error {
	if len(f.ProfileName()) == 0 {
		return errors.New("scheduler name is needed")
	}
	if cfg.Plugins == nil {
		return fmt.Errorf("plugins required for profile with scheduler name %q", f.ProfileName())
	}
	if v.m[f.ProfileName()] != nil {
		return fmt.Errorf("duplicate profile with scheduler name %q", f.ProfileName())
	}

	queueSort := f.ListPlugins().QueueSort.Enabled[0].Name
	var queueSortArgs runtime.Object
	for _, plCfg := range cfg.PluginConfig {
		if plCfg.Name == queueSort {
			queueSortArgs = plCfg.Args
		}
	}
	if len(v.queueSort) == 0 {
		v.queueSort = queueSort
		v.queueSortArgs = queueSortArgs
		return nil
	}
	if v.queueSort != queueSort {
		return fmt.Errorf("different queue sort plugins for profile %q: %q, first: %q", cfg.SchedulerName, queueSort, v.queueSort)
	}
	if !cmp.Equal(v.queueSortArgs, queueSortArgs) {
		return fmt.Errorf("different queue sort plugin args for profile %q", cfg.SchedulerName)
	}
	return nil
}
<|endoftext|>


<commit_before>package stringreplace

import (
	"encoding/json"
	"fmt"
	"reflect"

	"github.com/golang/glog"
)

// VisitObjectStrings recursively visits all string fields in the object and calls the
// visitor function on them. The visitor function can be used to modify the
// value of the string fields.
func VisitObjectStrings(obj interface{}, visitor func(string) (string, bool)) error {
	return visitValue(reflect.ValueOf(obj), visitor)
}

func visitValue(v reflect.Value, visitor func(string) (string, bool)) error {
	// you'll never be able to substitute on a nil.  Check the kind first or you'll accidentally
	// end up panic-ing
	switch v.Kind() {
	case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Ptr, reflect.Slice:
		if v.IsNil() {
			return nil
		}
	}

	switch v.Kind() {

	case reflect.Ptr, reflect.Interface:
		err := visitValue(v.Elem(), visitor)
		if err != nil {
			return err
		}
	case reflect.Slice, reflect.Array:
		vt := v.Type().Elem()
		for i := 0; i < v.Len(); i++ {
			val, err := visitUnsettableValues(vt, v.Index(i), visitor)
			if err != nil {
				return err
			}
			v.Index(i).Set(val)
		}
	case reflect.Struct:
		for i := 0; i < v.NumField(); i++ {
			err := visitValue(v.Field(i), visitor)
			if err != nil {
				return err
			}
		}
	case reflect.Map:
		vt := v.Type().Elem()
		for _, oldKey := range v.MapKeys() {
			newKey, err := visitUnsettableValues(oldKey.Type(), oldKey, visitor)
			if err != nil {
				return err
			}

			oldValue := v.MapIndex(oldKey)
			newValue, err := visitUnsettableValues(vt, oldValue, visitor)
			if err != nil {
				return err
			}
			v.SetMapIndex(oldKey, reflect.Value{})
			v.SetMapIndex(newKey, newValue)
		}
	case reflect.String:
		if !v.CanSet() {
			return fmt.Errorf("unable to set String value '%v'", v)
		}
		s, asString := visitor(v.String())
		if !asString {
			return fmt.Errorf("attempted to set String field to non-string value '%v'", s)
		}
		v.SetString(s)
	default:
		glog.V(5).Infof("Unknown field type '%s': %v", v.Kind(), v)
		return nil
	}
	return nil
}

// visitUnsettableValues creates a copy of the object you want to modify and returns the modified result
func visitUnsettableValues(typeOf reflect.Type, original reflect.Value, visitor func(string) (string, bool)) (reflect.Value, error) {
	val := reflect.New(typeOf).Elem()
	existing := original
	// if the value type is interface, we must resolve it to a concrete value prior to setting it back.
	if existing.CanInterface() {
		existing = reflect.ValueOf(existing.Interface())
	}
	switch existing.Kind() {
	case reflect.String:
		s, asString := visitor(existing.String())

		if asString {
			val = reflect.ValueOf(s)
		} else {
			b := []byte(s)
			var data interface{}
			err := json.Unmarshal(b, &data)
			if err != nil {
				// the result of substitution may have been an unquoted string value,
				// which is an error when decoding in json(only "true", "false", and numeric
				// values can be unquoted), so try wrapping the value in quotes so it will be
				// properly converted to a string type during decoding.
				val = reflect.ValueOf(s)
			} else {
				val = reflect.ValueOf(data)
			}
		}

	default:
		if existing.IsValid() && existing.Kind() != reflect.Invalid {
			val.Set(existing)
		}
		visitValue(val, visitor)
	}

	return val, nil
}
<commit_msg>clarify unknown type log message<commit_after>package stringreplace

import (
	"encoding/json"
	"fmt"
	"reflect"

	"github.com/golang/glog"
)

// VisitObjectStrings recursively visits all string fields in the object and calls the
// visitor function on them. The visitor function can be used to modify the
// value of the string fields.
func VisitObjectStrings(obj interface{}, visitor func(string) (string, bool)) error {
	return visitValue(reflect.ValueOf(obj), visitor)
}

func visitValue(v reflect.Value, visitor func(string) (string, bool)) error {
	// you'll never be able to substitute on a nil.  Check the kind first or you'll accidentally
	// end up panic-ing
	switch v.Kind() {
	case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Ptr, reflect.Slice:
		if v.IsNil() {
			return nil
		}
	}

	switch v.Kind() {

	case reflect.Ptr, reflect.Interface:
		err := visitValue(v.Elem(), visitor)
		if err != nil {
			return err
		}
	case reflect.Slice, reflect.Array:
		vt := v.Type().Elem()
		for i := 0; i < v.Len(); i++ {
			val, err := visitUnsettableValues(vt, v.Index(i), visitor)
			if err != nil {
				return err
			}
			v.Index(i).Set(val)
		}
	case reflect.Struct:
		for i := 0; i < v.NumField(); i++ {
			err := visitValue(v.Field(i), visitor)
			if err != nil {
				return err
			}
		}
	case reflect.Map:
		vt := v.Type().Elem()
		for _, oldKey := range v.MapKeys() {
			newKey, err := visitUnsettableValues(oldKey.Type(), oldKey, visitor)
			if err != nil {
				return err
			}

			oldValue := v.MapIndex(oldKey)
			newValue, err := visitUnsettableValues(vt, oldValue, visitor)
			if err != nil {
				return err
			}
			v.SetMapIndex(oldKey, reflect.Value{})
			v.SetMapIndex(newKey, newValue)
		}
	case reflect.String:
		if !v.CanSet() {
			return fmt.Errorf("unable to set String value '%v'", v)
		}
		s, asString := visitor(v.String())
		if !asString {
			return fmt.Errorf("attempted to set String field to non-string value '%v'", s)
		}
		v.SetString(s)
	default:
		glog.V(5).Infof("Ignoring non-parameterizable field type '%s': %v", v.Kind(), v)
		return nil
	}
	return nil
}

// visitUnsettableValues creates a copy of the object you want to modify and returns the modified result
func visitUnsettableValues(typeOf reflect.Type, original reflect.Value, visitor func(string) (string, bool)) (reflect.Value, error) {
	val := reflect.New(typeOf).Elem()
	existing := original
	// if the value type is interface, we must resolve it to a concrete value prior to setting it back.
	if existing.CanInterface() {
		existing = reflect.ValueOf(existing.Interface())
	}
	switch existing.Kind() {
	case reflect.String:
		s, asString := visitor(existing.String())

		if asString {
			val = reflect.ValueOf(s)
		} else {
			b := []byte(s)
			var data interface{}
			err := json.Unmarshal(b, &data)
			if err != nil {
				// the result of substitution may have been an unquoted string value,
				// which is an error when decoding in json(only "true", "false", and numeric
				// values can be unquoted), so try wrapping the value in quotes so it will be
				// properly converted to a string type during decoding.
				val = reflect.ValueOf(s)
			} else {
				val = reflect.ValueOf(data)
			}
		}

	default:
		if existing.IsValid() && existing.Kind() != reflect.Invalid {
			val.Set(existing)
		}
		visitValue(val, visitor)
	}

	return val, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Daniel Oaks <daniel@danieloaks.net>
// released under the MIT license

package irc

import (
	"strings"

	"github.com/goshuirc/irc-go/ircmsg"
)

// nsHandler handles the /NS and /NICKSERV commands
func nsHandler(server *Server, client *Client, msg ircmsg.IrcMessage) bool {
	server.nickservReceivePrivmsg(client, strings.Join(msg.Params, " "))
	return false
}

func (server *Server) nickservReceiveNotice(client *Client, message string) {
	// do nothing
}

func (server *Server) nickservReceivePrivmsg(client *Client, message string) {
	client.Notice(client.t("NickServ is not yet implemented, sorry! To register an account, check /HELPOP ACC"))
}
<commit_msg>nickserv: Add NS REGISTER and NS IDENTIFY<commit_after>// Copyright (c) 2017 Daniel Oaks <daniel@danieloaks.net>
// released under the MIT license

package irc

import (
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/goshuirc/irc-go/ircfmt"
	"github.com/goshuirc/irc-go/ircmsg"
	"github.com/oragono/oragono/irc/passwd"
	"github.com/oragono/oragono/irc/sno"
	"github.com/tidwall/buntdb"
)

const nickservHelp = `NickServ lets you register and log into a user account.

To register an account:
	/NS REGISTER username [password]
Leave out [password] if you're registering using your client certificate fingerprint.

To login to an account:
	/NS IDENTIFY [username password]
Leave out [username password] to use your client certificate fingerprint. Otherwise,
the given username and password will be used.`

// nsHandler handles the /NS and /NICKSERV commands
func nsHandler(server *Server, client *Client, msg ircmsg.IrcMessage) bool {
	server.nickservReceivePrivmsg(client, strings.Join(msg.Params, " "))
	return false
}

func (server *Server) nickservReceiveNotice(client *Client, message string) {
	// do nothing
}

// extractParam extracts a parameter from the given string, returning the param and the rest of the string.
func extractParam(line string) (string, string) {
	rawParams := strings.SplitN(strings.TrimSpace(line), " ", 2)
	param0 := rawParams[0]
	var param1 string
	if 1 < len(rawParams) {
		param1 = strings.TrimSpace(rawParams[1])
	}
	return param0, param1
}

func (server *Server) nickservReceivePrivmsg(client *Client, message string) {
	command, params := extractParam(message)
	command = strings.ToLower(command)

	if command == "help" {
		for _, line := range strings.Split(nickservHelp, "\n") {
			client.Notice(line)
		}
	} else if command == "register" {
		// get params
		username, passphrase := extractParam(params)

		// fail out if we need to
		if username == "" {
			client.Notice(client.t("No username supplied"))
			return
		}

		certfp := client.certfp
		if passphrase == "" && certfp == "" {
			client.Notice(client.t("You need to either supply a passphrase or be connected via TLS with a client cert"))
			return
		}

		if !server.accountRegistration.Enabled {
			client.Notice(client.t("Account registration has been disabled"))
			return
		}

		if client.LoggedIntoAccount() {
			if server.accountRegistration.AllowMultiplePerConnection {
				client.LogoutOfAccount()
			} else {
				client.Notice(client.t("You're already logged into an account"))
				return
			}
		}

		// get and sanitise account name
		account := strings.TrimSpace(username)
		casefoldedAccount, err := CasefoldName(account)
		// probably don't need explicit check for "*" here... but let's do it anyway just to make sure
		if err != nil || username == "*" {
			client.Notice(client.t("Account name is not valid"))
			return
		}

		// check whether account exists
		// do it all in one write tx to prevent races
		err = server.store.Update(func(tx *buntdb.Tx) error {
			accountKey := fmt.Sprintf(keyAccountExists, casefoldedAccount)

			_, err := tx.Get(accountKey)
			if err != buntdb.ErrNotFound {
				//TODO(dan): if account verified key doesn't exist account is not verified, calc the maximum time without verification and expire and continue if need be
				client.Notice(client.t("Account already exists"))
				return errAccountCreation
			}

			registeredTimeKey := fmt.Sprintf(keyAccountRegTime, casefoldedAccount)

			tx.Set(accountKey, "1", nil)
			tx.Set(fmt.Sprintf(keyAccountName, casefoldedAccount), account, nil)
			tx.Set(registeredTimeKey, strconv.FormatInt(time.Now().Unix(), 10), nil)
			return nil
		})

		// account could not be created and relevant numerics have been dispatched, abort
		if err != nil {
			if err != errAccountCreation {
				client.Notice(client.t("Account registration failed"))
			}
			return
		}

		// store details
		err = server.store.Update(func(tx *buntdb.Tx) error {
			// certfp special lookup key
			if passphrase == "" {
				assembledKeyCertToAccount := fmt.Sprintf(keyCertToAccount, client.certfp)

				// make sure certfp doesn't already exist because that'd be silly
				_, err := tx.Get(assembledKeyCertToAccount)
				if err != buntdb.ErrNotFound {
					return errCertfpAlreadyExists
				}

				tx.Set(assembledKeyCertToAccount, casefoldedAccount, nil)
			}

			// make creds
			var creds AccountCredentials

			// always set passphrase salt
			creds.PassphraseSalt, err = passwd.NewSalt()
			if err != nil {
				return fmt.Errorf("Could not create passphrase salt: %s", err.Error())
			}

			if passphrase == "" {
				creds.Certificate = client.certfp
			} else {
				creds.PassphraseHash, err = server.passwords.GenerateFromPassword(creds.PassphraseSalt, passphrase)
				if err != nil {
					return fmt.Errorf("Could not hash password: %s", err)
				}
			}
			credText, err := json.Marshal(creds)
			if err != nil {
				return fmt.Errorf("Could not marshal creds: %s", err)
			}
			tx.Set(fmt.Sprintf(keyAccountCredentials, account), string(credText), nil)

			return nil
		})

		// details could not be stored and relevant numerics have been dispatched, abort
		if err != nil {
			errMsg := "Could not register"
			if err == errCertfpAlreadyExists {
				errMsg = "An account already exists for your certificate fingerprint"
			}
			client.Notice(errMsg)
			removeFailedAccRegisterData(server.store, casefoldedAccount)
			return
		}

		err = server.store.Update(func(tx *buntdb.Tx) error {
			tx.Set(fmt.Sprintf(keyAccountVerified, casefoldedAccount), "1", nil)

			// load acct info inside store tx
			account := ClientAccount{
				Name:         username,
				RegisteredAt: time.Now(),
				Clients:      []*Client{client},
			}
			//TODO(dan): Consider creating ircd-wide account adding/removing/affecting lock for protecting access to these sorts of variables
			server.accounts[casefoldedAccount] = &account
			client.account = &account

			client.Notice(client.t("Account created"))
			client.Send(nil, server.name, RPL_LOGGEDIN, client.nick, client.nickMaskString, account.Name, fmt.Sprintf(client.t("You are now logged in as %s"), account.Name))
			client.Send(nil, server.name, RPL_SASLSUCCESS, client.nick, client.t("Authentication successful"))
			server.snomasks.Send(sno.LocalAccounts, fmt.Sprintf(ircfmt.Unescape("Account registered $c[grey][$r%s$c[grey]] by $c[grey][$r%s$c[grey]]"), account.Name, client.nickMaskString))
			return nil
		})
		if err != nil {
			client.Notice(client.t("Account registration failed"))
			removeFailedAccRegisterData(server.store, casefoldedAccount)
			return
		}

	} else if command == "identify" {
		// fail out if we need to
		if !server.accountAuthenticationEnabled {
			client.Notice(client.t("Login has been disabled"))
			return
		}

		// try passphrase
		username, passphrase := extractParam(params)
		if username != "" && passphrase != "" {
			// keep it the same as in the ACC CREATE stage
			accountKey, err := CasefoldName(username)
			if err != nil {
				client.Notice(client.t("Could not login with your username/password"))
				return
			}

			// load and check acct data all in one update to prevent races.
			// as noted elsewhere, change to proper locking for Account type later probably
			var accountName string
			err = server.store.Update(func(tx *buntdb.Tx) error {
				// confirm account is verified
				_, err = tx.Get(fmt.Sprintf(keyAccountVerified, accountKey))
				if err != nil {
					return errSaslFail
				}

				creds, err := loadAccountCredentials(tx, accountKey)
				if err != nil {
					return err
				}

				// ensure creds are valid
				if len(creds.PassphraseHash) < 1 || len(creds.PassphraseSalt) < 1 || len(passphrase) < 1 {
					return errSaslFail
				}
				err = server.passwords.CompareHashAndPassword(creds.PassphraseHash, creds.PassphraseSalt, passphrase)

				// succeeded, load account info if necessary
				account, exists := server.accounts[accountKey]
				if !exists {
					account = loadAccount(server, tx, accountKey)
				}

				client.LoginToAccount(account)
				accountName = account.Name

				return err
			})

			if err == nil {
				client.Notice(fmt.Sprintf(client.t("You're now logged in as %s"), accountName))
				return
			}
		}

		// try certfp
		certfp := client.certfp
		if certfp != "" {
			var accountName string
			err := server.store.Update(func(tx *buntdb.Tx) error {
				// certfp lookup key
				accountKey, err := tx.Get(fmt.Sprintf(keyCertToAccount, certfp))
				if err != nil {
					return errSaslFail
				}

				// confirm account exists
				_, err = tx.Get(fmt.Sprintf(keyAccountExists, accountKey))
				if err != nil {
					return errSaslFail
				}

				// confirm account is verified
				_, err = tx.Get(fmt.Sprintf(keyAccountVerified, accountKey))
				if err != nil {
					return errSaslFail
				}

				// confirm the certfp in that account's credentials
				creds, err := loadAccountCredentials(tx, accountKey)
				if err != nil || creds.Certificate != client.certfp {
					return errSaslFail
				}

				// succeeded, load account info if necessary
				account, exists := server.accounts[accountKey]
				if !exists {
					account = loadAccount(server, tx, accountKey)
				}

				client.LoginToAccount(account)
				accountName = account.Name

				return nil
			})

			if err == nil {
				client.Notice(fmt.Sprintf(client.t("You're now logged in as %s"), accountName))
				return
			}
		}

		client.Notice(client.t("Could not login with your TLS certificate or supplied username/password"))
	}
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package miniprofiler_gae

import (
	"appengine"
	"appengine/memcache"
	"appengine/user"
	"appengine_internal"
	"fmt"
	"github.com/mjibson/MiniProfiler/go/miniprofiler"
	"github.com/mjibson/appstats"
	"net/http"
)

func init() {
	miniprofiler.Enable = EnableIfAdminOrDev
	miniprofiler.Get = GetMemcache
	miniprofiler.Store = StoreMemcache
	miniprofiler.MachineName = Instance
}

// EnableIfAdminOrDev returns true if this is the dev server or the current
// user is an admin. This is the default for miniprofiler.Enable.
func EnableIfAdminOrDev(r *http.Request) bool {
	if appengine.IsDevAppServer() {
		return true
	}
	c := appengine.NewContext(r)
	if u := user.Current(c); u != nil {
		return u.Admin
	}
	return false
}

// Instance returns the app engine instance id, or the hostname on dev.
// This is the default for miniprofiler.MachineName.
func Instance() string {
	if i := appengine.InstanceID(); i != "" {
		return i[len(i)-8:]
	}
	return miniprofiler.Hostname()
}

// StoreMemcache stores the Profile in memcache. This is the default for
// miniprofiler.Store.
func StoreMemcache(r *http.Request, p *miniprofiler.Profile) {
	item := &memcache.Item{
		Key:   mp_key(string(p.Id)),
		Value: p.Json(),
	}
	c := appengine.NewContext(r)
	memcache.Set(c, item)
}

// GetMemcache gets the Profile from memcache. This is the default for
// miniprofiler.Get.
func GetMemcache(r *http.Request, id string) *miniprofiler.Profile {
	c := appengine.NewContext(r)
	item, err := memcache.Get(c, mp_key(id))
	if err != nil {
		return nil
	}
	return miniprofiler.ProfileFromJson(item.Value)
}

type Context struct {
	appstats.Context
	P *miniprofiler.Profile
}

func (c Context) Call(service, method string, in, out appengine_internal.ProtoMessage, opts *appengine_internal.CallOptions) error {
	err := c.Context.Call(service, method, in, out, opts)
	v := c.Context.Stats.RPCStats[len(c.Context.Stats.RPCStats)-1]
	c.P.AddCustomTiming("RPC", &miniprofiler.CustomTiming{
		StartMilliseconds:    float64(v.Offset.Nanoseconds()) / 1000000,
		DurationMilliseconds: float64(v.Duration.Nanoseconds()) / 1000000,
	})
	return err
}

// NewHandler returns a profiled, appstats-aware appengine.Context.
func NewHandler(f func(Context, http.ResponseWriter, *http.Request)) appstats.Handler {
	return appstats.NewHandler(func(c appengine.Context, w http.ResponseWriter, r *http.Request) {
		pc := Context{
			Context: c.(appstats.Context),
		}
		pc.P = miniprofiler.NewProfile(w, r, miniprofiler.FuncName(f))
		f(pc, w, r)

		if pc.P.Root != nil {
			pc.P.CustomLink = pc.URL()
			pc.P.CustomLinkName = "appstats"
			pc.P.Finalize()
		}
	})
}

func mp_key(id string) string {
	return fmt.Sprintf("mini-profiler-results:%s", id)
}
<commit_msg>Ignore __go__ calls<commit_after>/*
 * Copyright (c) 2013 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package miniprofiler_gae

import (
	"appengine"
	"appengine/memcache"
	"appengine/user"
	"appengine_internal"
	"fmt"
	"github.com/mjibson/MiniProfiler/go/miniprofiler"
	"github.com/mjibson/appstats"
	"net/http"
)

func init() {
	miniprofiler.Enable = EnableIfAdminOrDev
	miniprofiler.Get = GetMemcache
	miniprofiler.Store = StoreMemcache
	miniprofiler.MachineName = Instance
}

// EnableIfAdminOrDev returns true if this is the dev server or the current
// user is an admin. This is the default for miniprofiler.Enable.
func EnableIfAdminOrDev(r *http.Request) bool {
	if appengine.IsDevAppServer() {
		return true
	}
	c := appengine.NewContext(r)
	if u := user.Current(c); u != nil {
		return u.Admin
	}
	return false
}

// Instance returns the app engine instance id, or the hostname on dev.
// This is the default for miniprofiler.MachineName.
func Instance() string {
	if i := appengine.InstanceID(); i != "" {
		return i[len(i)-8:]
	}
	return miniprofiler.Hostname()
}

// StoreMemcache stores the Profile in memcache. This is the default for
// miniprofiler.Store.
func StoreMemcache(r *http.Request, p *miniprofiler.Profile) {
	item := &memcache.Item{
		Key:   mp_key(string(p.Id)),
		Value: p.Json(),
	}
	c := appengine.NewContext(r)
	memcache.Set(c, item)
}

// GetMemcache gets the Profile from memcache. This is the default for
// miniprofiler.Get.
func GetMemcache(r *http.Request, id string) *miniprofiler.Profile {
	c := appengine.NewContext(r)
	item, err := memcache.Get(c, mp_key(id))
	if err != nil {
		return nil
	}
	return miniprofiler.ProfileFromJson(item.Value)
}

type Context struct {
	appstats.Context
	P *miniprofiler.Profile
}

func (c Context) Call(service, method string, in, out appengine_internal.ProtoMessage, opts *appengine_internal.CallOptions) error {
	err := c.Context.Call(service, method, in, out, opts)
	if service == "__go__" {
		v := c.Context.Stats.RPCStats[len(c.Context.Stats.RPCStats)-1]
		c.P.AddCustomTiming("RPC", &miniprofiler.CustomTiming{
			StartMilliseconds:    float64(v.Offset.Nanoseconds()) / 1000000,
			DurationMilliseconds: float64(v.Duration.Nanoseconds()) / 1000000,
		})
	}
	return err
}

// NewHandler returns a profiled, appstats-aware appengine.Context.
func NewHandler(f func(Context, http.ResponseWriter, *http.Request)) appstats.Handler {
	return appstats.NewHandler(func(c appengine.Context, w http.ResponseWriter, r *http.Request) {
		pc := Context{
			Context: c.(appstats.Context),
		}
		pc.P = miniprofiler.NewProfile(w, r, miniprofiler.FuncName(f))
		f(pc, w, r)

		if pc.P.Root != nil {
			pc.P.CustomLink = pc.URL()
			pc.P.CustomLinkName = "appstats"
			pc.P.Finalize()
		}
	})
}

func mp_key(id string) string {
	return fmt.Sprintf("mini-profiler-results:%s", id)
}
<|endoftext|>


<commit_before>package terminal

import (
	"bytes"
	"fmt"
	"koding/tools/dnode"
	"koding/tools/kite"
	"koding/tools/pty"
	"koding/tools/utils"
	"koding/virt"
	"os"
	"strconv"
	"strings"
	"syscall"
	"time"
	"unicode/utf8"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

const (
	sessionPrefix     = "koding"
	kodingScreenPath  = "/opt/koding/bin/screen"
	kodingScreenrc    = "/opt/koding/etc/screenrc"
	defaultScreenPath = "/usr/bin/screen"
)

var (
	ErrNoSession      = "ErrNoSession"
	ErrInvalidSession = "ErrInvalidSession"
)

type WebtermServer struct {
	Session          string `json:"session"`
	remote           WebtermRemote
	vm               *virt.VM
	user             *virt.User
	isForeignSession bool
	pty              *pty.PTY
	currentSecond    int64
	messageCounter   int
	byteCounter      int
	lineFeeedCounter int
	screenPath       string
}

type WebtermRemote struct {
	Output       dnode.Callback
	SessionEnded dnode.Callback
}

func webtermPing(args *dnode.Partial, channel *kite.Channel, vos *virt.VOS) (interface{}, error) {
	return "pong", nil
}

func webtermKillSession(args *dnode.Partial, channel *kite.Channel, vos *virt.VOS) (interface{}, error) {
	var params struct {
		Session string
	}

	if args == nil {
		return nil, &kite.ArgumentError{Expected: "empty argument passed"}
	}

	if args.Unmarshal(&params) != nil {
		return nil, &kite.ArgumentError{Expected: "{ session: [string]}"}
	}

	if err := killSession(vos, params.Session); err != nil {
		return nil, err
	}

	return true, nil
}

func webtermGetSessions(args *dnode.Partial, channel *kite.Channel, vos *virt.VOS) (interface{}, error) {
	sessions := screenSessions(vos)
	if len(sessions) == 0 {
		return nil, &kite.BaseError{
			Message: "No sessions available",
			CodeErr: ErrNoSession,
		}
	}

	return sessions, nil
}

// this method is special cased in oskite.go to allow foreign access
func webtermConnect(args *dnode.Partial, channel *kite.Channel, vos *virt.VOS) (interface{}, error) {
	var params struct {
		Remote       WebtermRemote
		Session      string
		SizeX, SizeY int
		Mode         string
		JoinUser     string
	}

	if args == nil {
		return nil, &kite.ArgumentError{Expected: "empty argument passed"}
	}

	if err := args.Unmarshal(&params); err != nil {
		return nil, kite.NewKiteErr(err)
	}

	if params.JoinUser != "" {
		if len(params.Session) != utils.RandomStringLength {
			return nil, &kite.BaseError{
				Message: "Invalid session identifier",
				CodeErr: ErrInvalidSession,
			}
		}

		user := new(virt.User)
		if err := mongodbConn.Run("jUsers", func(c *mgo.Collection) error {
			return c.Find(bson.M{"username": params.JoinUser}).One(user)
		}); err != nil {
			return nil, err
		}

		// vos.VM is replaced already in registerMethod via
		// channel.CorrelationName which is the remote VM hostnameAlias
		vos.User = user
	}

	screen, err := newScreen(vos, params.Mode, params.Session)
	if err != nil {
		return nil, err
	}

	server := &WebtermServer{
		Session:          screen.Session,
		remote:           params.Remote,
		vm:               vos.VM,
		user:             vos.User,
		isForeignSession: vos.User.Name != channel.Username,
		pty:              pty.New(vos.VM.PtsDir()),
		screenPath:       screen.ScreenPath,
	}

	if params.Mode != "resume" || params.Mode != "shared" {
		if params.SizeX <= 0 || params.SizeY <= 0 {
			return nil, &kite.ArgumentError{Expected: "{ sizeX: [integer], sizeY: [integer] }"}
		}

		server.SetSize(float64(params.SizeX), float64(params.SizeY))
	}

	server.pty.Slave.Chown(vos.User.Uid, -1)

	cmd := vos.VM.AttachCommand(vos.User.Uid, "/dev/pts/"+strconv.Itoa(server.pty.No), screen.Command...)
	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	go func() {
		defer log.RecoverAndLog()

		cmd.Wait()
		server.pty.Slave.Close()
		server.pty.Master.Close()
		server.remote.SessionEnded()
	}()

	go func() {
		defer log.RecoverAndLog()

		buf := make([]byte, (1<<12)-utf8.UTFMax, 1<<12)
		for {
			n, err := server.pty.Master.Read(buf)
			for n < cap(buf)-1 {
				r, _ := utf8.DecodeLastRune(buf[:n])
				if r != utf8.RuneError {
					break
				}
				server.pty.Master.Read(buf[n : n+1])
				n++
			}

			s := time.Now().Unix()
			if server.currentSecond != s {
				server.currentSecond = s
				server.messageCounter = 0
				server.byteCounter = 0
				server.lineFeeedCounter = 0
			}
			server.messageCounter += 1
			server.byteCounter += n
			server.lineFeeedCounter += bytes.Count(buf[:n], []byte{'\n'})
			if server.messageCounter > 100 || server.byteCounter > 1<<18 || server.lineFeeedCounter > 300 {
				time.Sleep(time.Second / 100)
			}

			server.remote.Output(string(utils.FilterInvalidUTF8(buf[:n])))
			if err != nil {
				break
			}
		}
	}()

	channel.OnDisconnect(func() { server.Close() })

	return server, nil
}

func (server *WebtermServer) Input(data string) {
	server.pty.Master.Write([]byte(data))
}

func (server *WebtermServer) ControlSequence(data string) {
	server.pty.MasterEncoded.Write([]byte(data))
}

func (server *WebtermServer) SetSize(x, y float64) {
	server.pty.SetSize(uint16(x), uint16(y))
}

func (server *WebtermServer) Close() error {
	server.pty.Signal(syscall.SIGHUP)
	return nil
}

func (server *WebtermServer) Terminate() error {
	server.Close()
	if !server.isForeignSession {
		server.vm.AttachCommand(server.user.Uid, "", server.screenPath, "-S", sessionPrefix+"."+server.Session, "-X", "quit").Run()
	}
	return nil
}

type screen struct {
	// Binary used for starting the screen
	ScreenPath string

	// Used for remote or multiuser mode, defines the custom session name
	Session string

	// the final command to be executed
	Command []string
}

func getScreenPath(vos *virt.VOS) (string, error) {
	screenPath := kodingScreenPath

	// it can happen that the user deleted our screen binary
	// accidently, if this happens fallback to default screen binary
	_, err := vos.Stat(kodingScreenPath)
	if err != nil {
		log.Warning("vos.Stat kodingScreenPath %v, failing over to default screen path", err)
	}

	if os.IsNotExist(err) {
		// check if the default screen binary exists too
		_, err := vos.Stat(defaultScreenPath)
		if err != nil {
			log.Error("vos.Stat defaultScreenPath %v", err)
		} else {
			log.Info("vos.Stat success defaultScreenPath found %s", defaultScreenPath)
		}

		if os.IsNotExist(err) {
			return "", &kite.BaseError{
				Message: fmt.Sprintf("neither %s nor %s does exist.", kodingScreenPath, defaultScreenPath),
				CodeErr: ErrInvalidSession,
			}
		}

		screenPath = defaultScreenPath
	}

	// debugging, remove later
	screenfile := vos.VM.File("rootfs" + defaultScreenPath)
	log.Info("DEBUG: trying os.Stat defaultScreenPath %s", screenfile)
	_, err = os.Stat(screenfile)
	if err != nil {
		log.Error("DEBUG: couldn't stat defaultScreenPath %s : %v", screenfile, err)
	} else {
		log.Info("DEBUG: result os.Stat defaultScreenPath %s", screenfile)
	}

	screenfile = vos.VM.File("rootfs" + kodingScreenPath)
	log.Info("DEBUG: trying os.Stat kodingScreenPath %s", screenfile)
	_, err = os.Stat(screenfile)
	if err != nil {
		log.Error("DEBUG: error couldn't stat %s : %v", screenfile, err)
	} else {
		log.Info("DEBUG: success os.Stat kodingScreenPath %s", screenfile)
	}

	return screenPath, nil
}

// newScreen returns a new screen instance that is used to start screen. The
// screen command line is created differently based on the incoming mode.
func newScreen(vos *virt.VOS, mode, session string) (*screen, error) {
	var screenPath string
	var err error
	attempts := 0

	// we do try several trimes to get the binary path because the VM might not
	// up immedieately.
	for {
		screenPath, err = getScreenPath(vos)
		if err == nil {
			break
		}

		// try 4 times before we hit our 15 sec timeout limit
		if attempts != 4 {
			time.Sleep(time.Second * 3) // wait a little bit ...
			attempts++
			continue
		}

		return nil, fmt.Errorf("tried five times: %s", err)
	}

	cmdArgs := []string{screenPath, "-c", kodingScreenrc, "-S"}

	// check also if our custom screenrc exists before we continue
	_, err = vos.Stat(kodingScreenrc)
	if os.IsNotExist(err) {
		log.Warning("Screenrc %s does not exist. Starting screen without screenrc.", kodingScreenrc)
		cmdArgs = []string{screenPath, "-S"}
	}

	log.Info("Mode: %s", mode)

	switch mode {
	case "shared", "resume":
		if session == "" {
			return nil, &kite.ArgumentError{Expected: "{ session: [string] }"}
		}

		if !sessionExists(vos, session) {
			return nil, &kite.BaseError{
				Message: fmt.Sprintf("The given session '%s' is not available.", session),
				CodeErr: ErrInvalidSession,
			}
		}

		cmdArgs = append(cmdArgs, sessionPrefix+"."+session)
		if mode == "shared" {
			cmdArgs = append(cmdArgs, "-x") // multiuser mode
		} else if mode == "resume" {
			cmdArgs = append(cmdArgs, "-raAd") // resume
		}
	case "noscreen":
		cmdArgs = nil
	case "create":
		session = utils.RandomString()
		cmdArgs = append(cmdArgs, sessionPrefix+"."+session)
	default:
		return nil, &kite.ArgumentError{Expected: "{ mode: [shared|noscreen|resume|create] }"}
	}

	s := &screen{
		ScreenPath: screenPath,
		Session:    session,
		Command:    cmdArgs,
	}

	return s, nil
}

// screenSessions returns a list of sessions that belongs to the given vos
// context. The sessions are in the form of ["k7sdjv12344", "askIj12sas12", ...]
func screenSessions(vos *virt.VOS) []string {
	// Do not include dead sessions in our result
	vos.VM.AttachCommand(vos.User.Uid, "", "screen", "-wipe").Output()

	// We need to use ls here, because /var/run/screen mount is only
	// visible from inside of container. Errors are ignored.
	out, _ := vos.VM.AttachCommand(vos.User.Uid, "", "ls", "/var/run/screen/S-"+vos.User.Name).Output()
	shellOut := string(bytes.TrimSpace(out))
	if shellOut == "" {
		return []string{}
	}

	names := strings.Split(shellOut, "\n")
	sessions := make([]string, len(names))

	prefix := sessionPrefix + "."
	for i, name := range names {
		segments := strings.SplitN(name, ".", 2)
		sessions[i] = strings.TrimPrefix(segments[1], prefix)
	}

	return sessions
}

// screenExists checks whether the given session exists in the running list of
// screen sessions.
func sessionExists(vos *virt.VOS, session string) bool {
	for _, s := range screenSessions(vos) {
		if s == session {
			return true
		}
	}

	return false
}

// killSession kills the given SessionID
func killSession(vos *virt.VOS, sessionID string) error {
	screenPath, err := getScreenPath(vos)
	if err != nil {
		return err
	}

	out, err := vos.VM.AttachCommand(vos.User.Uid, "", screenPath, "-X", "-S", sessionPrefix+"."+sessionID, "kill").Output()
	if err != nil {
		return commandError("screen kill failed", err, out)
	}

	return nil
}

func commandError(message string, err error, out []byte) error {
	return fmt.Errorf("%s\n%s\n%s", message, err.Error(), string(out))
}
<commit_msg>terminal: remove debug statements<commit_after>package terminal

import (
	"bytes"
	"fmt"
	"koding/tools/dnode"
	"koding/tools/kite"
	"koding/tools/pty"
	"koding/tools/utils"
	"koding/virt"
	"os"
	"strconv"
	"strings"
	"syscall"
	"time"
	"unicode/utf8"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

const (
	sessionPrefix     = "koding"
	kodingScreenPath  = "/opt/koding/bin/screen"
	kodingScreenrc    = "/opt/koding/etc/screenrc"
	defaultScreenPath = "/usr/bin/screen"
)

var (
	ErrNoSession      = "ErrNoSession"
	ErrInvalidSession = "ErrInvalidSession"
)

type WebtermServer struct {
	Session          string `json:"session"`
	remote           WebtermRemote
	vm               *virt.VM
	user             *virt.User
	isForeignSession bool
	pty              *pty.PTY
	currentSecond    int64
	messageCounter   int
	byteCounter      int
	lineFeeedCounter int
	screenPath       string
}

type WebtermRemote struct {
	Output       dnode.Callback
	SessionEnded dnode.Callback
}

func webtermPing(args *dnode.Partial, channel *kite.Channel, vos *virt.VOS) (interface{}, error) {
	return "pong", nil
}

func webtermKillSession(args *dnode.Partial, channel *kite.Channel, vos *virt.VOS) (interface{}, error) {
	var params struct {
		Session string
	}

	if args == nil {
		return nil, &kite.ArgumentError{Expected: "empty argument passed"}
	}

	if args.Unmarshal(&params) != nil {
		return nil, &kite.ArgumentError{Expected: "{ session: [string]}"}
	}

	if err := killSession(vos, params.Session); err != nil {
		return nil, err
	}

	return true, nil
}

func webtermGetSessions(args *dnode.Partial, channel *kite.Channel, vos *virt.VOS) (interface{}, error) {
	sessions := screenSessions(vos)
	if len(sessions) == 0 {
		return nil, &kite.BaseError{
			Message: "No sessions available",
			CodeErr: ErrNoSession,
		}
	}

	return sessions, nil
}

// this method is special cased in oskite.go to allow foreign access
func webtermConnect(args *dnode.Partial, channel *kite.Channel, vos *virt.VOS) (interface{}, error) {
	var params struct {
		Remote       WebtermRemote
		Session      string
		SizeX, SizeY int
		Mode         string
		JoinUser     string
	}

	if args == nil {
		return nil, &kite.ArgumentError{Expected: "empty argument passed"}
	}

	if err := args.Unmarshal(&params); err != nil {
		return nil, kite.NewKiteErr(err)
	}

	if params.JoinUser != "" {
		if len(params.Session) != utils.RandomStringLength {
			return nil, &kite.BaseError{
				Message: "Invalid session identifier",
				CodeErr: ErrInvalidSession,
			}
		}

		user := new(virt.User)
		if err := mongodbConn.Run("jUsers", func(c *mgo.Collection) error {
			return c.Find(bson.M{"username": params.JoinUser}).One(user)
		}); err != nil {
			return nil, err
		}

		// vos.VM is replaced already in registerMethod via
		// channel.CorrelationName which is the remote VM hostnameAlias
		vos.User = user
	}

	screen, err := newScreen(vos, params.Mode, params.Session)
	if err != nil {
		return nil, err
	}

	server := &WebtermServer{
		Session:          screen.Session,
		remote:           params.Remote,
		vm:               vos.VM,
		user:             vos.User,
		isForeignSession: vos.User.Name != channel.Username,
		pty:              pty.New(vos.VM.PtsDir()),
		screenPath:       screen.ScreenPath,
	}

	if params.Mode != "resume" || params.Mode != "shared" {
		if params.SizeX <= 0 || params.SizeY <= 0 {
			return nil, &kite.ArgumentError{Expected: "{ sizeX: [integer], sizeY: [integer] }"}
		}

		server.SetSize(float64(params.SizeX), float64(params.SizeY))
	}

	server.pty.Slave.Chown(vos.User.Uid, -1)

	cmd := vos.VM.AttachCommand(vos.User.Uid, "/dev/pts/"+strconv.Itoa(server.pty.No), screen.Command...)
	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	go func() {
		defer log.RecoverAndLog()

		cmd.Wait()
		server.pty.Slave.Close()
		server.pty.Master.Close()
		server.remote.SessionEnded()
	}()

	go func() {
		defer log.RecoverAndLog()

		buf := make([]byte, (1<<12)-utf8.UTFMax, 1<<12)
		for {
			n, err := server.pty.Master.Read(buf)
			for n < cap(buf)-1 {
				r, _ := utf8.DecodeLastRune(buf[:n])
				if r != utf8.RuneError {
					break
				}
				server.pty.Master.Read(buf[n : n+1])
				n++
			}

			s := time.Now().Unix()
			if server.currentSecond != s {
				server.currentSecond = s
				server.messageCounter = 0
				server.byteCounter = 0
				server.lineFeeedCounter = 0
			}
			server.messageCounter += 1
			server.byteCounter += n
			server.lineFeeedCounter += bytes.Count(buf[:n], []byte{'\n'})
			if server.messageCounter > 100 || server.byteCounter > 1<<18 || server.lineFeeedCounter > 300 {
				time.Sleep(time.Second / 100)
			}

			server.remote.Output(string(utils.FilterInvalidUTF8(buf[:n])))
			if err != nil {
				break
			}
		}
	}()

	channel.OnDisconnect(func() { server.Close() })

	return server, nil
}

func (server *WebtermServer) Input(data string) {
	server.pty.Master.Write([]byte(data))
}

func (server *WebtermServer) ControlSequence(data string) {
	server.pty.MasterEncoded.Write([]byte(data))
}

func (server *WebtermServer) SetSize(x, y float64) {
	server.pty.SetSize(uint16(x), uint16(y))
}

func (server *WebtermServer) Close() error {
	server.pty.Signal(syscall.SIGHUP)
	return nil
}

func (server *WebtermServer) Terminate() error {
	server.Close()
	if !server.isForeignSession {
		server.vm.AttachCommand(server.user.Uid, "", server.screenPath, "-S", sessionPrefix+"."+server.Session, "-X", "quit").Run()
	}
	return nil
}

type screen struct {
	// Binary used for starting the screen
	ScreenPath string

	// Used for remote or multiuser mode, defines the custom session name
	Session string

	// the final command to be executed
	Command []string
}

func getScreenPath(vos *virt.VOS) (string, error) {
	screenPath := kodingScreenPath

	// it can happen that the user deleted our screen binary
	// accidently, if this happens fallback to default screen binary
	_, err := vos.Stat(kodingScreenPath)
	if err != nil {
		log.Warning("vos.Stat kodingScreenPath %v, failing over to default screen path", err)
	}

	if os.IsNotExist(err) {
		// check if the default screen binary exists too
		_, err := vos.Stat(defaultScreenPath)
		if err != nil {
			log.Error("vos.Stat defaultScreenPath %v", err)
		} else {
			log.Info("vos.Stat success defaultScreenPath found %s", defaultScreenPath)
		}

		if os.IsNotExist(err) {
			return "", &kite.BaseError{
				Message: fmt.Sprintf("neither %s nor %s does exist.", kodingScreenPath, defaultScreenPath),
				CodeErr: ErrInvalidSession,
			}
		}

		screenPath = defaultScreenPath
	}

	return screenPath, nil
}

// newScreen returns a new screen instance that is used to start screen. The
// screen command line is created differently based on the incoming mode.
func newScreen(vos *virt.VOS, mode, session string) (*screen, error) {
	var screenPath string
	var err error
	attempts := 0

	// we do try several trimes to get the binary path because the VM might not
	// up immedieately.
	for {
		screenPath, err = getScreenPath(vos)
		if err == nil {
			break
		}

		// try 4 times before we hit our 15 sec timeout limit
		if attempts != 4 {
			time.Sleep(time.Second * 3) // wait a little bit ...
			attempts++
			continue
		}

		return nil, fmt.Errorf("tried five times: %s", err)
	}

	cmdArgs := []string{screenPath, "-c", kodingScreenrc, "-S"}

	// check also if our custom screenrc exists before we continue
	_, err = vos.Stat(kodingScreenrc)
	if os.IsNotExist(err) {
		log.Warning("Screenrc %s does not exist. Starting screen without screenrc.", kodingScreenrc)
		cmdArgs = []string{screenPath, "-S"}
	}

	log.Info("Mode: %s", mode)

	switch mode {
	case "shared", "resume":
		if session == "" {
			return nil, &kite.ArgumentError{Expected: "{ session: [string] }"}
		}

		if !sessionExists(vos, session) {
			return nil, &kite.BaseError{
				Message: fmt.Sprintf("The given session '%s' is not available.", session),
				CodeErr: ErrInvalidSession,
			}
		}

		cmdArgs = append(cmdArgs, sessionPrefix+"."+session)
		if mode == "shared" {
			cmdArgs = append(cmdArgs, "-x") // multiuser mode
		} else if mode == "resume" {
			cmdArgs = append(cmdArgs, "-raAd") // resume
		}
	case "noscreen":
		cmdArgs = nil
	case "create":
		session = utils.RandomString()
		cmdArgs = append(cmdArgs, sessionPrefix+"."+session)
	default:
		return nil, &kite.ArgumentError{Expected: "{ mode: [shared|noscreen|resume|create] }"}
	}

	s := &screen{
		ScreenPath: screenPath,
		Session:    session,
		Command:    cmdArgs,
	}

	return s, nil
}

// screenSessions returns a list of sessions that belongs to the given vos
// context. The sessions are in the form of ["k7sdjv12344", "askIj12sas12", ...]
func screenSessions(vos *virt.VOS) []string {
	// Do not include dead sessions in our result
	vos.VM.AttachCommand(vos.User.Uid, "", "screen", "-wipe").Output()

	// We need to use ls here, because /var/run/screen mount is only
	// visible from inside of container. Errors are ignored.
	out, _ := vos.VM.AttachCommand(vos.User.Uid, "", "ls", "/var/run/screen/S-"+vos.User.Name).Output()
	shellOut := string(bytes.TrimSpace(out))
	if shellOut == "" {
		return []string{}
	}

	names := strings.Split(shellOut, "\n")
	sessions := make([]string, len(names))

	prefix := sessionPrefix + "."
	for i, name := range names {
		segments := strings.SplitN(name, ".", 2)
		sessions[i] = strings.TrimPrefix(segments[1], prefix)
	}

	return sessions
}

// screenExists checks whether the given session exists in the running list of
// screen sessions.
func sessionExists(vos *virt.VOS, session string) bool {
	for _, s := range screenSessions(vos) {
		if s == session {
			return true
		}
	}

	return false
}

// killSession kills the given SessionID
func killSession(vos *virt.VOS, sessionID string) error {
	screenPath, err := getScreenPath(vos)
	if err != nil {
		return err
	}

	out, err := vos.VM.AttachCommand(vos.User.Uid, "", screenPath, "-X", "-S", sessionPrefix+"."+sessionID, "kill").Output()
	if err != nil {
		return commandError("screen kill failed", err, out)
	}

	return nil
}

func commandError(message string, err error, out []byte) error {
	return fmt.Errorf("%s\n%s\n%s", message, err.Error(), string(out))
}
<|endoftext|>