<commit_before>package beater

import (
	"sync"
	"sync/atomic"
	"time"

	"github.com/elastic/beats/filebeat/input"
	"github.com/elastic/beats/libbeat/common"
	"github.com/elastic/beats/libbeat/logp"
	"github.com/elastic/beats/libbeat/publisher"
)

type logPublisher interface {
	Start()
	Stop()
}

type syncLogPublisher struct {
	client  publisher.Client
	in, out chan []*input.FileEvent

	done chan struct{}
	wg   sync.WaitGroup
}

type asyncLogPublisher struct {
	client  publisher.Client
	in, out chan []*input.FileEvent

	// list of in-flight batches
	active  batchList
	failing bool

	done chan struct{}
	wg   sync.WaitGroup
}

// eventsBatch is used to store sorted list of actively published log lines.
// Implements `outputs.Signalerr` interface for marking batch as finished
type eventsBatch struct {
	next   *eventsBatch
	flag   int32
	events []*input.FileEvent
}

type batchList struct {
	head, tail *eventsBatch
}

const (
	defaultGCTimeout = 1 * time.Second
)

const (
	batchSuccess int32 = 1
	batchFailed  int32 = 2
)

func newPublisher(
	async bool,
	in, out chan []*input.FileEvent,
	client publisher.Client,
) logPublisher {
	if async {
		return newAsyncLogPublisher(in, out, client)
	}
	return newSyncLogPublisher(in, out, client)
}

func newSyncLogPublisher(
	in, out chan []*input.FileEvent,
	client publisher.Client,
) *syncLogPublisher {
	return &syncLogPublisher{
		in:     in,
		out:    out,
		client: client,
		done:   make(chan struct{}),
	}
}

func (p *syncLogPublisher) Start() {
	p.wg.Add(1)
	go func() {
		defer p.wg.Done()

		logp.Info("Start sending events to output")

		for {
			var events []*input.FileEvent
			select {
			case <-p.done:
				return
			case events = <-p.in:
			}

			pubEvents := make([]common.MapStr, 0, len(events))
			for _, event := range events {
				pubEvents = append(pubEvents, event.ToMapStr())
			}

			p.client.PublishEvents(pubEvents, publisher.Sync, publisher.Guaranteed)
			logp.Info("Events sent: %d", len(events))

			// Tell the registrar that we've successfully sent these events
			select {
			case <-p.done:
				return
			case p.out <- events:
			}
		}
	}()
}

func (p *syncLogPublisher) Stop() {
	close(p.done)
	p.wg.Wait()
}

func newAsyncLogPublisher(
	in, out chan []*input.FileEvent,
	client publisher.Client,
) *asyncLogPublisher {
	return &asyncLogPublisher{
		in:     in,
		out:    out,
		client: client,
		done:   make(chan struct{}),
	}
}

func (p *asyncLogPublisher) Start() {
	p.wg.Add(1)
	go func() {
		defer p.wg.Done()

		logp.Info("Start sending events to output")

		// short gc timer, in case no logs are received from spooler the queued
		// bulkEvents can still be cleaned up and forwarded to the registrar
		ticker := time.NewTicker(defaultGCTimeout)

		for {
			select {
			case <-p.done:
				return
			case events := <-p.in:
				pubEvents := make([]common.MapStr, len(events))
				for i, event := range events {
					pubEvents[i] = event.ToMapStr()
				}

				batch := &eventsBatch{
					flag:   0,
					events: events,
				}
				p.client.PublishEvents(pubEvents,
					publisher.Signal(batch), publisher.Guaranteed)

				p.active.append(batch)
			case <-ticker.C:
			}

			p.collect()
		}
	}()
}

func (p *asyncLogPublisher) Stop() {
	close(p.done)
	p.wg.Wait()
}

// collect collects finished bulk-Events in order and forward processed batches
// to registrar. Reports to registrar are guaranteed to be in same order
// as bulk-Events have been received by the spooler
func (p *asyncLogPublisher) collect() bool {
	for batch := p.active.head; batch != nil; batch = batch.next {
		state := atomic.LoadInt32(&batch.flag)
		if state == 0 && !p.failing {
			break
		}

		// remove batch from active list
		p.active.head = batch.next
		if batch.next == nil {
			p.active.tail = nil
		}

		// Batches get marked as failed, if publisher pipeline is shutting down
		// In this case we do not want to send any more batches to the registrar
		if state == batchFailed {
			p.failing = true
		}

		if p.failing {
			// if in failing state keep cleaning up queue
			continue
		}

		// Tell the registrar that we've successfully sent these events
		select {
		case <-p.done:
			return false
		case p.out <- batch.events:
		}
	}
	return true
}

func (b *eventsBatch) Completed() { atomic.StoreInt32(&b.flag, batchSuccess) }
func (b *eventsBatch) Failed()    { atomic.StoreInt32(&b.flag, batchFailed) }

func (l *batchList) append(b *eventsBatch) {
	if l.head == nil {
		l.head = b
	} else {
		l.tail.next = b
	}
	b.next = nil
	l.tail = b
}
<commit_msg>Log filebeat publish update fails (missing registrar updates)<commit_after>package beater

import (
	"sync"
	"sync/atomic"
	"time"

	"github.com/elastic/beats/filebeat/input"
	"github.com/elastic/beats/libbeat/common"
	"github.com/elastic/beats/libbeat/logp"
	"github.com/elastic/beats/libbeat/publisher"
)

type logPublisher interface {
	Start()
	Stop()
}

type syncLogPublisher struct {
	client  publisher.Client
	in, out chan []*input.FileEvent

	done chan struct{}
	wg   sync.WaitGroup
}

type asyncLogPublisher struct {
	client  publisher.Client
	in, out chan []*input.FileEvent

	// list of in-flight batches
	active  batchList
	failing bool

	done chan struct{}
	wg   sync.WaitGroup
}

// eventsBatch is used to store sorted list of actively published log lines.
// Implements `outputs.Signalerr` interface for marking batch as finished
type eventsBatch struct {
	next   *eventsBatch
	flag   int32
	events []*input.FileEvent
}

type batchList struct {
	head, tail *eventsBatch
}

const (
	defaultGCTimeout = 1 * time.Second
)

const (
	batchSuccess int32 = 1
	batchFailed  int32 = 2
)

func newPublisher(
	async bool,
	in, out chan []*input.FileEvent,
	client publisher.Client,
) logPublisher {
	if async {
		return newAsyncLogPublisher(in, out, client)
	}
	return newSyncLogPublisher(in, out, client)
}

func newSyncLogPublisher(
	in, out chan []*input.FileEvent,
	client publisher.Client,
) *syncLogPublisher {
	return &syncLogPublisher{
		in:     in,
		out:    out,
		client: client,
		done:   make(chan struct{}),
	}
}

func (p *syncLogPublisher) Start() {
	p.wg.Add(1)
	go func() {
		defer p.wg.Done()

		logp.Info("Start sending events to output")

		for {
			var events []*input.FileEvent
			select {
			case <-p.done:
				return
			case events = <-p.in:
			}

			pubEvents := make([]common.MapStr, 0, len(events))
			for _, event := range events {
				pubEvents = append(pubEvents, event.ToMapStr())
			}

			p.client.PublishEvents(pubEvents, publisher.Sync, publisher.Guaranteed)
			logp.Info("Events sent: %d", len(events))

			// Tell the registrar that we've successfully sent these events
			select {
			case <-p.done:
				return
			case p.out <- events:
			}
		}
	}()
}

func (p *syncLogPublisher) Stop() {
	close(p.done)
	p.wg.Wait()
}

func newAsyncLogPublisher(
	in, out chan []*input.FileEvent,
	client publisher.Client,
) *asyncLogPublisher {
	return &asyncLogPublisher{
		in:     in,
		out:    out,
		client: client,
		done:   make(chan struct{}),
	}
}

func (p *asyncLogPublisher) Start() {
	p.wg.Add(1)
	go func() {
		defer p.wg.Done()

		logp.Info("Start sending events to output")

		// short gc timer, in case no logs are received from spooler the queued
		// bulkEvents can still be cleaned up and forwarded to the registrar
		ticker := time.NewTicker(defaultGCTimeout)

		for {
			select {
			case <-p.done:
				return
			case events := <-p.in:
				pubEvents := make([]common.MapStr, len(events))
				for i, event := range events {
					pubEvents[i] = event.ToMapStr()
				}

				batch := &eventsBatch{
					flag:   0,
					events: events,
				}
				p.client.PublishEvents(pubEvents,
					publisher.Signal(batch), publisher.Guaranteed)

				p.active.append(batch)
			case <-ticker.C:
			}

			p.collect()
		}
	}()
}

func (p *asyncLogPublisher) Stop() {
	close(p.done)
	p.wg.Wait()
}

// collect collects finished bulk-Events in order and forward processed batches
// to registrar. Reports to registrar are guaranteed to be in same order
// as bulk-Events have been received by the spooler
func (p *asyncLogPublisher) collect() bool {
	for batch := p.active.head; batch != nil; batch = batch.next {
		state := atomic.LoadInt32(&batch.flag)
		if state == 0 && !p.failing {
			break
		}

		// remove batch from active list
		p.active.head = batch.next
		if batch.next == nil {
			p.active.tail = nil
		}

		// Batches get marked as failed, if publisher pipeline is shutting down
		// In this case we do not want to send any more batches to the registrar
		if state == batchFailed {
			p.failing = true
		}

		if p.failing {
			logp.Warn("No registrar update for potentially published batch.")

			// if in failing state keep cleaning up queue
			continue
		}

		// Tell the registrar that we've successfully sent these events
		select {
		case <-p.done:
			return false
		case p.out <- batch.events:
		}
	}
	return true
}

func (b *eventsBatch) Completed() {
	atomic.StoreInt32(&b.flag, batchSuccess)
}

func (b *eventsBatch) Failed() {
	logp.Err("Failed to publish batch. Stop updating registrar.")
	atomic.StoreInt32(&b.flag, batchFailed)
}

func (l *batchList) append(b *eventsBatch) {
	if l.head == nil {
		l.head = b
	} else {
		l.tail.next = b
	}
	b.next = nil
	l.tail = b
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess 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 tlstest contains utility methods to create test certificates.
// It is not meant to be used in production.
package tlstest

import (
	"fmt"
	"os"
	"os/exec"
	"path"

	"vitess.io/vitess/go/vt/log"
)

const (
	// CA is the name of the CA toplevel cert.
	CA = "ca"

	caConfigTemplate = `
[ ca ]
default_ca = default_ca

[ default_ca ]
database = %s
default_md = default
default_crl_days = 30

[ req ]
 default_bits           = 4096
 default_keyfile        = keyfile.pem
 distinguished_name     = req_distinguished_name
 attributes             = req_attributes
 x509_extensions       = req_ext
 prompt                 = no
 output_password        = mypass
[ req_distinguished_name ]
 C                      = US
 ST                     = California
 L                      = Mountain View
 O                      = Vitessio
 OU                     = Vitess
 CN                     = CA
 emailAddress           = test@email.address
[ req_attributes ]
 challengePassword      = A challenge password
[ req_ext ]
 basicConstraints       = CA:TRUE

`

	certConfig = `
[ req ]
 default_bits           = 4096
 default_keyfile        = keyfile.pem
 distinguished_name     = req_distinguished_name
 attributes             = req_attributes
 req_extensions        = req_ext
 prompt                 = no
 output_password        = mypass
[ req_distinguished_name ]
 C                      = US
 ST                     = California
 L                      = Mountain View
 O                      = Vitessio
 OU                     = Vitess
 CN                     = %s
 emailAddress           = test@email.address
[ req_attributes ]
 challengePassword      = A challenge password
[ req_ext ]
 basicConstraints       = CA:TRUE
 subjectAltName         = @alternate_names
 extendedKeyUsage       =serverAuth,clientAuth
[ alternate_names ]
 DNS.1                  = localhost
 DNS.2                  = 127.0.0.1
 DNS.3                  = %s
`
)

// openssl runs the openssl binary with the provided command.
func openssl(argv ...string) {
	cmd := exec.Command("openssl", argv...)
	output, err := cmd.CombinedOutput()
	if err != nil {
		if len(output) > 0 {
			log.Errorf("openssl %v returned:\n%v", argv, string(output))
		}
		log.Fatalf("openssl %v failed: %v", argv, err)
	}
}

// createKeyDBAndCAConfig creates a key database and ca config file
// for the passed in CA (possibly an intermediate CA)
func createKeyDBAndCAConfig(root, parent string) {
	databasePath := path.Join(root, parent+"-keys.db")
	if _, err := os.Stat(databasePath); os.IsNotExist(err) {
		if err := os.WriteFile(databasePath, []byte{}, os.ModePerm); err != nil {
			log.Fatalf("cannot write file %v: %v", databasePath, err)
		}
	}

	config := path.Join(root, parent+"-ca.config")
	if _, err := os.Stat(config); os.IsNotExist(err) {
		if err := os.WriteFile(config, []byte(fmt.Sprintf(caConfigTemplate, databasePath)), os.ModePerm); err != nil {
			log.Fatalf("cannot write file %v: %v", config, err)
		}
	}
}

// CreateCA creates the toplevel 'ca' certificate and key, and places it
// in the provided directory. Temporary files are also created in that
// directory.
func CreateCA(root string) {
	log.Infof("Creating test root CA in %v", root)
	key := path.Join(root, "ca-key.pem")
	cert := path.Join(root, "ca-cert.pem")
	config := path.Join(root, "ca-ca.config")
	createKeyDBAndCAConfig(root, "ca")

	openssl("genrsa", "-out", key)

	openssl("req", "-new", "-x509", "-nodes", "-days", "3600", "-batch",
		"-config", config,
		"-key", key,
		"-out", cert)
}

// CreateSignedCert creates a new certificate signed by the provided parent,
// with the provided serial number, name and common name.
// name is the file name to use. Common Name is the certificate common name.
func CreateSignedCert(root, parent, serial, name, commonName string) {
	log.Infof("Creating signed cert and key %v", commonName)
	caKey := path.Join(root, parent+"-key.pem")
	caCert := path.Join(root, parent+"-cert.pem")
	key := path.Join(root, name+"-key.pem")
	cert := path.Join(root, name+"-cert.pem")
	req := path.Join(root, name+"-req.pem")

	config := path.Join(root, name+".config")
	if err := os.WriteFile(config, []byte(fmt.Sprintf(certConfig, commonName, commonName)), os.ModePerm); err != nil {
		log.Fatalf("cannot write file %v: %v", config, err)
	}
	openssl("req", "-newkey", "rsa:2048", "-days", "3600", "-nodes",
		"-batch",
		"-config", config,
		"-keyout", key, "-out", req)
	openssl("rsa", "-in", key, "-out", key)
	openssl("x509", "-req",
		"-in", req,
		"-days", "3600",
		"-CA", caCert,
		"-CAkey", caKey,
		"-set_serial", serial,
		"-extensions", "req_ext",
		"-extfile", config,
		"-out", cert)
}

// CreateCRL creates a new empty certificate revocation list
// for the provided parent
func CreateCRL(root, parent string) {
	log.Infof("Creating CRL for root CA in %v", root)
	caKey := path.Join(root, parent+"-key.pem")
	caCert := path.Join(root, parent+"-cert.pem")
	configPath := path.Join(root, parent+"-ca.config")
	crlPath := path.Join(root, parent+"-crl.pem")
	createKeyDBAndCAConfig(root, parent)

	openssl("ca", "-gencrl",
		"-keyfile", caKey,
		"-cert", caCert,
		"-config", configPath,
		"-out",
		crlPath,
	)
}

// RevokeCertAndRegenerateCRL revokes a provided certificate under the
// provided parent CA and regenerates the CRL file for that parent
func RevokeCertAndRegenerateCRL(root, parent, name string) {
	log.Infof("Revoking certificate %s", name)
	caKey := path.Join(root, parent+"-key.pem")
	caCert := path.Join(root, parent+"-cert.pem")
	cert := path.Join(root, name+"-cert.pem")
	configPath := path.Join(root, parent+"-ca.config")
	createKeyDBAndCAConfig(root, parent)

	openssl("ca", "-revoke", cert,
		"-keyfile", caKey,
		"-cert", caCert,
		"-config", configPath,
	)

	CreateCRL(root, parent)
}

// ClientServerKeyPairs is used in tests
type ClientServerKeyPairs struct {
	ServerCert        string
	ServerKey         string
	ServerCA          string
	ServerName        string
	ServerCRL         string
	RevokedServerCert string
	RevokedServerKey  string
	RevokedServerName string
	ClientCert        string
	ClientKey         string
	ClientCA          string
	ClientCRL         string
	RevokedClientCert string
	RevokedClientKey  string
	RevokedClientName string
	CombinedCRL       string
}

var serialCounter = 0

// CreateClientServerCertPairs creates certificate pairs for use in test
func CreateClientServerCertPairs(root string) ClientServerKeyPairs {
	// Create the certs and configs.
	CreateCA(root)

	serverCASerial := fmt.Sprintf("%03d", serialCounter*2+1)
	serverSerial := fmt.Sprintf("%03d", serialCounter*2+3)
	revokedServerSerial := fmt.Sprintf("%03d", serialCounter*2+5)
	clientCASerial := fmt.Sprintf("%03d", serialCounter*2+2)
	clientCertSerial := fmt.Sprintf("%03d", serialCounter*2+4)
	revokedClientSerial := fmt.Sprintf("%03d", serialCounter*2+6)

	serialCounter = serialCounter + 3

	serverCAName := fmt.Sprintf("servers-ca-%s", serverCASerial)
	serverCACommonName := fmt.Sprintf("Servers %s CA", serverCASerial)
	serverCertName := fmt.Sprintf("server-instance-%s", serverSerial)
	serverCertCommonName := fmt.Sprintf("server%s.example.com", serverSerial)
	revokedServerCertName := fmt.Sprintf("server-instance-%s", revokedServerSerial)
	revokedServerCertCommonName := fmt.Sprintf("server%s.example.com", revokedServerSerial)

	clientCAName := fmt.Sprintf("clients-ca-%s", clientCASerial)
	clientCACommonName := fmt.Sprintf("Clients %s CA", clientCASerial)
	clientCertName := fmt.Sprintf("client-instance-%s", clientCertSerial)
	clientCertCommonName := fmt.Sprintf("client%s.example.com", clientCertSerial)
	revokedClientCertName := fmt.Sprintf("client-instance-%s", revokedClientSerial)
	revokedClientCertCommonName := fmt.Sprintf("client%s.example.com", revokedClientSerial)

	CreateSignedCert(root, CA, serverCASerial, serverCAName, serverCACommonName)
	CreateSignedCert(root, serverCAName, serverSerial, serverCertName, serverCertCommonName)
	CreateSignedCert(root, serverCAName, revokedServerSerial, revokedServerCertName, revokedServerCertCommonName)
	RevokeCertAndRegenerateCRL(root, serverCAName, revokedServerCertName)

	CreateSignedCert(root, CA, clientCASerial, clientCAName, clientCACommonName)
	CreateSignedCert(root, clientCAName, clientCertSerial, clientCertName, clientCertCommonName)
	CreateSignedCert(root, clientCAName, revokedClientSerial, revokedClientCertName, revokedClientCertCommonName)
	RevokeCertAndRegenerateCRL(root, clientCAName, revokedClientCertName)

	serverCRLPath := path.Join(root, fmt.Sprintf("%s-crl.pem", serverCAName))
	clientCRLPath := path.Join(root, fmt.Sprintf("%s-crl.pem", clientCAName))
	combinedCRLPath := path.Join(root, fmt.Sprintf("%s-%s-combined-crl.pem", serverCAName, clientCAName))

	serverCRLBytes, err := os.ReadFile(serverCRLPath)
	if err != nil {
		log.Fatalf("Could not read server CRL file")
	}

	clientCRLBytes, err := os.ReadFile(clientCRLPath)
	if err != nil {
		log.Fatalf("Could not read client CRL file")
	}

	err = os.WriteFile(combinedCRLPath, append(serverCRLBytes, clientCRLBytes...), 0777)
	if err != nil {
		log.Fatalf("Could not write combined CRL file")
	}

	return ClientServerKeyPairs{
		ServerCert:        path.Join(root, fmt.Sprintf("%s-cert.pem", serverCertName)),
		ServerKey:         path.Join(root, fmt.Sprintf("%s-key.pem", serverCertName)),
		ServerCA:          path.Join(root, fmt.Sprintf("%s-cert.pem", serverCAName)),
		ServerCRL:         serverCRLPath,
		RevokedServerCert: path.Join(root, fmt.Sprintf("%s-cert.pem", revokedServerCertName)),
		RevokedServerKey:  path.Join(root, fmt.Sprintf("%s-key.pem", revokedServerCertName)),
		ClientCert:        path.Join(root, fmt.Sprintf("%s-cert.pem", clientCertName)),
		ClientKey:         path.Join(root, fmt.Sprintf("%s-key.pem", clientCertName)),
		ClientCA:          path.Join(root, fmt.Sprintf("%s-cert.pem", clientCAName)),
		ClientCRL:         clientCRLPath,
		RevokedClientCert: path.Join(root, fmt.Sprintf("%s-cert.pem", revokedClientCertName)),
		RevokedClientKey:  path.Join(root, fmt.Sprintf("%s-key.pem", revokedClientCertName)),
		CombinedCRL:       combinedCRLPath,
		ServerName:        serverCertCommonName,
		RevokedServerName: revokedServerCertCommonName,
		RevokedClientName: revokedClientCertCommonName,
	}
}
<commit_msg>tlstest: use sha256 for signing<commit_after>/*
Copyright 2019 The Vitess 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 tlstest contains utility methods to create test certificates.
// It is not meant to be used in production.
package tlstest

import (
	"fmt"
	"os"
	"os/exec"
	"path"

	"vitess.io/vitess/go/vt/log"
)

const (
	// CA is the name of the CA toplevel cert.
	CA = "ca"

	caConfigTemplate = `
[ ca ]
default_ca = default_ca

[ default_ca ]
database = %s
default_md = default
default_crl_days = 30

[ req ]
 default_bits           = 4096
 default_keyfile        = keyfile.pem
 distinguished_name     = req_distinguished_name
 attributes             = req_attributes
 x509_extensions       = req_ext
 prompt                 = no
 output_password        = mypass
[ req_distinguished_name ]
 C                      = US
 ST                     = California
 L                      = Mountain View
 O                      = Vitessio
 OU                     = Vitess
 CN                     = CA
 emailAddress           = test@email.address
[ req_attributes ]
 challengePassword      = A challenge password
[ req_ext ]
 basicConstraints       = CA:TRUE

`

	certConfig = `
[ req ]
 default_bits           = 4096
 default_keyfile        = keyfile.pem
 distinguished_name     = req_distinguished_name
 attributes             = req_attributes
 req_extensions        = req_ext
 prompt                 = no
 output_password        = mypass
[ req_distinguished_name ]
 C                      = US
 ST                     = California
 L                      = Mountain View
 O                      = Vitessio
 OU                     = Vitess
 CN                     = %s
 emailAddress           = test@email.address
[ req_attributes ]
 challengePassword      = A challenge password
[ req_ext ]
 basicConstraints       = CA:TRUE
 subjectAltName         = @alternate_names
 extendedKeyUsage       =serverAuth,clientAuth
[ alternate_names ]
 DNS.1                  = localhost
 DNS.2                  = 127.0.0.1
 DNS.3                  = %s
`
)

// openssl runs the openssl binary with the provided command.
func openssl(argv ...string) {
	cmd := exec.Command("openssl", argv...)
	output, err := cmd.CombinedOutput()
	if err != nil {
		if len(output) > 0 {
			log.Errorf("openssl %v returned:\n%v", argv, string(output))
		}
		log.Fatalf("openssl %v failed: %v", argv, err)
	}
}

// createKeyDBAndCAConfig creates a key database and ca config file
// for the passed in CA (possibly an intermediate CA)
func createKeyDBAndCAConfig(root, parent string) {
	databasePath := path.Join(root, parent+"-keys.db")
	if _, err := os.Stat(databasePath); os.IsNotExist(err) {
		if err := os.WriteFile(databasePath, []byte{}, os.ModePerm); err != nil {
			log.Fatalf("cannot write file %v: %v", databasePath, err)
		}
	}

	config := path.Join(root, parent+"-ca.config")
	if _, err := os.Stat(config); os.IsNotExist(err) {
		if err := os.WriteFile(config, []byte(fmt.Sprintf(caConfigTemplate, databasePath)), os.ModePerm); err != nil {
			log.Fatalf("cannot write file %v: %v", config, err)
		}
	}
}

// CreateCA creates the toplevel 'ca' certificate and key, and places it
// in the provided directory. Temporary files are also created in that
// directory.
func CreateCA(root string) {
	log.Infof("Creating test root CA in %v", root)
	key := path.Join(root, "ca-key.pem")
	cert := path.Join(root, "ca-cert.pem")
	config := path.Join(root, "ca-ca.config")
	createKeyDBAndCAConfig(root, "ca")

	openssl("genrsa", "-out", key)

	openssl("req",
		"-new",
		"-x509",
		"-days", "3600",
		"-nodes",
		"-batch",
		"-sha256",
		"-config", config,
		"-key", key,
		"-out", cert)
}

// CreateSignedCert creates a new certificate signed by the provided parent,
// with the provided serial number, name and common name.
// name is the file name to use. Common Name is the certificate common name.
func CreateSignedCert(root, parent, serial, name, commonName string) {
	log.Infof("Creating signed cert and key %v", commonName)
	caKey := path.Join(root, parent+"-key.pem")
	caCert := path.Join(root, parent+"-cert.pem")
	key := path.Join(root, name+"-key.pem")
	cert := path.Join(root, name+"-cert.pem")
	req := path.Join(root, name+"-req.pem")

	config := path.Join(root, name+".config")
	if err := os.WriteFile(config, []byte(fmt.Sprintf(certConfig, commonName, commonName)), os.ModePerm); err != nil {
		log.Fatalf("cannot write file %v: %v", config, err)
	}
	openssl("req",
		"-newkey", "rsa:2048",
		"-days", "3600",
		"-nodes",
		"-batch",
		"-sha256",
		"-config", config,
		"-keyout", key,
		"-out", req)
	openssl("rsa", "-in", key, "-out", key)
	openssl("x509", "-req",
		"-in", req,
		"-days", "3600",
		"-CA", caCert,
		"-CAkey", caKey,
		"-set_serial", serial,
		"-extensions", "req_ext",
		"-extfile", config,
		"-out", cert)
}

// CreateCRL creates a new empty certificate revocation list
// for the provided parent
func CreateCRL(root, parent string) {
	log.Infof("Creating CRL for root CA in %v", root)
	caKey := path.Join(root, parent+"-key.pem")
	caCert := path.Join(root, parent+"-cert.pem")
	configPath := path.Join(root, parent+"-ca.config")
	crlPath := path.Join(root, parent+"-crl.pem")
	createKeyDBAndCAConfig(root, parent)

	openssl("ca", "-gencrl",
		"-keyfile", caKey,
		"-cert", caCert,
		"-config", configPath,
		"-out",
		crlPath,
	)
}

// RevokeCertAndRegenerateCRL revokes a provided certificate under the
// provided parent CA and regenerates the CRL file for that parent
func RevokeCertAndRegenerateCRL(root, parent, name string) {
	log.Infof("Revoking certificate %s", name)
	caKey := path.Join(root, parent+"-key.pem")
	caCert := path.Join(root, parent+"-cert.pem")
	cert := path.Join(root, name+"-cert.pem")
	configPath := path.Join(root, parent+"-ca.config")
	createKeyDBAndCAConfig(root, parent)

	openssl("ca", "-revoke", cert,
		"-keyfile", caKey,
		"-cert", caCert,
		"-config", configPath,
	)

	CreateCRL(root, parent)
}

// ClientServerKeyPairs is used in tests
type ClientServerKeyPairs struct {
	ServerCert        string
	ServerKey         string
	ServerCA          string
	ServerName        string
	ServerCRL         string
	RevokedServerCert string
	RevokedServerKey  string
	RevokedServerName string
	ClientCert        string
	ClientKey         string
	ClientCA          string
	ClientCRL         string
	RevokedClientCert string
	RevokedClientKey  string
	RevokedClientName string
	CombinedCRL       string
}

var serialCounter = 0

// CreateClientServerCertPairs creates certificate pairs for use in test
func CreateClientServerCertPairs(root string) ClientServerKeyPairs {
	// Create the certs and configs.
	CreateCA(root)

	serverCASerial := fmt.Sprintf("%03d", serialCounter*2+1)
	serverSerial := fmt.Sprintf("%03d", serialCounter*2+3)
	revokedServerSerial := fmt.Sprintf("%03d", serialCounter*2+5)
	clientCASerial := fmt.Sprintf("%03d", serialCounter*2+2)
	clientCertSerial := fmt.Sprintf("%03d", serialCounter*2+4)
	revokedClientSerial := fmt.Sprintf("%03d", serialCounter*2+6)

	serialCounter = serialCounter + 3

	serverCAName := fmt.Sprintf("servers-ca-%s", serverCASerial)
	serverCACommonName := fmt.Sprintf("Servers %s CA", serverCASerial)
	serverCertName := fmt.Sprintf("server-instance-%s", serverSerial)
	serverCertCommonName := fmt.Sprintf("server%s.example.com", serverSerial)
	revokedServerCertName := fmt.Sprintf("server-instance-%s", revokedServerSerial)
	revokedServerCertCommonName := fmt.Sprintf("server%s.example.com", revokedServerSerial)

	clientCAName := fmt.Sprintf("clients-ca-%s", clientCASerial)
	clientCACommonName := fmt.Sprintf("Clients %s CA", clientCASerial)
	clientCertName := fmt.Sprintf("client-instance-%s", clientCertSerial)
	clientCertCommonName := fmt.Sprintf("client%s.example.com", clientCertSerial)
	revokedClientCertName := fmt.Sprintf("client-instance-%s", revokedClientSerial)
	revokedClientCertCommonName := fmt.Sprintf("client%s.example.com", revokedClientSerial)

	CreateSignedCert(root, CA, serverCASerial, serverCAName, serverCACommonName)
	CreateSignedCert(root, serverCAName, serverSerial, serverCertName, serverCertCommonName)
	CreateSignedCert(root, serverCAName, revokedServerSerial, revokedServerCertName, revokedServerCertCommonName)
	RevokeCertAndRegenerateCRL(root, serverCAName, revokedServerCertName)

	CreateSignedCert(root, CA, clientCASerial, clientCAName, clientCACommonName)
	CreateSignedCert(root, clientCAName, clientCertSerial, clientCertName, clientCertCommonName)
	CreateSignedCert(root, clientCAName, revokedClientSerial, revokedClientCertName, revokedClientCertCommonName)
	RevokeCertAndRegenerateCRL(root, clientCAName, revokedClientCertName)

	serverCRLPath := path.Join(root, fmt.Sprintf("%s-crl.pem", serverCAName))
	clientCRLPath := path.Join(root, fmt.Sprintf("%s-crl.pem", clientCAName))
	combinedCRLPath := path.Join(root, fmt.Sprintf("%s-%s-combined-crl.pem", serverCAName, clientCAName))

	serverCRLBytes, err := os.ReadFile(serverCRLPath)
	if err != nil {
		log.Fatalf("Could not read server CRL file")
	}

	clientCRLBytes, err := os.ReadFile(clientCRLPath)
	if err != nil {
		log.Fatalf("Could not read client CRL file")
	}

	err = os.WriteFile(combinedCRLPath, append(serverCRLBytes, clientCRLBytes...), 0777)
	if err != nil {
		log.Fatalf("Could not write combined CRL file")
	}

	return ClientServerKeyPairs{
		ServerCert:        path.Join(root, fmt.Sprintf("%s-cert.pem", serverCertName)),
		ServerKey:         path.Join(root, fmt.Sprintf("%s-key.pem", serverCertName)),
		ServerCA:          path.Join(root, fmt.Sprintf("%s-cert.pem", serverCAName)),
		ServerCRL:         serverCRLPath,
		RevokedServerCert: path.Join(root, fmt.Sprintf("%s-cert.pem", revokedServerCertName)),
		RevokedServerKey:  path.Join(root, fmt.Sprintf("%s-key.pem", revokedServerCertName)),
		ClientCert:        path.Join(root, fmt.Sprintf("%s-cert.pem", clientCertName)),
		ClientKey:         path.Join(root, fmt.Sprintf("%s-key.pem", clientCertName)),
		ClientCA:          path.Join(root, fmt.Sprintf("%s-cert.pem", clientCAName)),
		ClientCRL:         clientCRLPath,
		RevokedClientCert: path.Join(root, fmt.Sprintf("%s-cert.pem", revokedClientCertName)),
		RevokedClientKey:  path.Join(root, fmt.Sprintf("%s-key.pem", revokedClientCertName)),
		CombinedCRL:       combinedCRLPath,
		ServerName:        serverCertCommonName,
		RevokedServerName: revokedServerCertCommonName,
		RevokedClientName: revokedClientCertCommonName,
	}
}
<|endoftext|>


<commit_before>// Copyright 2013, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package vttablet

// This file handles the binlog players launched on masters for filtered
// replication

import (
	"bytes"
	"fmt"
	"time"

	"github.com/youtube/vitess/go/mysql"
	"github.com/youtube/vitess/go/relog"
	"github.com/youtube/vitess/go/stats"
	"github.com/youtube/vitess/go/vt/mysqlctl"
	"github.com/youtube/vitess/go/vt/topo"
)

// BinlogPlayerController controls one player
type BinlogPlayerController struct {
	ts       topo.Server
	dbConfig *mysql.ConnectionParams
	keyspace string
	source   topo.SourceShard

	// Player is the BinlogPlayer when we have one
	player *mysqlctl.BinlogPlayer

	// interrupted is the channel to close to stop the playback
	interrupted chan struct{}

	// TODO(alainjobart): add state
	// TODO(alainjobart): add statsJson, include stats and player stats if any
	// TODO(alainjobart): figure out if we need a lock on structure (for stats)
}

func NewBinlogController(ts topo.Server, dbConfig *mysql.ConnectionParams, keyspace string, source topo.SourceShard) *BinlogPlayerController {
	return &BinlogPlayerController{
		ts:          ts,
		dbConfig:    dbConfig,
		keyspace:    keyspace,
		source:      source,
		interrupted: make(chan struct{}, 1),
	}
}

func (blc *BinlogPlayerController) statsJSON() string {
	buf := bytes.NewBuffer(make([]byte, 0, 128))
	fmt.Fprintf(buf, "{")
	fmt.Fprintf(buf, "\n \"Player\": %v", blc.player.StatsJSON())
	fmt.Fprintf(buf, "\n}")
	return buf.String()
}

func (bpc *BinlogPlayerController) Start() {
	relog.Info("Starting binlog player for %v", bpc.source)
	go bpc.Loop()
}

func (bpc *BinlogPlayerController) Stop() {
	relog.Info("Stopping binlog player for %v", bpc.source)
	close(bpc.interrupted)
}

func (bpc *BinlogPlayerController) Loop() {
	for {
		// create the db connection, connect it
		vtClient := mysqlctl.NewDbClient(bpc.dbConfig)
		if err := vtClient.Connect(); err != nil {
			relog.Error("BinlogPlayerMap: can't connect to database: %v", err)
			continue
		}

		// Read the start position
		startPosition, err := mysqlctl.ReadStartPosition(vtClient, string(bpc.source.KeyRange.Start.Hex()), string(bpc.source.KeyRange.End.Hex()))
		if err != nil {
			relog.Warning("BinlogPlayerController: can't read startPosition: %v", err)
			time.Sleep(5)
			continue
		}

		// TODO(alainjobart): Find the server list

		// TODO(alainjobart): Pick a server (same if it's available,
		// if not clear master file / pos and keep only group id)

		// Create the player.
		bpc.player, err = mysqlctl.NewBinlogPlayer(vtClient, startPosition, nil /*tables*/, 1 /*txnBatch*/, 30*time.Second /*maxTxnInterval*/, false /*execDdl*/)
		if err != nil {
			relog.Warning("BinlogPlayerController: can't create player: %v", err)
			time.Sleep(5)
			continue
		}

		// Run player loop until it's done.
		err = bpc.player.ApplyBinlogEvents(bpc.interrupted)
		if err != nil {
			relog.Error("Error in applying binlog events, err %v", err)
			time.Sleep(5)
		} else {
			// We were interrupted.
			break
		}
	}
	relog.Info("Exited main binlog player loop for %v", bpc.source)
}

// BinlogPlayerMap controls all the players
type BinlogPlayerMap struct {
	ts       topo.Server
	dbConfig *mysql.ConnectionParams
	players  map[topo.SourceShard]*BinlogPlayerController
}

func NewBinlogPlayerMap(ts topo.Server, dbConfig *mysql.ConnectionParams) *BinlogPlayerMap {
	return &BinlogPlayerMap{
		ts:       ts,
		dbConfig: dbConfig,
		players:  make(map[topo.SourceShard]*BinlogPlayerController),
	}
}

func RegisterBinlogPlayerMap(blm *BinlogPlayerMap) {
	stats.PublishFunc("BinlogPlayerMap", func() string { return blm.statsJSON() })
}

func (blm *BinlogPlayerMap) statsJSON() string {
	buf := bytes.NewBuffer(make([]byte, 0, 128))
	fmt.Fprintf(buf, "{")
	for source, player := range blm.players {
		fmt.Fprintf(buf, "\n \"%v-%v\": %v,", source.KeyRange.Start, source.KeyRange.End, player.statsJSON())
	}
	fmt.Fprintf(buf, "\n \"Enabled\": 1")
	fmt.Fprintf(buf, "\n}")
	return buf.String()
}

func (blm *BinlogPlayerMap) AddPlayer(keyspace string, source topo.SourceShard) {
	bpc, ok := blm.players[source]
	if ok {
		relog.Info("Already playing logs for %v", source)
		return
	}

	bpc = NewBinlogController(blm.ts, blm.dbConfig, keyspace, source)
	blm.players[source] = bpc
	bpc.Start()
}

func (blm *BinlogPlayerMap) StopAllPlayers() {
	for _, bpc := range blm.players {
		bpc.Stop()
	}
	blm.players = make(map[topo.SourceShard]*BinlogPlayerController)
}

// RefreshMap reads the right data from topo.Server and makes sure
// we're playing the right logs
func (blm *BinlogPlayerMap) RefreshMap(tablet topo.Tablet) {
	relog.Info("Refreshing map of binlog players")

	// read the shard to get SourceShards
	shardInfo, err := blm.ts.GetShard(tablet.Keyspace, tablet.Shard)
	if err != nil {
		relog.Error("Cannot read shard for this tablet: %v", tablet.Alias())
		return
	}

	// get the existing sources and build a map of sources to remove
	toRemove := make(map[topo.SourceShard]bool)
	for source, _ := range blm.players {
		toRemove[source] = true
	}

	// for each source, add it if not there, and delete from toRemove
	for _, source := range shardInfo.SourceShards {
		blm.AddPlayer(tablet.Keyspace, source)
		delete(toRemove, source)
	}

	// remove all entries from toRemove
	for source, _ := range toRemove {
		blm.players[source].Stop()
		delete(blm.players, source)
	}
}
<commit_msg>Adding states, fixing mutex access.<commit_after>// Copyright 2013, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package vttablet

// This file handles the binlog players launched on masters for filtered
// replication

import (
	"bytes"
	"fmt"
	"sync"
	"time"

	"github.com/youtube/vitess/go/mysql"
	"github.com/youtube/vitess/go/relog"
	"github.com/youtube/vitess/go/stats"
	"github.com/youtube/vitess/go/vt/mysqlctl"
	"github.com/youtube/vitess/go/vt/topo"
)

const (
	BINLOG_PLAYER_CONNECTING = iota
	BINLOG_PLAYER_PLAYING
	BINLOG_PLAYER_SLEEPING
)

// BinlogPlayerController controls one player
type BinlogPlayerController struct {
	ts       topo.Server
	dbConfig *mysql.ConnectionParams
	keyspace string
	source   topo.SourceShard

	// Player is the BinlogPlayer when we have one, protected by mu
	player *mysqlctl.BinlogPlayer
	mu     sync.Mutex

	// states is our state accounting variable, for stats
	states *stats.States

	// interrupted is the channel to close to stop the playback
	interrupted chan struct{}
}

func NewBinlogController(ts topo.Server, dbConfig *mysql.ConnectionParams, keyspace string, source topo.SourceShard) *BinlogPlayerController {
	blc := &BinlogPlayerController{
		ts:          ts,
		dbConfig:    dbConfig,
		keyspace:    keyspace,
		source:      source,
		interrupted: make(chan struct{}, 1),
	}
	blc.states = stats.NewStates("", []string{
		"Connecting",
		"Playing",
		"Sleeping",
	}, time.Now(), BINLOG_PLAYER_CONNECTING)
	return blc
}

func (blc *BinlogPlayerController) statsJSON() string {
	buf := bytes.NewBuffer(make([]byte, 0, 128))
	fmt.Fprintf(buf, "{")
	blc.mu.Lock()
	if blc.player != nil {
		fmt.Fprintf(buf, "\n \"Player\": %v", blc.player.StatsJSON())
	}
	blc.mu.Unlock()
	fmt.Fprintf(buf, "\n \"States\": %v", blc.states.String())
	fmt.Fprintf(buf, "\n}")
	return buf.String()
}

func (bpc *BinlogPlayerController) Start() {
	relog.Info("Starting binlog player for %v", bpc.source)
	go bpc.Loop()
}

func (bpc *BinlogPlayerController) Stop() {
	relog.Info("Stopping binlog player for %v", bpc.source)
	close(bpc.interrupted)
}

func (bpc *BinlogPlayerController) Loop() {
	for {
		err := bpc.Iteration()
		if err == nil {
			// this happens when we get interrupted
			break
		}
		relog.Warning("BinlogPlayerController(%v-%v): %v", bpc.source.KeyRange.Start, bpc.source.KeyRange.End, err)

		// sleep for a bit before retrying to connect
		bpc.states.SetState(BINLOG_PLAYER_SLEEPING)
		time.Sleep(5)
	}

	relog.Info("Exited main binlog player loop for %v", bpc.source)
}

func (bpc *BinlogPlayerController) Iteration() (err error) {
	defer func() {
		if x := recover(); x != nil {
			relog.Error("BinlogPlayerController caught panic: %v", x)
			err = fmt.Errorf("panic: %v", x)
		}
	}()
	bpc.states.SetState(BINLOG_PLAYER_CONNECTING)

	// create the db connection, connect it
	vtClient := mysqlctl.NewDbClient(bpc.dbConfig)
	if err := vtClient.Connect(); err != nil {
		return fmt.Errorf("can't connect to database: %v", err)
	}

	// Read the start position
	startPosition, err := mysqlctl.ReadStartPosition(vtClient, string(bpc.source.KeyRange.Start.Hex()), string(bpc.source.KeyRange.End.Hex()))
	if err != nil {
		return fmt.Errorf("can't read startPosition: %v", err)
	}

	// TODO(alainjobart): Find the server list

	// TODO(alainjobart): Pick a server (same if it's available,
	// if not clear master file / pos and keep only group id)

	// Create the player.
	bpc.mu.Lock()
	bpc.player, err = mysqlctl.NewBinlogPlayer(vtClient, startPosition, nil /*tables*/, 1 /*txnBatch*/, 30*time.Second /*maxTxnInterval*/, false /*execDdl*/)
	bpc.mu.Unlock()
	if err != nil {
		return fmt.Errorf("can't create player: %v", err)
	}

	// Run player loop until it's done.
	bpc.states.SetState(BINLOG_PLAYER_PLAYING)
	err = bpc.player.ApplyBinlogEvents(bpc.interrupted)

	bpc.mu.Lock()
	bpc.player = nil
	bpc.mu.Unlock()

	return err
}

// BinlogPlayerMap controls all the players
type BinlogPlayerMap struct {
	ts       topo.Server
	dbConfig *mysql.ConnectionParams
	players  map[topo.SourceShard]*BinlogPlayerController
}

func NewBinlogPlayerMap(ts topo.Server, dbConfig *mysql.ConnectionParams) *BinlogPlayerMap {
	return &BinlogPlayerMap{
		ts:       ts,
		dbConfig: dbConfig,
		players:  make(map[topo.SourceShard]*BinlogPlayerController),
	}
}

func RegisterBinlogPlayerMap(blm *BinlogPlayerMap) {
	stats.PublishFunc("BinlogPlayerMap", func() string { return blm.statsJSON() })
}

func (blm *BinlogPlayerMap) statsJSON() string {
	buf := bytes.NewBuffer(make([]byte, 0, 128))
	fmt.Fprintf(buf, "{")
	for source, player := range blm.players {
		fmt.Fprintf(buf, "\n \"%v-%v\": %v,", source.KeyRange.Start, source.KeyRange.End, player.statsJSON())
	}
	fmt.Fprintf(buf, "\n \"Enabled\": 1")
	fmt.Fprintf(buf, "\n}")
	return buf.String()
}

func (blm *BinlogPlayerMap) AddPlayer(keyspace string, source topo.SourceShard) {
	bpc, ok := blm.players[source]
	if ok {
		relog.Info("Already playing logs for %v", source)
		return
	}

	bpc = NewBinlogController(blm.ts, blm.dbConfig, keyspace, source)
	blm.players[source] = bpc
	bpc.Start()
}

func (blm *BinlogPlayerMap) StopAllPlayers() {
	for _, bpc := range blm.players {
		bpc.Stop()
	}
	blm.players = make(map[topo.SourceShard]*BinlogPlayerController)
}

// RefreshMap reads the right data from topo.Server and makes sure
// we're playing the right logs
func (blm *BinlogPlayerMap) RefreshMap(tablet topo.Tablet) {
	relog.Info("Refreshing map of binlog players")

	// read the shard to get SourceShards
	shardInfo, err := blm.ts.GetShard(tablet.Keyspace, tablet.Shard)
	if err != nil {
		relog.Error("Cannot read shard for this tablet: %v", tablet.Alias())
		return
	}

	// get the existing sources and build a map of sources to remove
	toRemove := make(map[topo.SourceShard]bool)
	for source, _ := range blm.players {
		toRemove[source] = true
	}

	// for each source, add it if not there, and delete from toRemove
	for _, source := range shardInfo.SourceShards {
		blm.AddPlayer(tablet.Keyspace, source)
		delete(toRemove, source)
	}

	// remove all entries from toRemove
	for source, _ := range toRemove {
		blm.players[source].Stop()
		delete(blm.players, source)
	}
}
<|endoftext|>


<commit_before>// Copyright 2012 Aryan Naraghi (aryan.naraghi@gmail.com)
//
// 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.

// A demo for difflib. This program accepts the paths to two files and
// launches a web server at port 8080 that serves the diff results.
package main

import (
	"fmt"
	"github.com/aryann/difflib"
	"html"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
)

var hostPort = "localhost:8080"

var templateString = `
<!doctype html>

  <meta charset="utf-8" />
  <title>difflib results</title>

      {{.Filename1}}
      {{.Filename2}}

    {{.Diff}}

`

func main() {
	if len(os.Args) != 3 {
		fmt.Fprintf(os.Stderr, "USAGE: %s <file-1> <file-2>\n", os.Args[0])
		os.Exit(1)
	}
	http.HandleFunc("/", diffHandler(os.Args[1], os.Args[2]))
	fmt.Printf("Starting server at %s.\n", hostPort)
	err := http.ListenAndServe(hostPort, nil)
	if err != nil {
		panic(err)
	}
}

// diffHandler returns an http.HandlerFunc that serves the diff of the
// two given files.
func diffHandler(filename1, filename2 string) http.HandlerFunc {
	diff := difflib.HTMLDiff(fileToLines(filename1), fileToLines(filename2))
	tmpl, _ := template.New("diffTemplate").Parse(templateString)
	return func(w http.ResponseWriter, r *http.Request) {
		err := tmpl.Execute(w, map[string]interface{}{
			"Diff":      template.HTML(diff),
			"Filename1": filename1,
			"Filename2": filename2,
		})
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
		}
	}
}

// filesToLines takes a path to a file and returns a string array of
// the lines in the file. Any HTML in the file is escaped.
func fileToLines(filename string) []string {
	contents, _ := ioutil.ReadFile(filename)
	return strings.Split(html.EscapeString(string(contents)), "\n")
}
<commit_msg>Removed some extraneous CSS.<commit_after>// Copyright 2012 Aryan Naraghi (aryan.naraghi@gmail.com)
//
// 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.

// A demo for difflib. This program accepts the paths to two files and
// launches a web server at port 8080 that serves the diff results.
package main

import (
	"fmt"
	"github.com/aryann/difflib"
	"html"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
)

var hostPort = "localhost:8080"

var templateString = `
<!doctype html>

  <meta charset="utf-8" />
  <title>difflib results</title>

      {{.Filename1}}
      {{.Filename2}}

    {{.Diff}}

`

func main() {
	if len(os.Args) != 3 {
		fmt.Fprintf(os.Stderr, "USAGE: %s <file-1> <file-2>\n", os.Args[0])
		os.Exit(1)
	}
	http.HandleFunc("/", diffHandler(os.Args[1], os.Args[2]))
	fmt.Printf("Starting server at %s.\n", hostPort)
	err := http.ListenAndServe(hostPort, nil)
	if err != nil {
		panic(err)
	}
}

// diffHandler returns an http.HandlerFunc that serves the diff of the
// two given files.
func diffHandler(filename1, filename2 string) http.HandlerFunc {
	diff := difflib.HTMLDiff(fileToLines(filename1), fileToLines(filename2))
	tmpl, _ := template.New("diffTemplate").Parse(templateString)
	return func(w http.ResponseWriter, r *http.Request) {
		err := tmpl.Execute(w, map[string]interface{}{
			"Diff":      template.HTML(diff),
			"Filename1": filename1,
			"Filename2": filename2,
		})
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
		}
	}
}

// filesToLines takes a path to a file and returns a string array of
// the lines in the file. Any HTML in the file is escaped.
func fileToLines(filename string) []string {
	contents, _ := ioutil.ReadFile(filename)
	return strings.Split(html.EscapeString(string(contents)), "\n")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"testing"
)

func TestSourceJsonFiles(t *testing.T) {
	sources := GetSources()

	fmt.Printf("Testing %v source files\n", len(sources.Sources))

	for _, sourceFile := range sources.Sources {
		source, err := GetSourceContent(sourceFile)

		if err != nil {
			t.Logf("Source file %v contains invalid Json", sourceFile)
			t.FailNow()
		}

		if source.Source == "" {
			t.Errorf("Source field inside the %v file is empty", sourceFile)
		}

		if len(source.Paragraphs) == 0 {
			t.Errorf("Source file %v contains no paragraphs", sourceFile)
		}

		for index, paragraph := range source.Paragraphs {
			if paragraph == "" {
				t.Errorf("Paragraph %v the %v file is empty", index, sourceFile)
			}
		}
	}
}

func TestNumLinesFunction(t *testing.T) {
	source := GetNumLines(5, false)

	if len(source.Paragraphs) != 5 {
		t.Errorf("%v paragraphs returned, exptected %v", len(source.Paragraphs), 5)
	}

	source2 := GetNumLines(50, false)

	if len(source2.Paragraphs) != 50 {
		t.Errorf("%v paragraphs returned, exptected %v", len(source2.Paragraphs), 50)
	}
}
<commit_msg>Correct the sources test<commit_after>package main

import (
	"fmt"
	"testing"
)

func TestSourceJsonFiles(t *testing.T) {
	sources := GetSources()

	fmt.Printf("Testing %v source files\n", len(sources))

	for _, sourceFile := range sources {
		source, err := GetSourceContent(sourceFile)

		if err != nil {
			t.Logf("Source file %v contains invalid Json", sourceFile)
			t.FailNow()
		}

		if source.Source == "" {
			t.Errorf("Source field inside the %v file is empty", sourceFile)
		}

		if len(source.Paragraphs) == 0 {
			t.Errorf("Source file %v contains no paragraphs", sourceFile)
		}

		for index, paragraph := range source.Paragraphs {
			if paragraph == "" {
				t.Errorf("Paragraph %v the %v file is empty", index, sourceFile)
			}
		}
	}
}

func TestNumLinesFunction(t *testing.T) {
	source := GetNumLines(5, false)

	if len(source.Paragraphs) != 5 {
		t.Errorf("%v paragraphs returned, exptected %v", len(source.Paragraphs), 5)
	}

	source2 := GetNumLines(50, false)

	if len(source2.Paragraphs) != 50 {
		t.Errorf("%v paragraphs returned, exptected %v", len(source2.Paragraphs), 50)
	}
}
<|endoftext|>


<commit_before>package multiaddr

// You **MUST** register your multicodecs with
// https://github.com/multiformats/multicodec before adding them here.
//
// TODO: Use a single source of truth for all multicodecs instead of
// distributing them like this...
const (
	P_IP4               = 0x0004
	P_TCP               = 0x0006
	P_UDP               = 0x0111
	P_DCCP              = 0x0021
	P_IP6               = 0x0029
	P_IP6ZONE           = 0x002A
	P_QUIC              = 0x01CC
	P_SCTP              = 0x0084
	P_UDT               = 0x012D
	P_UTP               = 0x012E
	P_UNIX              = 0x0190
	P_P2P               = 0x01A5
	P_IPFS              = 0x01A5 // alias for backwards compatability
	P_HTTP              = 0x01E0
	P_HTTPS             = 0x01BB
	P_ONION             = 0x01BC // also for backwards compatibility
	P_ONION3            = 0x01BD
	P_GARLIC64          = 0x01CA
	P_P2P_WEBRTC_DIRECT = 0x0114
)

var (
	protoIP4 = Protocol{
		Name:       "ip4",
		Code:       P_IP4,
		VCode:      CodeToVarint(P_IP4),
		Size:       32,
		Path:       false,
		Transcoder: TranscoderIP4,
	}
	protoTCP = Protocol{
		Name:       "tcp",
		Code:       P_TCP,
		VCode:      CodeToVarint(P_TCP),
		Size:       16,
		Path:       false,
		Transcoder: TranscoderPort,
	}
	protoUDP = Protocol{
		Name:       "udp",
		Code:       P_UDP,
		VCode:      CodeToVarint(P_UDP),
		Size:       16,
		Path:       false,
		Transcoder: TranscoderPort,
	}
	protoDCCP = Protocol{
		Name:       "dccp",
		Code:       P_DCCP,
		VCode:      CodeToVarint(P_DCCP),
		Size:       16,
		Path:       false,
		Transcoder: TranscoderPort,
	}
	protoIP6 = Protocol{
		Name:       "ip6",
		Code:       P_IP6,
		VCode:      CodeToVarint(P_IP6),
		Size:       128,
		Transcoder: TranscoderIP6,
	}
	// these require varint
	protoIP6ZONE = Protocol{
		Name:       "ip6zone",
		Code:       P_IP6ZONE,
		VCode:      CodeToVarint(P_IP6ZONE),
		Size:       LengthPrefixedVarSize,
		Path:       false,
		Transcoder: TranscoderIP6Zone,
	}
	protoSCTP = Protocol{
		Name:       "sctp",
		Code:       P_SCTP,
		VCode:      CodeToVarint(P_SCTP),
		Size:       16,
		Transcoder: TranscoderPort,
	}
	protoONION2 = Protocol{
		Name:       "onion",
		Code:       P_ONION,
		VCode:      CodeToVarint(P_ONION),
		Size:       96,
		Transcoder: TranscoderOnion,
	}
	protoONION3 = Protocol{
		Name:       "onion3",
		Code:       P_ONION3,
		VCode:      CodeToVarint(P_ONION3),
		Size:       296,
		Transcoder: TranscoderOnion3,
	}
	protoGARLIC64 = Protocol{
		Name:       "garlic64",
		Code:       P_GARLIC64,
		VCode:      CodeToVarint(P_GARLIC64),
		Size:       LengthPrefixedVarSize,
		Transcoder: TranscoderGarlic64,
	}
	protoUTP = Protocol{
		Name:  "utp",
		Code:  P_UTP,
		VCode: CodeToVarint(P_UTP),
	}
	protoUDT = Protocol{
		Name:  "udt",
		Code:  P_UDT,
		VCode: CodeToVarint(P_UDT),
	}
	protoQUIC = Protocol{
		Name:  "quic",
		Code:  P_QUIC,
		VCode: CodeToVarint(P_QUIC),
	}
	protoHTTP = Protocol{
		Name:  "http",
		Code:  P_HTTP,
		VCode: CodeToVarint(P_HTTP),
	}
	protoHTTPS = Protocol{
		Name:  "https",
		Code:  P_HTTPS,
		VCode: CodeToVarint(P_HTTPS),
	}
	protoP2P = Protocol{
		Name:       "ipfs",
		Code:       P_P2P,
		VCode:      CodeToVarint(P_P2P),
		Size:       LengthPrefixedVarSize,
		Transcoder: TranscoderP2P,
	}
	protoUNIX = Protocol{
		Name:       "unix",
		Code:       P_UNIX,
		VCode:      CodeToVarint(P_UNIX),
		Size:       LengthPrefixedVarSize,
		Path:       true,
		Transcoder: TranscoderUnix,
	}
	protoP2P_WEBRTC_DIRECT = Protocol{
		Name:  "p2p-webrtc-direct",
		Code:  P_P2P_WEBRTC_DIRECT,
		VCode: CodeToVarint(P_P2P_WEBRTC_DIRECT),
	}
)

func init() {
	for _, p := range []Protocol{
		protoIP4,
		protoTCP,
		protoUDP,
		protoDCCP,
		protoIP6,
		protoIP6ZONE,
		protoSCTP,
		protoONION2,
		protoONION3,
		protoGARLIC64,
		protoUTP,
		protoUDT,
		protoQUIC,
		protoHTTP,
		protoHTTPS,
		protoP2P,
		protoUNIX,
		protoP2P_WEBRTC_DIRECT,
	} {
		if err := AddProtocol(p); err != nil {
			panic(err)
		}
	}

	// explicitly set both of these
	protocolsByName["p2p"] = protoP2P
	protocolsByName["ipfs"] = protoP2P
}
<commit_msg>remove stale TODO.<commit_after>package multiaddr

// You **MUST** register your multicodecs with
// https://github.com/multiformats/multicodec before adding them here.
const (
	P_IP4               = 0x0004
	P_TCP               = 0x0006
	P_UDP               = 0x0111
	P_DCCP              = 0x0021
	P_IP6               = 0x0029
	P_IP6ZONE           = 0x002A
	P_QUIC              = 0x01CC
	P_SCTP              = 0x0084
	P_UDT               = 0x012D
	P_UTP               = 0x012E
	P_UNIX              = 0x0190
	P_P2P               = 0x01A5
	P_IPFS              = 0x01A5 // alias for backwards compatability
	P_HTTP              = 0x01E0
	P_HTTPS             = 0x01BB
	P_ONION             = 0x01BC // also for backwards compatibility
	P_ONION3            = 0x01BD
	P_GARLIC64          = 0x01CA
	P_P2P_WEBRTC_DIRECT = 0x0114
)

var (
	protoIP4 = Protocol{
		Name:       "ip4",
		Code:       P_IP4,
		VCode:      CodeToVarint(P_IP4),
		Size:       32,
		Path:       false,
		Transcoder: TranscoderIP4,
	}
	protoTCP = Protocol{
		Name:       "tcp",
		Code:       P_TCP,
		VCode:      CodeToVarint(P_TCP),
		Size:       16,
		Path:       false,
		Transcoder: TranscoderPort,
	}
	protoUDP = Protocol{
		Name:       "udp",
		Code:       P_UDP,
		VCode:      CodeToVarint(P_UDP),
		Size:       16,
		Path:       false,
		Transcoder: TranscoderPort,
	}
	protoDCCP = Protocol{
		Name:       "dccp",
		Code:       P_DCCP,
		VCode:      CodeToVarint(P_DCCP),
		Size:       16,
		Path:       false,
		Transcoder: TranscoderPort,
	}
	protoIP6 = Protocol{
		Name:       "ip6",
		Code:       P_IP6,
		VCode:      CodeToVarint(P_IP6),
		Size:       128,
		Transcoder: TranscoderIP6,
	}
	// these require varint
	protoIP6ZONE = Protocol{
		Name:       "ip6zone",
		Code:       P_IP6ZONE,
		VCode:      CodeToVarint(P_IP6ZONE),
		Size:       LengthPrefixedVarSize,
		Path:       false,
		Transcoder: TranscoderIP6Zone,
	}
	protoSCTP = Protocol{
		Name:       "sctp",
		Code:       P_SCTP,
		VCode:      CodeToVarint(P_SCTP),
		Size:       16,
		Transcoder: TranscoderPort,
	}
	protoONION2 = Protocol{
		Name:       "onion",
		Code:       P_ONION,
		VCode:      CodeToVarint(P_ONION),
		Size:       96,
		Transcoder: TranscoderOnion,
	}
	protoONION3 = Protocol{
		Name:       "onion3",
		Code:       P_ONION3,
		VCode:      CodeToVarint(P_ONION3),
		Size:       296,
		Transcoder: TranscoderOnion3,
	}
	protoGARLIC64 = Protocol{
		Name:       "garlic64",
		Code:       P_GARLIC64,
		VCode:      CodeToVarint(P_GARLIC64),
		Size:       LengthPrefixedVarSize,
		Transcoder: TranscoderGarlic64,
	}
	protoUTP = Protocol{
		Name:  "utp",
		Code:  P_UTP,
		VCode: CodeToVarint(P_UTP),
	}
	protoUDT = Protocol{
		Name:  "udt",
		Code:  P_UDT,
		VCode: CodeToVarint(P_UDT),
	}
	protoQUIC = Protocol{
		Name:  "quic",
		Code:  P_QUIC,
		VCode: CodeToVarint(P_QUIC),
	}
	protoHTTP = Protocol{
		Name:  "http",
		Code:  P_HTTP,
		VCode: CodeToVarint(P_HTTP),
	}
	protoHTTPS = Protocol{
		Name:  "https",
		Code:  P_HTTPS,
		VCode: CodeToVarint(P_HTTPS),
	}
	protoP2P = Protocol{
		Name:       "ipfs",
		Code:       P_P2P,
		VCode:      CodeToVarint(P_P2P),
		Size:       LengthPrefixedVarSize,
		Transcoder: TranscoderP2P,
	}
	protoUNIX = Protocol{
		Name:       "unix",
		Code:       P_UNIX,
		VCode:      CodeToVarint(P_UNIX),
		Size:       LengthPrefixedVarSize,
		Path:       true,
		Transcoder: TranscoderUnix,
	}
	protoP2P_WEBRTC_DIRECT = Protocol{
		Name:  "p2p-webrtc-direct",
		Code:  P_P2P_WEBRTC_DIRECT,
		VCode: CodeToVarint(P_P2P_WEBRTC_DIRECT),
	}
)

func init() {
	for _, p := range []Protocol{
		protoIP4,
		protoTCP,
		protoUDP,
		protoDCCP,
		protoIP6,
		protoIP6ZONE,
		protoSCTP,
		protoONION2,
		protoONION3,
		protoGARLIC64,
		protoUTP,
		protoUDT,
		protoQUIC,
		protoHTTP,
		protoHTTPS,
		protoP2P,
		protoUNIX,
		protoP2P_WEBRTC_DIRECT,
	} {
		if err := AddProtocol(p); err != nil {
			panic(err)
		}
	}

	// explicitly set both of these
	protocolsByName["p2p"] = protoP2P
	protocolsByName["ipfs"] = protoP2P
}
<|endoftext|>


<commit_before>package chat

import (
	"context"
	"testing"
	"time"

	"github.com/keybase/client/go/chat/globals"
	"github.com/keybase/client/go/chat/storage"
	"github.com/keybase/client/go/chat/types"
	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
	"github.com/stretchr/testify/require"
)

func TestBackgroundPurge(t *testing.T) {
	ctx, tc, world, _, baseSender, listener, res := setupLoaderTest(t)
	defer world.Cleanup()

	g := globals.NewContext(tc.G, tc.ChatG)
	u := world.GetUsers()[0]
	uid := gregor1.UID(u.GetUID().ToBytes())
	trip := newConvTriple(ctx, t, tc, u.Username)
	clock := world.Fc
	chatStorage := storage.New(g, tc.ChatG.ConvSource)
	chatStorage.SetClock(clock)

	assertListener := func(convID chat1.ConversationID) {
		select {
		case convID := <-listener.bgConvLoads:
			require.Equal(t, res.ConvID, convID)
		case <-time.After(2 * time.Second):
			t.Fatal("timeout waiting for conversation load")
		}
	}

	sendEphemeral := func(lifetime gregor1.DurationSec) {
		_, _, err := baseSender.Send(ctx, res.ConvID, chat1.MessagePlaintext{
			ClientHeader: chat1.MessageClientHeader{
				Conv:              trip,
				Sender:            uid,
				TlfName:           u.Username,
				TlfPublic:         false,
				MessageType:       chat1.MessageType_TEXT,
				EphemeralMetadata: &chat1.MsgEphemeralMetadata{Lifetime: lifetime},
			},
			MessageBody: chat1.NewMessageBodyWithText(chat1.MessageText{
				Body: "hi",
			}),
		}, 0, nil)
		require.NoError(t, err)
	}

	assertTrackerState := func(convID chat1.ConversationID, expectedPurgeInfo chat1.EphemeralPurgeInfo) {
		allPurgeInfo, err := chatStorage.GetAllPurgeInfo(ctx, uid)
		require.NoError(t, err)
		purgeInfo, ok := allPurgeInfo[convID.String()]
		require.True(t, ok)
		require.Equal(t, expectedPurgeInfo, purgeInfo)
	}

	assertEphemeralPurgeNotifInfo := func(convID chat1.ConversationID, msgIDs []chat1.MessageID) {
		info := listener.consumeEphemeralPurge(t)
		require.Equal(t, info.ConvID, convID)
		if msgIDs == nil {
			require.Nil(t, info.Msgs)
		} else {
			purgedIDs := []chat1.MessageID{}
			for _, purgedMsg := range info.Msgs {
				purgedIDs = append(purgedIDs, purgedMsg.GetMessageID())
			}
			require.Equal(t, msgIDs, purgedIDs)
			require.Equal(t, convID.String(), info.Conv.ConvID)
		}
	}

	// Load our conv with the initial tlf msg
	require.NoError(t, tc.Context().ConvLoader.Queue(context.TODO(),
		types.NewConvLoaderJob(res.ConvID, &chat1.Pagination{Num: 3}, types.ConvLoaderPriorityHigh, nil)))
	t.Logf("assert listener 0")
	assertListener(res.ConvID)

	// Nothing is up for purging yet
	expectedPurgeInfo := chat1.EphemeralPurgeInfo{
		ConvID:          res.ConvID,
		MinUnexplodedID: 1,
		NextPurgeTime:   0,
		IsActive:        false,
	}
	assertTrackerState(res.ConvID, expectedPurgeInfo)

	// Send two ephemeral messages, and ensure both get purged
	lifetime := gregor1.DurationSec(1)
	lifetimeDuration := time.Second
	sendEphemeral(lifetime)
	sendEphemeral(lifetime * 2)

	thread, err := tc.ChatG.ConvSource.Pull(ctx, res.ConvID, uid, &chat1.GetThreadQuery{
		MessageTypes: []chat1.MessageType{chat1.MessageType_TEXT},
	}, nil)
	require.NoError(t, err)
	msgs := thread.Messages
	require.Len(t, msgs, 2)
	msgUnboxed2 := msgs[0]
	msgUnboxed1 := msgs[1]

	t.Logf("assert listener 1")
	world.Fc.Advance(lifetimeDuration)
	assertListener(res.ConvID)
	assertTrackerState(res.ConvID, chat1.EphemeralPurgeInfo{
		ConvID:          res.ConvID,
		MinUnexplodedID: msgUnboxed2.GetMessageID(),
		NextPurgeTime:   msgUnboxed2.Valid().Etime(),
		IsActive:        true,
	})
	assertEphemeralPurgeNotifInfo(res.ConvID, []chat1.MessageID{msgUnboxed1.GetMessageID()})

	// Ensure things run smoothly even with extraneous start/stop calls to purger
	g.EphemeralPurger.Start(context.Background(), uid)
	<-g.EphemeralPurger.Stop(context.Background())
	<-g.EphemeralPurger.Stop(context.Background())
	g.EphemeralPurger.Start(context.Background(), uid)
	g.EphemeralPurger.Start(context.Background(), uid)

	t.Logf("assert listener 2")
	world.Fc.Advance(lifetimeDuration)
	assertListener(res.ConvID)
	assertTrackerState(res.ConvID, chat1.EphemeralPurgeInfo{
		ConvID:          res.ConvID,
		MinUnexplodedID: msgUnboxed2.GetMessageID(),
		NextPurgeTime:   0,
		IsActive:        false,
	})
	assertEphemeralPurgeNotifInfo(res.ConvID, []chat1.MessageID{msgUnboxed2.GetMessageID()})
}
<commit_msg>Update ephemeral_purger_test.go (#11967)<commit_after>package chat

import (
	"context"
	"testing"
	"time"

	"github.com/keybase/client/go/chat/globals"
	"github.com/keybase/client/go/chat/storage"
	"github.com/keybase/client/go/chat/types"
	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
	"github.com/stretchr/testify/require"
)

func TestBackgroundPurge(t *testing.T) {
	t.Skip()
	ctx, tc, world, _, baseSender, listener, res := setupLoaderTest(t)
	defer world.Cleanup()

	g := globals.NewContext(tc.G, tc.ChatG)
	u := world.GetUsers()[0]
	uid := gregor1.UID(u.GetUID().ToBytes())
	trip := newConvTriple(ctx, t, tc, u.Username)
	clock := world.Fc
	chatStorage := storage.New(g, tc.ChatG.ConvSource)
	chatStorage.SetClock(clock)

	assertListener := func(convID chat1.ConversationID) {
		select {
		case convID := <-listener.bgConvLoads:
			require.Equal(t, res.ConvID, convID)
		case <-time.After(2 * time.Second):
			t.Fatal("timeout waiting for conversation load")
		}
	}

	sendEphemeral := func(lifetime gregor1.DurationSec) {
		_, _, err := baseSender.Send(ctx, res.ConvID, chat1.MessagePlaintext{
			ClientHeader: chat1.MessageClientHeader{
				Conv:              trip,
				Sender:            uid,
				TlfName:           u.Username,
				TlfPublic:         false,
				MessageType:       chat1.MessageType_TEXT,
				EphemeralMetadata: &chat1.MsgEphemeralMetadata{Lifetime: lifetime},
			},
			MessageBody: chat1.NewMessageBodyWithText(chat1.MessageText{
				Body: "hi",
			}),
		}, 0, nil)
		require.NoError(t, err)
	}

	assertTrackerState := func(convID chat1.ConversationID, expectedPurgeInfo chat1.EphemeralPurgeInfo) {
		allPurgeInfo, err := chatStorage.GetAllPurgeInfo(ctx, uid)
		require.NoError(t, err)
		purgeInfo, ok := allPurgeInfo[convID.String()]
		require.True(t, ok)
		require.Equal(t, expectedPurgeInfo, purgeInfo)
	}

	assertEphemeralPurgeNotifInfo := func(convID chat1.ConversationID, msgIDs []chat1.MessageID) {
		info := listener.consumeEphemeralPurge(t)
		require.Equal(t, info.ConvID, convID)
		if msgIDs == nil {
			require.Nil(t, info.Msgs)
		} else {
			purgedIDs := []chat1.MessageID{}
			for _, purgedMsg := range info.Msgs {
				purgedIDs = append(purgedIDs, purgedMsg.GetMessageID())
			}
			require.Equal(t, msgIDs, purgedIDs)
			require.Equal(t, convID.String(), info.Conv.ConvID)
		}
	}

	// Load our conv with the initial tlf msg
	require.NoError(t, tc.Context().ConvLoader.Queue(context.TODO(),
		types.NewConvLoaderJob(res.ConvID, &chat1.Pagination{Num: 3}, types.ConvLoaderPriorityHigh, nil)))
	t.Logf("assert listener 0")
	assertListener(res.ConvID)

	// Nothing is up for purging yet
	expectedPurgeInfo := chat1.EphemeralPurgeInfo{
		ConvID:          res.ConvID,
		MinUnexplodedID: 1,
		NextPurgeTime:   0,
		IsActive:        false,
	}
	assertTrackerState(res.ConvID, expectedPurgeInfo)

	// Send two ephemeral messages, and ensure both get purged
	lifetime := gregor1.DurationSec(1)
	lifetimeDuration := time.Second
	sendEphemeral(lifetime)
	sendEphemeral(lifetime * 2)

	thread, err := tc.ChatG.ConvSource.Pull(ctx, res.ConvID, uid, &chat1.GetThreadQuery{
		MessageTypes: []chat1.MessageType{chat1.MessageType_TEXT},
	}, nil)
	require.NoError(t, err)
	msgs := thread.Messages
	require.Len(t, msgs, 2)
	msgUnboxed2 := msgs[0]
	msgUnboxed1 := msgs[1]

	t.Logf("assert listener 1")
	world.Fc.Advance(lifetimeDuration)
	assertListener(res.ConvID)
	assertTrackerState(res.ConvID, chat1.EphemeralPurgeInfo{
		ConvID:          res.ConvID,
		MinUnexplodedID: msgUnboxed2.GetMessageID(),
		NextPurgeTime:   msgUnboxed2.Valid().Etime(),
		IsActive:        true,
	})
	assertEphemeralPurgeNotifInfo(res.ConvID, []chat1.MessageID{msgUnboxed1.GetMessageID()})

	// Ensure things run smoothly even with extraneous start/stop calls to purger
	g.EphemeralPurger.Start(context.Background(), uid)
	<-g.EphemeralPurger.Stop(context.Background())
	<-g.EphemeralPurger.Stop(context.Background())
	g.EphemeralPurger.Start(context.Background(), uid)
	g.EphemeralPurger.Start(context.Background(), uid)

	t.Logf("assert listener 2")
	world.Fc.Advance(lifetimeDuration)
	assertListener(res.ConvID)
	assertTrackerState(res.ConvID, chat1.EphemeralPurgeInfo{
		ConvID:          res.ConvID,
		MinUnexplodedID: msgUnboxed2.GetMessageID(),
		NextPurgeTime:   0,
		IsActive:        false,
	})
	assertEphemeralPurgeNotifInfo(res.ConvID, []chat1.MessageID{msgUnboxed2.GetMessageID()})
}
<|endoftext|>


<commit_before>package resources

import (
	"bytes"
	"encoding/json"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	gocf "github.com/mweagle/go-cloudformation"
	"github.com/sirupsen/logrus"
)

// S3ArtifactPublisherResourceRequest is what the UserProperties
// should be set to in the CustomResource invocation
type S3ArtifactPublisherResourceRequest struct {
	Bucket *gocf.StringExpr
	Key    *gocf.StringExpr
	Body   *gocf.StringExpr
}

// S3ArtifactPublisherResource is a simple POC showing how to create custom resources
type S3ArtifactPublisherResource struct {
	gocf.CloudFormationCustomResource
	S3ArtifactPublisherResourceRequest
}

// IAMPrivileges returns the IAM privs for this custom action
func (command *S3ArtifactPublisherResource) IAMPrivileges() []string {
	return []string{"s3:PutObject",
		"s3:DeleteObject"}
}

// Create implements the S3 create operation
func (command S3ArtifactPublisherResource) Create(awsSession *session.Session,
	event *CloudFormationLambdaEvent,
	logger *logrus.Logger) (map[string]interface{}, error) {

	unmarshalErr := json.Unmarshal(event.ResourceProperties, &command)
	if unmarshalErr != nil {
		return nil, unmarshalErr
	}
	itemInput := bytes.NewReader([]byte(command.Body.Literal))
	s3PutObjectParams := &s3.PutObjectInput{
		Body:   itemInput,
		Bucket: aws.String(command.Bucket.Literal),
		Key:    aws.String(command.Key.Literal),
	}
	s3Svc := s3.New(awsSession)
	s3Response, s3ResponseErr := s3Svc.PutObject(s3PutObjectParams)
	if s3ResponseErr != nil {
		return nil, s3ResponseErr
	}
	return map[string]interface{}{
		"ObjectVersion": s3Response.VersionId,
	}, nil
}

// Update implements the S3 update operation
func (command S3ArtifactPublisherResource) Update(awsSession *session.Session,
	event *CloudFormationLambdaEvent,
	logger *logrus.Logger) (map[string]interface{}, error) {
	return command.Create(awsSession, event, logger)
}

// Delete implements the S3 delete operation
func (command S3ArtifactPublisherResource) Delete(awsSession *session.Session,
	event *CloudFormationLambdaEvent,
	logger *logrus.Logger) (map[string]interface{}, error) {

	unmarshalErr := json.Unmarshal(event.ResourceProperties, &command)
	if unmarshalErr != nil {
		return nil, unmarshalErr
	}
	s3DeleteObjectParams := &s3.DeleteObjectInput{
		Bucket: aws.String(command.Bucket.Literal),
		Key:    aws.String(command.Key.Literal),
	}
	s3Svc := s3.New(awsSession)
	_, s3ResponseErr := s3Svc.DeleteObject(s3DeleteObjectParams)
	if s3ResponseErr != nil {
		return nil, s3ResponseErr
	}
	logger.WithFields(logrus.Fields{
		"Bucket": command.Bucket.Literal,
		"Key":    command.Key.Literal,
	}).Info("Object deleted")
	return nil, nil
}
<commit_msg>Prefer map[string]interface{} for S3 body<commit_after>package resources

import (
	"bytes"
	"encoding/json"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	gocf "github.com/mweagle/go-cloudformation"
	"github.com/sirupsen/logrus"
)

// S3ArtifactPublisherResourceRequest is what the UserProperties
// should be set to in the CustomResource invocation
type S3ArtifactPublisherResourceRequest struct {
	Bucket *gocf.StringExpr
	Key    *gocf.StringExpr
	Body   map[string]interface{}
}

// S3ArtifactPublisherResource is a simple POC showing how to create custom resources
type S3ArtifactPublisherResource struct {
	gocf.CloudFormationCustomResource
	S3ArtifactPublisherResourceRequest
}

// IAMPrivileges returns the IAM privs for this custom action
func (command *S3ArtifactPublisherResource) IAMPrivileges() []string {
	return []string{"s3:PutObject",
		"s3:DeleteObject"}
}

// Create implements the S3 create operation
func (command S3ArtifactPublisherResource) Create(awsSession *session.Session,
	event *CloudFormationLambdaEvent,
	logger *logrus.Logger) (map[string]interface{}, error) {

	unmarshalErr := json.Unmarshal(event.ResourceProperties, &command)
	if unmarshalErr != nil {
		return nil, unmarshalErr
	}
	mapData, mapDataErr := json.Marshal(command.Body)
	if mapDataErr != nil {
		return nil, mapDataErr
	}
	itemInput := bytes.NewReader(mapData)
	s3PutObjectParams := &s3.PutObjectInput{
		Body:   itemInput,
		Bucket: aws.String(command.Bucket.Literal),
		Key:    aws.String(command.Key.Literal),
	}
	s3Svc := s3.New(awsSession)
	s3Response, s3ResponseErr := s3Svc.PutObject(s3PutObjectParams)
	if s3ResponseErr != nil {
		return nil, s3ResponseErr
	}
	return map[string]interface{}{
		"ObjectVersion": s3Response.VersionId,
	}, nil
}

// Update implements the S3 update operation
func (command S3ArtifactPublisherResource) Update(awsSession *session.Session,
	event *CloudFormationLambdaEvent,
	logger *logrus.Logger) (map[string]interface{}, error) {
	return command.Create(awsSession, event, logger)
}

// Delete implements the S3 delete operation
func (command S3ArtifactPublisherResource) Delete(awsSession *session.Session,
	event *CloudFormationLambdaEvent,
	logger *logrus.Logger) (map[string]interface{}, error) {

	unmarshalErr := json.Unmarshal(event.ResourceProperties, &command)
	if unmarshalErr != nil {
		return nil, unmarshalErr
	}
	s3DeleteObjectParams := &s3.DeleteObjectInput{
		Bucket: aws.String(command.Bucket.Literal),
		Key:    aws.String(command.Key.Literal),
	}
	s3Svc := s3.New(awsSession)
	_, s3ResponseErr := s3Svc.DeleteObject(s3DeleteObjectParams)
	if s3ResponseErr != nil {
		return nil, s3ResponseErr
	}
	logger.WithFields(logrus.Fields{
		"Bucket": command.Bucket.Literal,
		"Key":    command.Key.Literal,
	}).Info("Object deleted")
	return nil, nil
}
<|endoftext|>


<commit_before>package vagrant

import (
	"bytes"
	"compress/gzip"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/url"
	"strconv"
	"time"

	"koding/db/models"
	"koding/db/mongodb/modelhelper"
	"koding/kites/kloud/api/vagrantapi"
	puser "koding/kites/kloud/scripts/provisionklient/userdata"
	"koding/kites/kloud/stackplan"
	"koding/kites/kloud/utils"

	"github.com/satori/go.uuid"
	"gopkg.in/mgo.v2/bson"
)

// VagrantResource represents vagrant_instance Terraform resource.
type VagrantResource struct {
	Build map[string]map[string]interface{} `hcl:"vagrant_instance"`
}

func (s *Stack) updateCredential(cred *stackplan.Credential) {
	meta := cred.Meta.(*VagrantMeta)
	meta.SetDefaults()
}

type Tunnel struct {
	Name    string
	KiteURL string
}

func (s *Stack) newTunnel(resourceName string) *Tunnel {
	t := &Tunnel{
		Name:    utils.RandString(12),
		KiteURL: s.TunnelURL.String(),
	}

	if m, ok := s.Builder.Machines[resourceName]; ok {
		t.Name = m.Uid
	}

	return t
}

// InjectVagrantData sets default properties for vagrant_instance Terraform template.
//
// TODO(rjeczalik): move out hostQueryString outside this method.
func (s *Stack) InjectVagrantData() (string, stackplan.KiteMap, error) {
	// TODO(rjeczalik): add ByProvider to stackplan.Credentials
	for _, c := range s.Builder.Credentials {
		if c.Provider == "vagrant" {
			s.Credential = c
			break
		}
	}

	if s.Credential == nil {
		return "", nil, errors.New("vagrant credential not found")
	}

	s.updateCredential(s.Credential)

	t := s.Builder.Template
	meta := s.Credential.Meta.(*VagrantMeta)

	s.Log.Debug("Injecting vagrant credentials: %# v", meta)

	if err := t.InjectVariables(s.Credential.Provider, s.Credential.Meta); err != nil {
		return "", nil, err
	}

	var res VagrantResource

	if err := t.DecodeResource(&res); err != nil {
		return "", nil, err
	}

	if len(res.Build) == 0 {
		return "", nil, errors.New("no vagrant instances specified")
	}

	kiteIDs := make(stackplan.KiteMap)

	uids := s.Builder.MachineUIDs()

	s.Log.Debug("machine uids (%d): %v", len(uids), uids)

	klientURL, err := s.Session.Userdata.LookupKlientURL()
	if err != nil {
		return "", nil, err
	}

	// queryString is taken from jCredentialData.meta.queryString,
	// for debugging purposes it can be overwritten in the template,
	// however if template has multiple machines, all of them
	// are required to overwrite the queryString to the same value
	// to match current implementation of terraformplugins/vagrant
	// provider.
	//
	// Otherwise we fail early to show problem with the template.
	var queryString string
	for resourceName, box := range res.Build {
		kiteID := uuid.NewV4().String()

		kiteKey, err := s.Session.Userdata.Keycreator.Create(s.Req.Username, kiteID)
		if err != nil {
			return "", nil, err
		}

		// set kontrolURL if not provided via template
		kontrolURL := s.Session.Userdata.Keycreator.KontrolURL
		if k, ok := box["kontrolURL"].(string); ok {
			kontrolURL = k
		} else {
			box["kontrolURL"] = kontrolURL
		}

		if q, ok := box["queryString"].(string); ok {
			q = utils.QueryString(q)
			if queryString != "" && queryString != q {
				return "", nil, fmt.Errorf("mismatched queryString provided for multiple instances; want %q, got %q", queryString, q)
			}
			queryString = q
		} else {
			box["queryString"] = "${var.vagrant_queryString}"
			queryString = meta.QueryString
		}

		// set default filePath to relative <stackdir>/<boxname>; for
		// default configured klient it resolves to ~/.vagrant.d/<stackdir>/<boxname>
		if _, ok := box["filePath"]; !ok {
			uid, ok := uids[resourceName]
			if !ok {
				// For Plan call we return random uid as it won't be returned
				// as a part of meta; the filePath is inserted into meta by
				// the apply method.
				uid = resourceName + "-" + utils.RandString(6)
			}
			box["filePath"] = "koding/${var.koding_group_slug}/" + uid
		}

		// set default CPU number
		if _, ok := box["cpus"]; !ok {
			box["cpus"] = "${var.vagrant_cpus}"
		}

		// set default RAM in MiB
		if _, ok := box["memory"]; !ok {
			box["memory"] = "${var.vagrant_memory}"
		}

		// set default box type
		if _, ok := box["box"]; !ok {
			box["box"] = "${var.vagrant_box}"
		}

		ports, ok := box["forwarded_ports"].([]interface{})
		if !ok {
			ports = make([]interface{}, 0)
		}

		// klient kite port
		kitePort := &vagrantapi.ForwardedPort{
			HostPort:  2200,
			GuestPort: 56789,
		}

		// tlsproxy port
		kitesPort := &vagrantapi.ForwardedPort{
			HostPort:  2201,
			GuestPort: 56790,
		}

		ports = append(ports, kitePort, kitesPort)

		box["forwarded_ports"] = ports
		box["username"] = s.Req.Username

		tunnel := s.newTunnel(resourceName)

		s.Builder.BuildUserData(box, resourceName)

		if b, ok := box["debug"].(bool); ok && b {
			s.Debug = true
		}

		data := puser.Value{
			Username:        s.Req.Username,
			Groups:          []string{"sudo"},
			Hostname:        s.Req.Username, // no typo here. hostname = username
			KiteKey:         kiteKey,
			LatestKlientURL: klientURL,
			TunnelName:      tunnel.Name,
			TunnelKiteURL:   tunnel.KiteURL,
			KontrolURL:      kontrolURL,
		}

		// pass the values as a JSON encoded as bae64. Our script will decode
		// and unmarshall and use it inside the Vagrant box
		val, err := json.Marshal(&data)
		if err != nil {
			return "", nil, err
		}

		// Compressing the provision data isn't doing any serious optimizations,
		// it's just here so the debug output does not take half a screen.
		//
		// The provisionclient handles both compressed and uncompressed JSONs.
		var buf bytes.Buffer
		if cw, err := gzip.NewWriterLevel(&buf, 9); err == nil {
			if _, err = io.Copy(cw, bytes.NewReader(val)); err == nil && cw.Close() == nil {
				s.Log.Debug("using compressed provision data: %d vs %d", len(val), len(buf.Bytes()))

				val = buf.Bytes()
			}
		}

		box["provisionData"] = base64.StdEncoding.EncodeToString(val)
		kiteIDs[resourceName] = kiteID
		res.Build[resourceName] = box
	}

	t.Resource["vagrant_instance"] = res.Build

	if err := t.Flush(); err != nil {
		return "", nil, err
	}

	return queryString, kiteIDs, nil
}

func (s *Stack) machinesFromTemplate(t *stackplan.Template) (stackplan.Machines, error) {
	var res VagrantResource
	if err := t.DecodeResource(&res); err != nil {
		return nil, err
	}

	machines := make(stackplan.Machines, len(res.Build))

	for label, box := range res.Build {
		m := &stackplan.Machine{
			Provider:   "vagrant",
			Label:      label,
			Attributes: make(map[string]string),
		}

		if cpus, ok := box["cpus"].(string); ok && !stackplan.IsVariable(cpus) {
			m.Attributes["cpus"] = cpus
		}
		if mem, ok := box["memory"].(string); ok && !stackplan.IsVariable(mem) {
			m.Attributes["memory"] = mem
		}
		if typ, ok := box["box"].(string); ok && !stackplan.IsVariable(typ) {
			m.Attributes["box"] = typ
		}

		if _, ok := machines[label]; ok {
			return nil, errors.New("duplicate instance labels: " + label)
		}

		machines[label] = m
	}

	return machines, nil
}

func (s *Stack) updateMachines(machines stackplan.Machines, jMachines map[string]*models.Machine) error {
	for label, machine := range jMachines {
		s.Log.Debug("Updating machine with %q label and %q provider", label, machine.Provider)

		tf, ok := machines[label]
		if !ok {
			return fmt.Errorf("machine label '%s' doesn't exist in terraform output", label)
		}

		if tf.Provider == "vagrant" {
			if err := updateVagrant(tf, machine.ObjectId, s.Credential.Identifier); err != nil {
				return stackplan.ResError(err, "jMachine")
			}
		}
	}

	return nil
}

func updateVagrant(tf *stackplan.Machine, machineId bson.ObjectId, credential string) error {
	machine := bson.M{
		"provider":           tf.Provider,
		"queryString":        tf.QueryString,
		"credential":         credential,
		"ipAddress":          tf.Attributes["ipAddress"],
		"meta.filePath":      tf.Attributes["filePath"],
		"meta.box":           tf.Attributes["box"],
		"meta.hostname":      tf.Attributes["hostname"],
		"meta.klientHostURL": tf.Attributes["klientHostURL"],
		"status.modifiedAt":  time.Now().UTC(),
		"status.state":       tf.State.String(),
		"status.reason":      tf.StateReason,
	}

	if u, err := url.Parse(tf.RegisterURL); tf.RegisterURL != "" && err == nil {
		u.Path = "/klient/kite"
		machine["meta.klientGuestURL"] = u.String()
		machine["domain"] = u.Host
	}

	if n, err := strconv.Atoi(tf.Attributes["memory"]); err == nil {
		machine["meta.memory"] = n
	}

	if n, err := strconv.Atoi(tf.Attributes["cpus"]); err == nil {
		machine["meta.cpus"] = n
	}

	return modelhelper.UpdateMachine(machineId, bson.M{"$set": machine})
}
<commit_msg>provider/vagrant: fix appending to user-defined ports<commit_after>package vagrant

import (
	"bytes"
	"compress/gzip"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/url"
	"strconv"
	"time"

	"koding/db/models"
	"koding/db/mongodb/modelhelper"
	"koding/kites/kloud/api/vagrantapi"
	puser "koding/kites/kloud/scripts/provisionklient/userdata"
	"koding/kites/kloud/stackplan"
	"koding/kites/kloud/utils"

	"github.com/satori/go.uuid"
	"gopkg.in/mgo.v2/bson"
)

// VagrantResource represents vagrant_instance Terraform resource.
type VagrantResource struct {
	Build map[string]map[string]interface{} `hcl:"vagrant_instance"`
}

func (s *Stack) updateCredential(cred *stackplan.Credential) {
	meta := cred.Meta.(*VagrantMeta)
	meta.SetDefaults()
}

type Tunnel struct {
	Name    string
	KiteURL string
}

func (s *Stack) newTunnel(resourceName string) *Tunnel {
	t := &Tunnel{
		Name:    utils.RandString(12),
		KiteURL: s.TunnelURL.String(),
	}

	if m, ok := s.Builder.Machines[resourceName]; ok {
		t.Name = m.Uid
	}

	return t
}

// InjectVagrantData sets default properties for vagrant_instance Terraform template.
//
// TODO(rjeczalik): move out hostQueryString outside this method.
func (s *Stack) InjectVagrantData() (string, stackplan.KiteMap, error) {
	// TODO(rjeczalik): add ByProvider to stackplan.Credentials
	for _, c := range s.Builder.Credentials {
		if c.Provider == "vagrant" {
			s.Credential = c
			break
		}
	}

	if s.Credential == nil {
		return "", nil, errors.New("vagrant credential not found")
	}

	s.updateCredential(s.Credential)

	t := s.Builder.Template
	meta := s.Credential.Meta.(*VagrantMeta)

	s.Log.Debug("Injecting vagrant credentials: %# v", meta)

	if err := t.InjectVariables(s.Credential.Provider, s.Credential.Meta); err != nil {
		return "", nil, err
	}

	var res VagrantResource

	if err := t.DecodeResource(&res); err != nil {
		return "", nil, err
	}

	if len(res.Build) == 0 {
		return "", nil, errors.New("no vagrant instances specified")
	}

	kiteIDs := make(stackplan.KiteMap)

	uids := s.Builder.MachineUIDs()

	s.Log.Debug("machine uids (%d): %v", len(uids), uids)

	klientURL, err := s.Session.Userdata.LookupKlientURL()
	if err != nil {
		return "", nil, err
	}

	// queryString is taken from jCredentialData.meta.queryString,
	// for debugging purposes it can be overwritten in the template,
	// however if template has multiple machines, all of them
	// are required to overwrite the queryString to the same value
	// to match current implementation of terraformplugins/vagrant
	// provider.
	//
	// Otherwise we fail early to show problem with the template.
	var queryString string
	for resourceName, box := range res.Build {
		kiteID := uuid.NewV4().String()

		kiteKey, err := s.Session.Userdata.Keycreator.Create(s.Req.Username, kiteID)
		if err != nil {
			return "", nil, err
		}

		// set kontrolURL if not provided via template
		kontrolURL := s.Session.Userdata.Keycreator.KontrolURL
		if k, ok := box["kontrolURL"].(string); ok {
			kontrolURL = k
		} else {
			box["kontrolURL"] = kontrolURL
		}

		if q, ok := box["queryString"].(string); ok {
			q = utils.QueryString(q)
			if queryString != "" && queryString != q {
				return "", nil, fmt.Errorf("mismatched queryString provided for multiple instances; want %q, got %q", queryString, q)
			}
			queryString = q
		} else {
			box["queryString"] = "${var.vagrant_queryString}"
			queryString = meta.QueryString
		}

		// set default filePath to relative <stackdir>/<boxname>; for
		// default configured klient it resolves to ~/.vagrant.d/<stackdir>/<boxname>
		if _, ok := box["filePath"]; !ok {
			uid, ok := uids[resourceName]
			if !ok {
				// For Plan call we return random uid as it won't be returned
				// as a part of meta; the filePath is inserted into meta by
				// the apply method.
				uid = resourceName + "-" + utils.RandString(6)
			}
			box["filePath"] = "koding/${var.koding_group_slug}/" + uid
		}

		// set default CPU number
		if _, ok := box["cpus"]; !ok {
			box["cpus"] = "${var.vagrant_cpus}"
		}

		// set default RAM in MiB
		if _, ok := box["memory"]; !ok {
			box["memory"] = "${var.vagrant_memory}"
		}

		// set default box type
		if _, ok := box["box"]; !ok {
			box["box"] = "${var.vagrant_box}"
		}

		var ports []interface{}

		switch p := box["forwarded_ports"].(type) {
		case []interface{}:
			ports = p
		case []map[string]interface{}:
			ports = make([]interface{}, len(p))

			for i := range p {
				ports[i] = p[i]
			}
		}

		// klient kite port
		kitePort := &vagrantapi.ForwardedPort{
			HostPort:  2200,
			GuestPort: 56789,
		}

		// tlsproxy port
		kitesPort := &vagrantapi.ForwardedPort{
			HostPort:  2201,
			GuestPort: 56790,
		}

		ports = append(ports, kitePort, kitesPort)

		box["forwarded_ports"] = ports
		box["username"] = s.Req.Username

		tunnel := s.newTunnel(resourceName)

		s.Builder.BuildUserData(box, resourceName)

		if b, ok := box["debug"].(bool); ok && b {
			s.Debug = true
		}

		data := puser.Value{
			Username:        s.Req.Username,
			Groups:          []string{"sudo"},
			Hostname:        s.Req.Username, // no typo here. hostname = username
			KiteKey:         kiteKey,
			LatestKlientURL: klientURL,
			TunnelName:      tunnel.Name,
			TunnelKiteURL:   tunnel.KiteURL,
			KontrolURL:      kontrolURL,
		}

		// pass the values as a JSON encoded as bae64. Our script will decode
		// and unmarshall and use it inside the Vagrant box
		val, err := json.Marshal(&data)
		if err != nil {
			return "", nil, err
		}

		// Compressing the provision data isn't doing any serious optimizations,
		// it's just here so the debug output does not take half a screen.
		//
		// The provisionclient handles both compressed and uncompressed JSONs.
		var buf bytes.Buffer
		if cw, err := gzip.NewWriterLevel(&buf, 9); err == nil {
			if _, err = io.Copy(cw, bytes.NewReader(val)); err == nil && cw.Close() == nil {
				s.Log.Debug("using compressed provision data: %d vs %d", len(val), len(buf.Bytes()))

				val = buf.Bytes()
			}
		}

		box["provisionData"] = base64.StdEncoding.EncodeToString(val)
		kiteIDs[resourceName] = kiteID
		res.Build[resourceName] = box
	}

	t.Resource["vagrant_instance"] = res.Build

	if err := t.Flush(); err != nil {
		return "", nil, err
	}

	return queryString, kiteIDs, nil
}

func (s *Stack) machinesFromTemplate(t *stackplan.Template) (stackplan.Machines, error) {
	var res VagrantResource
	if err := t.DecodeResource(&res); err != nil {
		return nil, err
	}

	machines := make(stackplan.Machines, len(res.Build))

	for label, box := range res.Build {
		m := &stackplan.Machine{
			Provider:   "vagrant",
			Label:      label,
			Attributes: make(map[string]string),
		}

		if cpus, ok := box["cpus"].(string); ok && !stackplan.IsVariable(cpus) {
			m.Attributes["cpus"] = cpus
		}
		if mem, ok := box["memory"].(string); ok && !stackplan.IsVariable(mem) {
			m.Attributes["memory"] = mem
		}
		if typ, ok := box["box"].(string); ok && !stackplan.IsVariable(typ) {
			m.Attributes["box"] = typ
		}

		if _, ok := machines[label]; ok {
			return nil, errors.New("duplicate instance labels: " + label)
		}

		machines[label] = m
	}

	return machines, nil
}

func (s *Stack) updateMachines(machines stackplan.Machines, jMachines map[string]*models.Machine) error {
	for label, machine := range jMachines {
		s.Log.Debug("Updating machine with %q label and %q provider", label, machine.Provider)

		tf, ok := machines[label]
		if !ok {
			return fmt.Errorf("machine label '%s' doesn't exist in terraform output", label)
		}

		if tf.Provider == "vagrant" {
			if err := updateVagrant(tf, machine.ObjectId, s.Credential.Identifier); err != nil {
				return stackplan.ResError(err, "jMachine")
			}
		}
	}

	return nil
}

func updateVagrant(tf *stackplan.Machine, machineId bson.ObjectId, credential string) error {
	machine := bson.M{
		"provider":           tf.Provider,
		"queryString":        tf.QueryString,
		"credential":         credential,
		"ipAddress":          tf.Attributes["ipAddress"],
		"meta.filePath":      tf.Attributes["filePath"],
		"meta.box":           tf.Attributes["box"],
		"meta.hostname":      tf.Attributes["hostname"],
		"meta.klientHostURL": tf.Attributes["klientHostURL"],
		"status.modifiedAt":  time.Now().UTC(),
		"status.state":       tf.State.String(),
		"status.reason":      tf.StateReason,
	}

	if u, err := url.Parse(tf.RegisterURL); tf.RegisterURL != "" && err == nil {
		u.Path = "/klient/kite"
		machine["meta.klientGuestURL"] = u.String()
		machine["domain"] = u.Host
	}

	if n, err := strconv.Atoi(tf.Attributes["memory"]); err == nil {
		machine["meta.memory"] = n
	}

	if n, err := strconv.Atoi(tf.Attributes["cpus"]); err == nil {
		machine["meta.cpus"] = n
	}

	return modelhelper.UpdateMachine(machineId, bson.M{"$set": machine})
}
<|endoftext|>