<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"log"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/cloudfoundry-incubator/bbs"
	"github.com/cloudfoundry-incubator/cf-debug-server"
	cf_lager "github.com/cloudfoundry-incubator/cf-lager"
	"github.com/cloudfoundry-incubator/cf_http"
	"github.com/cloudfoundry-incubator/consuladapter"
	"github.com/cloudfoundry-incubator/executor"
	executorinit "github.com/cloudfoundry-incubator/executor/initializer"
	"github.com/cloudfoundry-incubator/locket"
	"github.com/cloudfoundry-incubator/rep"
	"github.com/cloudfoundry-incubator/rep/auction_cell_rep"
	"github.com/cloudfoundry-incubator/rep/evacuation"
	"github.com/cloudfoundry-incubator/rep/evacuation/evacuation_context"
	"github.com/cloudfoundry-incubator/rep/generator"
	"github.com/cloudfoundry-incubator/rep/handlers"
	"github.com/cloudfoundry-incubator/rep/harmonizer"
	"github.com/cloudfoundry-incubator/rep/lrp_stopper"
	"github.com/cloudfoundry-incubator/rep/maintain"
	"github.com/cloudfoundry/dropsonde"
	"github.com/nu7hatch/gouuid"
	"github.com/pivotal-golang/clock"
	"github.com/pivotal-golang/lager"
	"github.com/pivotal-golang/localip"
	"github.com/pivotal-golang/operationq"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/http_server"
	"github.com/tedsuo/ifrit/sigmon"
	"github.com/tedsuo/rata"
)

var sessionName = flag.String(
	"sessionName",
	"rep",
	"consul session name",
)

var consulCluster = flag.String(
	"consulCluster",
	"",
	"comma-separated list of consul server URLs (scheme://ip:port)",
)

var lockTTL = flag.Duration(
	"lockTTL",
	locket.LockTTL,
	"TTL for service lock",
)

var lockRetryInterval = flag.Duration(
	"lockRetryInterval",
	locket.RetryInterval,
	"interval to wait before retrying a failed lock acquisition",
)

var listenAddr = flag.String(
	"listenAddr",
	"0.0.0.0:1800",
	"host:port to serve auction and LRP stop requests on",
)

var cellID = flag.String(
	"cellID",
	"",
	"the ID used by the rep to identify itself to external systems - must be specified",
)

var zone = flag.String(
	"zone",
	"",
	"the availability zone associated with the rep",
)

var pollingInterval = flag.Duration(
	"pollingInterval",
	30*time.Second,
	"the interval on which to scan the executor",
)

var communicationTimeout = flag.Duration(
	"communicationTimeout",
	10*time.Second,
	"Timeout applied to all HTTP requests.",
)

var evacuationTimeout = flag.Duration(
	"evacuationTimeout",
	10*time.Minute,
	"Timeout to wait for evacuation to complete",
)

var evacuationPollingInterval = flag.Duration(
	"evacuationPollingInterval",
	10*time.Second,
	"the interval on which to scan the executor during evacuation",
)

var bbsAddress = flag.String(
	"bbsAddress",
	"",
	"Address to the BBS Server",
)

var bbsCACert = flag.String(
	"bbsCACert",
	"",
	"path to certificate authority cert used for mutually authenticated TLS BBS communication",
)

var bbsClientCert = flag.String(
	"bbsClientCert",
	"",
	"path to client cert used for mutually authenticated TLS BBS communication",
)

var bbsClientKey = flag.String(
	"bbsClientKey",
	"",
	"path to client key used for mutually authenticated TLS BBS communication",
)

type stackPathMap rep.StackPathMap

func (s *stackPathMap) String() string {
	return fmt.Sprintf("%v", *s)
}

func (s *stackPathMap) Set(value string) error {
	parts := strings.SplitN(value, ":", 2)
	if len(parts) != 2 {
		return errors.New("Invalid preloaded RootFS value: not of the form 'stack-name:path'")
	}

	if parts[0] == "" {
		return errors.New("Invalid preloaded RootFS value: blank stack")
	}

	if parts[1] == "" {
		return errors.New("Invalid preloaded RootFS value: blank path")
	}

	(*s)[parts[0]] = parts[1]
	return nil
}

type providers []string

func (p *providers) String() string {
	return fmt.Sprintf("%v", *p)
}

func (p *providers) Set(value string) error {
	if value == "" {
		return errors.New("Cannot set blank value for RootFS provider")
	}

	*p = append(*p, value)
	return nil
}

const (
	dropsondeDestination = "localhost:3457"
	dropsondeOrigin      = "rep"

	bbsPingTimeout = 5 * time.Minute
)

func main() {
	cf_debug_server.AddFlags(flag.CommandLine)
	cf_lager.AddFlags(flag.CommandLine)

	stackMap := stackPathMap{}
	supportedProviders := providers{}
	flag.Var(&stackMap, "preloadedRootFS", "List of preloaded RootFSes")
	flag.Var(&supportedProviders, "rootFSProvider", "List of RootFS providers")
	flag.Parse()

	cf_http.Initialize(*communicationTimeout)

	clock := clock.NewClock()
	logger, reconfigurableSink := cf_lager.New(*sessionName)

	executorConfiguration := executorConfig()
	if !executorinit.ValidateExecutor(logger, executorConfiguration) {
		os.Exit(1)
	}

	initializeDropsonde(logger)

	if *cellID == "" {
		log.Fatalf("-cellID must be specified")
	}

	executorClient, executorMembers, err := executorinit.Initialize(logger, executorConfiguration, clock)
	if err != nil {
		log.Fatalf("Failed to initialize executor: %s", err.Error())
	}
	defer executorClient.Cleanup()

	if err := validateBBSAddress(); err != nil {
		logger.Fatal("invalid-bbs-address", err)
	}

	locketClient := initializeLocketClient(logger)

	evacuatable, evacuationReporter, evacuationNotifier := evacuation_context.New()

	// only one outstanding operation per container is necessary
	queue := operationq.NewSlidingQueue(1)

	evacuator := evacuation.NewEvacuator(
		logger,
		clock,
		executorClient,
		evacuationNotifier,
		*cellID,
		*evacuationTimeout,
		*evacuationPollingInterval,
	)

	bbsClient := initializeBBSClient(logger)
	httpServer, address := initializeServer(bbsClient, executorClient, evacuatable, evacuationReporter, logger, rep.StackPathMap(stackMap), supportedProviders)
	opGenerator := generator.New(*cellID, bbsClient, executorClient, evacuationReporter, uint64(evacuationTimeout.Seconds()))

	preloadedRootFSes := []string{}
	for k := range stackMap {
		preloadedRootFSes = append(preloadedRootFSes, k)
	}

	members := grouper.Members{
		{"http_server", httpServer},
		{"presence", initializeCellPresence(address, locketClient, executorClient, logger, supportedProviders, preloadedRootFSes)},
		{"bulker", harmonizer.NewBulker(logger, *pollingInterval, *evacuationPollingInterval, evacuationNotifier, clock, opGenerator, queue)},
		{"event-consumer", harmonizer.NewEventConsumer(logger, opGenerator, queue)},
		{"evacuator", evacuator},
	}

	members = append(executorMembers, members...)

	if dbgAddr := cf_debug_server.DebugAddress(flag.CommandLine); dbgAddr != "" {
		members = append(grouper.Members{
			{"debug-server", cf_debug_server.Runner(dbgAddr, reconfigurableSink)},
		}, members...)
	}

	group := grouper.NewOrdered(os.Interrupt, members)

	monitor := ifrit.Invoke(sigmon.New(group))

	logger.Info("started", lager.Data{"cell-id": *cellID})

	err = <-monitor.Wait()
	if err != nil {
		logger.Error("exited-with-failure", err)
		os.Exit(1)
	}

	logger.Info("exited")
}

func initializeDropsonde(logger lager.Logger) {
	err := dropsonde.Initialize(dropsondeDestination, dropsondeOrigin)
	if err != nil {
		logger.Error("failed to initialize dropsonde: %v", err)
	}
}

func initializeCellPresence(address string, locketClient locket.Client, executorClient executor.Client, logger lager.Logger, rootFSProviders, preloadedRootFSes []string) ifrit.Runner {
	config := maintain.Config{
		CellID:            *cellID,
		RepAddress:        address,
		Zone:              *zone,
		RetryInterval:     *lockRetryInterval,
		RootFSProviders:   rootFSProviders,
		PreloadedRootFSes: preloadedRootFSes,
	}
	return maintain.New(config, executorClient, locketClient, logger, clock.NewClock())
}

func initializeLocketClient(logger lager.Logger) locket.Client {
	client, err := consuladapter.NewClient(*consulCluster)
	if err != nil {
		logger.Fatal("new-client-failed", err)
	}

	sessionMgr := consuladapter.NewSessionManager(client)
	consulSession, err := consuladapter.NewSessionNoChecks(*sessionName, *lockTTL, client, sessionMgr)
	if err != nil {
		logger.Fatal("consul-session-failed", err)
	}

	return locket.NewClient(consulSession, clock.NewClock(), logger)
}

func initializeLRPStopper(guid string, executorClient executor.Client, logger lager.Logger) lrp_stopper.LRPStopper {
	return lrp_stopper.New(guid, executorClient, logger)
}

func initializeServer(
	bbsClient bbs.Client,
	executorClient executor.Client,
	evacuatable evacuation_context.Evacuatable,
	evacuationReporter evacuation_context.EvacuationReporter,
	logger lager.Logger,
	stackMap rep.StackPathMap,
	supportedProviders []string,
) (ifrit.Runner, string) {
	lrpStopper := initializeLRPStopper(*cellID, executorClient, logger)

	auctionCellRep := auction_cell_rep.New(*cellID, stackMap, supportedProviders, *zone, generateGuid, executorClient, evacuationReporter, logger)
	handlers := handlers.New(auctionCellRep, lrpStopper, executorClient, evacuatable, logger)

	router, err := rata.NewRouter(rep.Routes, handlers)
	if err != nil {
		logger.Fatal("failed-to-construct-router", err)
	}

	ip, err := localip.LocalIP()
	if err != nil {
		logger.Fatal("failed-to-fetch-ip", err)
	}

	port := strings.Split(*listenAddr, ":")[1]
	address := fmt.Sprintf("http://%s:%s", ip, port)

	return http_server.New(*listenAddr, router), address
}

func validateBBSAddress() error {
	if *bbsAddress == "" {
		return errors.New("bbsAddress is required")
	}
	return nil
}

func generateGuid() (string, error) {
	guid, err := uuid.NewV4()
	if err != nil {
		return "", err
	}
	return guid.String(), nil
}

func initializeBBSClient(logger lager.Logger) bbs.Client {
	bbsURL, err := url.Parse(*bbsAddress)
	if err != nil {
		logger.Fatal("Invalid BBS URL", err)
	}

	if bbsURL.Scheme != "https" {
		return bbs.NewClient(*bbsAddress)
	}

	bbsClient, err := bbs.NewSecureClient(*bbsAddress, *bbsCACert, *bbsClientCert, *bbsClientKey)
	if err != nil {
		logger.Fatal("Failed to configure secure BBS client", err)
	}
	return bbsClient
}
<commit_msg>Add flags to configure BBS HTTP client<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"log"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/cloudfoundry-incubator/bbs"
	"github.com/cloudfoundry-incubator/cf-debug-server"
	cf_lager "github.com/cloudfoundry-incubator/cf-lager"
	"github.com/cloudfoundry-incubator/cf_http"
	"github.com/cloudfoundry-incubator/consuladapter"
	"github.com/cloudfoundry-incubator/executor"
	executorinit "github.com/cloudfoundry-incubator/executor/initializer"
	"github.com/cloudfoundry-incubator/locket"
	"github.com/cloudfoundry-incubator/rep"
	"github.com/cloudfoundry-incubator/rep/auction_cell_rep"
	"github.com/cloudfoundry-incubator/rep/evacuation"
	"github.com/cloudfoundry-incubator/rep/evacuation/evacuation_context"
	"github.com/cloudfoundry-incubator/rep/generator"
	"github.com/cloudfoundry-incubator/rep/handlers"
	"github.com/cloudfoundry-incubator/rep/harmonizer"
	"github.com/cloudfoundry-incubator/rep/lrp_stopper"
	"github.com/cloudfoundry-incubator/rep/maintain"
	"github.com/cloudfoundry/dropsonde"
	"github.com/nu7hatch/gouuid"
	"github.com/pivotal-golang/clock"
	"github.com/pivotal-golang/lager"
	"github.com/pivotal-golang/localip"
	"github.com/pivotal-golang/operationq"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/http_server"
	"github.com/tedsuo/ifrit/sigmon"
	"github.com/tedsuo/rata"
)

var sessionName = flag.String(
	"sessionName",
	"rep",
	"consul session name",
)

var consulCluster = flag.String(
	"consulCluster",
	"",
	"comma-separated list of consul server URLs (scheme://ip:port)",
)

var lockTTL = flag.Duration(
	"lockTTL",
	locket.LockTTL,
	"TTL for service lock",
)

var lockRetryInterval = flag.Duration(
	"lockRetryInterval",
	locket.RetryInterval,
	"interval to wait before retrying a failed lock acquisition",
)

var listenAddr = flag.String(
	"listenAddr",
	"0.0.0.0:1800",
	"host:port to serve auction and LRP stop requests on",
)

var cellID = flag.String(
	"cellID",
	"",
	"the ID used by the rep to identify itself to external systems - must be specified",
)

var zone = flag.String(
	"zone",
	"",
	"the availability zone associated with the rep",
)

var pollingInterval = flag.Duration(
	"pollingInterval",
	30*time.Second,
	"the interval on which to scan the executor",
)

var communicationTimeout = flag.Duration(
	"communicationTimeout",
	10*time.Second,
	"Timeout applied to all HTTP requests.",
)

var evacuationTimeout = flag.Duration(
	"evacuationTimeout",
	10*time.Minute,
	"Timeout to wait for evacuation to complete",
)

var evacuationPollingInterval = flag.Duration(
	"evacuationPollingInterval",
	10*time.Second,
	"the interval on which to scan the executor during evacuation",
)

var bbsAddress = flag.String(
	"bbsAddress",
	"",
	"Address to the BBS Server",
)

var bbsCACert = flag.String(
	"bbsCACert",
	"",
	"path to certificate authority cert used for mutually authenticated TLS BBS communication",
)

var bbsClientCert = flag.String(
	"bbsClientCert",
	"",
	"path to client cert used for mutually authenticated TLS BBS communication",
)

var bbsClientKey = flag.String(
	"bbsClientKey",
	"",
	"path to client key used for mutually authenticated TLS BBS communication",
)

var bbsClientSessionCacheSize = flag.Int(
	"bbsClientSessionCacheSize",
	0,
	"Capacity of the ClientSessionCache option on the TLS configuration. If zero, golang's default will be used",
)

var bbsMaxIdleConnsPerHost = flag.Int(
	"bbsMaxIdleConnsPerHost",
	0,
	"Controls the maximum number of idle (keep-alive) connctions per host. If zero, golang's default will be used",
)

type stackPathMap rep.StackPathMap

func (s *stackPathMap) String() string {
	return fmt.Sprintf("%v", *s)
}

func (s *stackPathMap) Set(value string) error {
	parts := strings.SplitN(value, ":", 2)
	if len(parts) != 2 {
		return errors.New("Invalid preloaded RootFS value: not of the form 'stack-name:path'")
	}

	if parts[0] == "" {
		return errors.New("Invalid preloaded RootFS value: blank stack")
	}

	if parts[1] == "" {
		return errors.New("Invalid preloaded RootFS value: blank path")
	}

	(*s)[parts[0]] = parts[1]
	return nil
}

type providers []string

func (p *providers) String() string {
	return fmt.Sprintf("%v", *p)
}

func (p *providers) Set(value string) error {
	if value == "" {
		return errors.New("Cannot set blank value for RootFS provider")
	}

	*p = append(*p, value)
	return nil
}

const (
	dropsondeDestination = "localhost:3457"
	dropsondeOrigin      = "rep"

	bbsPingTimeout = 5 * time.Minute
)

func main() {
	cf_debug_server.AddFlags(flag.CommandLine)
	cf_lager.AddFlags(flag.CommandLine)

	stackMap := stackPathMap{}
	supportedProviders := providers{}
	flag.Var(&stackMap, "preloadedRootFS", "List of preloaded RootFSes")
	flag.Var(&supportedProviders, "rootFSProvider", "List of RootFS providers")
	flag.Parse()

	cf_http.Initialize(*communicationTimeout)

	clock := clock.NewClock()
	logger, reconfigurableSink := cf_lager.New(*sessionName)

	executorConfiguration := executorConfig()
	if !executorinit.ValidateExecutor(logger, executorConfiguration) {
		os.Exit(1)
	}

	initializeDropsonde(logger)

	if *cellID == "" {
		log.Fatalf("-cellID must be specified")
	}

	executorClient, executorMembers, err := executorinit.Initialize(logger, executorConfiguration, clock)
	if err != nil {
		log.Fatalf("Failed to initialize executor: %s", err.Error())
	}
	defer executorClient.Cleanup()

	if err := validateBBSAddress(); err != nil {
		logger.Fatal("invalid-bbs-address", err)
	}

	locketClient := initializeLocketClient(logger)

	evacuatable, evacuationReporter, evacuationNotifier := evacuation_context.New()

	// only one outstanding operation per container is necessary
	queue := operationq.NewSlidingQueue(1)

	evacuator := evacuation.NewEvacuator(
		logger,
		clock,
		executorClient,
		evacuationNotifier,
		*cellID,
		*evacuationTimeout,
		*evacuationPollingInterval,
	)

	bbsClient := initializeBBSClient(logger)
	httpServer, address := initializeServer(bbsClient, executorClient, evacuatable, evacuationReporter, logger, rep.StackPathMap(stackMap), supportedProviders)
	opGenerator := generator.New(*cellID, bbsClient, executorClient, evacuationReporter, uint64(evacuationTimeout.Seconds()))

	preloadedRootFSes := []string{}
	for k := range stackMap {
		preloadedRootFSes = append(preloadedRootFSes, k)
	}

	members := grouper.Members{
		{"http_server", httpServer},
		{"presence", initializeCellPresence(address, locketClient, executorClient, logger, supportedProviders, preloadedRootFSes)},
		{"bulker", harmonizer.NewBulker(logger, *pollingInterval, *evacuationPollingInterval, evacuationNotifier, clock, opGenerator, queue)},
		{"event-consumer", harmonizer.NewEventConsumer(logger, opGenerator, queue)},
		{"evacuator", evacuator},
	}

	members = append(executorMembers, members...)

	if dbgAddr := cf_debug_server.DebugAddress(flag.CommandLine); dbgAddr != "" {
		members = append(grouper.Members{
			{"debug-server", cf_debug_server.Runner(dbgAddr, reconfigurableSink)},
		}, members...)
	}

	group := grouper.NewOrdered(os.Interrupt, members)

	monitor := ifrit.Invoke(sigmon.New(group))

	logger.Info("started", lager.Data{"cell-id": *cellID})

	err = <-monitor.Wait()
	if err != nil {
		logger.Error("exited-with-failure", err)
		os.Exit(1)
	}

	logger.Info("exited")
}

func initializeDropsonde(logger lager.Logger) {
	err := dropsonde.Initialize(dropsondeDestination, dropsondeOrigin)
	if err != nil {
		logger.Error("failed to initialize dropsonde: %v", err)
	}
}

func initializeCellPresence(address string, locketClient locket.Client, executorClient executor.Client, logger lager.Logger, rootFSProviders, preloadedRootFSes []string) ifrit.Runner {
	config := maintain.Config{
		CellID:            *cellID,
		RepAddress:        address,
		Zone:              *zone,
		RetryInterval:     *lockRetryInterval,
		RootFSProviders:   rootFSProviders,
		PreloadedRootFSes: preloadedRootFSes,
	}
	return maintain.New(config, executorClient, locketClient, logger, clock.NewClock())
}

func initializeLocketClient(logger lager.Logger) locket.Client {
	client, err := consuladapter.NewClient(*consulCluster)
	if err != nil {
		logger.Fatal("new-client-failed", err)
	}

	sessionMgr := consuladapter.NewSessionManager(client)
	consulSession, err := consuladapter.NewSessionNoChecks(*sessionName, *lockTTL, client, sessionMgr)
	if err != nil {
		logger.Fatal("consul-session-failed", err)
	}

	return locket.NewClient(consulSession, clock.NewClock(), logger)
}

func initializeLRPStopper(guid string, executorClient executor.Client, logger lager.Logger) lrp_stopper.LRPStopper {
	return lrp_stopper.New(guid, executorClient, logger)
}

func initializeServer(
	bbsClient bbs.Client,
	executorClient executor.Client,
	evacuatable evacuation_context.Evacuatable,
	evacuationReporter evacuation_context.EvacuationReporter,
	logger lager.Logger,
	stackMap rep.StackPathMap,
	supportedProviders []string,
) (ifrit.Runner, string) {
	lrpStopper := initializeLRPStopper(*cellID, executorClient, logger)

	auctionCellRep := auction_cell_rep.New(*cellID, stackMap, supportedProviders, *zone, generateGuid, executorClient, evacuationReporter, logger)
	handlers := handlers.New(auctionCellRep, lrpStopper, executorClient, evacuatable, logger)

	router, err := rata.NewRouter(rep.Routes, handlers)
	if err != nil {
		logger.Fatal("failed-to-construct-router", err)
	}

	ip, err := localip.LocalIP()
	if err != nil {
		logger.Fatal("failed-to-fetch-ip", err)
	}

	port := strings.Split(*listenAddr, ":")[1]
	address := fmt.Sprintf("http://%s:%s", ip, port)

	return http_server.New(*listenAddr, router), address
}

func validateBBSAddress() error {
	if *bbsAddress == "" {
		return errors.New("bbsAddress is required")
	}
	return nil
}

func generateGuid() (string, error) {
	guid, err := uuid.NewV4()
	if err != nil {
		return "", err
	}
	return guid.String(), nil
}

func initializeBBSClient(logger lager.Logger) bbs.Client {
	bbsURL, err := url.Parse(*bbsAddress)
	if err != nil {
		logger.Fatal("Invalid BBS URL", err)
	}

	if bbsURL.Scheme != "https" {
		return bbs.NewClient(*bbsAddress)
	}

	bbsClient, err := bbs.NewSecureClient(*bbsAddress, *bbsCACert, *bbsClientCert, *bbsClientKey, *bbsClientSessionCacheSize, *bbsMaxIdleConnsPerHost)
	if err != nil {
		logger.Fatal("Failed to configure secure BBS client", err)
	}
	return bbsClient
}
<|endoftext|>


<commit_before>package cmd

import (
	"log"
	"io/ioutil"
	"os"
	"fmt"
	"archive/zip"
	"github.com/gosuri/uilive"
	"time"
	"strings"
	"github.com/fatih/color"
	"path/filepath"
	"gopkg.in/yaml.v2"
)

var (
	allResFiles  map[string]bool
	patchFilesMap map[string]bool
	distFileMap map[string]bool
	addedFilesMap map[string]bool
)

func Validate(patchLocation, distributionLocation string, logsEnabled bool) {

	if (!logsEnabled) {
		log.SetOutput(ioutil.Discard)
	} else {
		log.Println("Logs enabled")
	}
	log.Println("validate command called")

	initialize()

	log.Println("Patch Loc: " + patchLocation)
	if !isAZipFile(patchLocation) {
		color.Set(color.FgRed)
		fmt.Println("Patch file should be a zip file.")
		color.Unset()
		os.Exit(1)
	}
	patchLocationExists := checkFile(patchLocation)
	if patchLocationExists {
		log.Println("Patch location exists.")
	} else {
		color.Set(color.FgRed)
		fmt.Println("Patch location does not exist. Enter a valid file location.")
		color.Unset()
		os.Exit(1)
	}

	log.Println("Reading patch zip...")
	readPatchZip(patchLocation, logsEnabled)
	log.Println("Patch zip successfully read.")
	log.Println("Entries in patch zip: ", patchFilesMap)

	log.Println("Distribution Loc: " + distributionLocation)
	if isAZipFile(distributionLocation) {
		zipFileExists := checkFile(distributionLocation)
		if zipFileExists {
			log.Println("Distribution location exists.")
			readDistZip(distributionLocation, logsEnabled)
		} else {
			color.Set(color.FgRed)
			fmt.Println("Distribution zip does not exist. Enter a valid location.")
			color.Unset()
			os.Exit(1)
		}
	} else {
		distributionLocationExists := checkDir(distributionLocation)
		if distributionLocationExists {
			log.Println("Distribution location exists.")
			readDistDir(distributionLocation, logsEnabled)
		} else {
			color.Set(color.FgRed)
			fmt.Println("Distribution location does not exist. Enter a valid location.")
			color.Unset()
			os.Exit(1)
		}
	}
	validate()
}

func initialize() {
	allResFiles = make(map[string]bool)
	allResFiles[_LICENSE_FILE_NAME] = true
	allResFiles[_NOT_A_CONTRIBUTION_FILE_NAME] = true
	allResFiles[_README_FILE_NAME] = true
	allResFiles[_UPDATE_DESCRIPTOR_FILE_NAME] = true
	allResFiles[_INSTRUCTIONS_FILE_NAME] = true

	patchFilesMap = make(map[string]bool)
	distFileMap = make(map[string]bool)
	addedFilesMap = make(map[string]bool)
}

func validate() {
	for patchLoc, _ := range patchFilesMap {
		log.Println("Checking location: ", patchLoc)
		_, found := distFileMap[patchLoc]
		if found {
			log.Println(patchLoc, "found in distFileMap")
		} else {
			log.Println(patchLoc, "not found in distFileMap")
			_, found := addedFilesMap[patchLoc]
			if found {
				log.Println(patchLoc, "found in addedFilesMap")
			} else {
				log.Println(patchLoc, "not found in addedFilesMap")
				log.Println("addedFilesMap: ", addedFilesMap)
				color.Set(color.FgRed)
				fmt.Println(patchLoc, "not found in distribution and it is not a newly added file.")
				fmt.Println("If it is a new file, please add an entry in", _UPDATE_DESCRIPTOR_FILE_NAME,
					"file")
				fmt.Println("\nValidation FAILED\n")
				color.Unset()
				os.Exit(1)
			}
		}
	}
	fmt.Println("\nValidation SUCCESSFUL\n")
}

//This function reads the files of the given patch zip
func readPatchZip(zipLocation string, logsEnabled bool) (bool, error) {
	log.Println("Zip file reading started: ", zipLocation)

	patchName := strings.TrimSuffix(zipLocation, ".zip")
	if lastIndex := strings.LastIndex(patchName, string(os.PathSeparator)); lastIndex > -1 {
		patchName = patchName[lastIndex + 1:]
	}

	log.Println("Patch name: ", patchName)

	if !strings.HasPrefix(patchName, _UPDATE_NAME_PREFIX) {
		color.Set(color.FgRed)
		fmt.Println("Patch file does not have", _UPDATE_NAME_PREFIX, "prefix")
		color.Unset()
		os.Exit(1)
	} else {
		log.Println("Patch file does have", _UPDATE_NAME_PREFIX, "prefix")
	}

	// Create a reader out of the zip archive
	zipReader, err := zip.OpenReader(zipLocation)

	if err != nil {
		color.Set(color.FgRed)
		fmt.Println("Error occurred while reading zip:", err)
		color.Unset()
		log.Fatal(err)
	}
	defer zipReader.Close()

	totalFiles := len(zipReader.Reader.File)
	log.Println("File count in zip: ", totalFiles)

	filesRead := 0

	writer := uilive.New()
	//start listening for updates and render
	writer.Start()

	// Iterate through each file/dir found in
	for _, file := range zipReader.Reader.File {
		filesRead++
		if (!logsEnabled) {
			fmt.Fprintf(writer, "Reading files from patch zip: (%d/%d)\n", filesRead, totalFiles)
			time.Sleep(time.Millisecond * 5)
		}

		log.Println("Checking file: ", file.Name)

		index := strings.Index(file.Name, string(os.PathSeparator))

		if index == -1 {
			color.Set(color.FgRed)
			fmt.Println("Patch zip file should have a root folder called", patchName)
			color.Unset()
			os.Exit(1)
		} else {
			rootFolder := file.Name[:index]
			log.Println("RootFolder:", rootFolder)
			if rootFolder != patchName {
				color.Set(color.FgRed)
				fmt.Println(file.Name, "should be in", patchName, "root folder. But it is in ",
					rootFolder, "folder")
				color.Unset()
				os.Exit(1)
			}
		}

		_, found := allResFiles[file.FileInfo().Name()]

		if !found {
			containsCarbonHome := strings.Contains(file.Name, _CARBON_HOME)
			if (!containsCarbonHome) {
				color.Set(color.FgRed)
				fmt.Println("Does not have a '" + _CARBON_HOME + "' folder in the root folder '" +
				patchName + "'")
				color.Unset()
				os.Exit(1)
			}
			log.Println("Have a", _CARBON_HOME, "folder")

			temp := strings.TrimPrefix(file.Name, patchName + string(os.PathSeparator) + _CARBON_HOME)
			log.Println("Entry: ", temp)
			patchFilesMap[temp] = true
		} else {
			log.Println(file.FileInfo().Name(), "was found in resource map")
			delete(allResFiles, file.FileInfo().Name())
			log.Println(file.FileInfo().Name(), "was removed from the map")

			if file.FileInfo().Name() == _UPDATE_DESCRIPTOR_FILE_NAME {
				yamlFile, err := file.Open()
				if err != nil {
					color.Set(color.FgRed)
					fmt.Println("Error occurred while reading the", _UPDATE_DESCRIPTOR_FILE_NAME, "file:",
						err)
					color.Unset()
				}
				data, err := ioutil.ReadAll(yamlFile)
				if err != nil {
					log.Fatal(err)
				}
				descriptor := update_descriptor{}

				err = yaml.Unmarshal(data, &descriptor)
				if err != nil {
					color.Set(color.FgRed)
					fmt.Println("Error occurred while unmarshalling the yaml:", err)
					color.Unset()
				}

				log.Println("descriptor:", descriptor)

				for _, addedFile := range descriptor.File_changes.Added_files {

					addedFilesMap[addedFile] = true
					//_, found := allResFiles[addedFile]
					//if found {
					//	log.Println(file, "found in patch")
					//
					//} else {
					//	fmt.Println(file, "not found in patch. But it is added as an entry in" +
					//	"'added_file' section in", _DESCRIPTOR_YAML_NAME)
					//	os.Exit(1)
					//}
				}
			}
		}
	}

	//Delete instructions.txt file if it is left in the map
	delete(allResFiles,_INSTRUCTIONS_FILE_NAME)

	log.Println("Resource map:", allResFiles)
	log.Println(patchFilesMap)
	if (len(allResFiles) != 0) {
		writer.Stop()
		color.Set(color.FgRed)
		fmt.Println("Following resource file(s) were not found in the patch: ")
		for key, _ := range allResFiles {
			fmt.Println("\t", "-", key)
		}
		color.Unset()
		os.Exit(1)
	}

	writer.Stop()
	log.Println("Zip file reading finished")
	log.Println("Total files read: ", filesRead)
	if totalFiles == filesRead {
		log.Println("All files read")
		return true, nil
	} else {
		color.Set(color.FgRed)
		fmt.Println("All files not read from zip file")
		color.Unset()
		os.Exit(1)
		return false, nil
	}
}

//This function reads the files of the given distribution zip
func readDistZip(zipLocation string, logsEnabled bool) (bool, error) {
	log.Println("Zip file reading started: ", zipLocation)

	distName := strings.TrimSuffix(zipLocation, ".zip")
	if lastIndex := strings.LastIndex(distName, string(os.PathSeparator)); lastIndex > -1 {
		distName = distName[lastIndex + 1:]
	}

	// Create a reader out of the zip archive
	zipReader, err := zip.OpenReader(zipLocation)

	if err != nil {
		color.Set(color.FgRed)
		fmt.Println("Error occurred while reading zip:", err)
		color.Unset()
		log.Fatal(err)
	}
	defer zipReader.Close()

	totalFiles := len(zipReader.Reader.File)
	log.Println("File count in zip: ", totalFiles)

	filesRead := 0

	writer := uilive.New()
	//start listening for updates and render
	writer.Start()

	// Iterate through each file/dir found in
	for _, file := range zipReader.Reader.File {
		filesRead++
		if (!logsEnabled) {
			fmt.Fprintf(writer, "Reading files from distribution zip: (%d/%d)\n", filesRead, totalFiles)
			time.Sleep(time.Millisecond * 2)
		}

		log.Println("Checking file: ", file.Name)

		temp := strings.TrimPrefix(file.Name, distName)
		log.Println("Entry: ", temp)
		distFileMap[temp] = true

	}

	writer.Stop()
	log.Println("Zip file reading finished")
	log.Println("Total files read: ", filesRead)
	if totalFiles == filesRead {
		log.Println("All files read")
		return true, nil
	} else {
		color.Set(color.FgRed)
		fmt.Println("All files not read from zip file")
		color.Unset()
		os.Exit(1)
		return false, nil
	}
}

func readDistDir(distributionLocation string, logsEnabled bool) {

	writer := uilive.New()
	//start listening for updates and render
	writer.Start()
	filesRead := 0
	err := filepath.Walk(distributionLocation, func(path string, fileInfo os.FileInfo, err error) error {
		filesRead++;
		if (!logsEnabled) {
			fmt.Fprintf(writer, "Reading files from distribution directory: %d files read\n", filesRead)
			time.Sleep(time.Millisecond * 2)
		}

		log.Println("Walking: ", path)

		//if !fileInfo.IsDir() {
		// Create and write files to the archive, which in turn
		// are getting written to the underlying writer to the
		// .zip file we created at the beginning

		//header, err := zip.FileInfoHeader(fileInfo)
		//if err != nil {
		//	color.Set(color.FgRed)
		//	fmt.Println("Error occurred while creating the zip file: ", err)
		//	return err
		//	color.Unset()
		//}
		//
		//header.Name = _PATCH_NAME + string(os.PathSeparator) + strings.TrimPrefix(path, _TEMP_DIR_NAME + string(os.PathSeparator))

		temp := strings.TrimPrefix(path, distributionLocation)
		log.Println("Entry: ", temp)
		distFileMap[temp] = true

		//}
		//log.Printf("Visited: %s\n", path)
		return nil
	})

	if err != nil {
		color.Set(color.FgRed)
		fmt.Println("Error occurred while reading the zip file: ", err)
		os.Exit(1)
		color.Unset()
	}
	log.Println("Total files read:", filesRead)
	writer.Stop()
}
<commit_msg>Comments added<commit_after>package cmd

import (
	"log"
	"io/ioutil"
	"os"
	"fmt"
	"archive/zip"
	"github.com/gosuri/uilive"
	"time"
	"strings"
	"github.com/fatih/color"
	"path/filepath"
	"gopkg.in/yaml.v2"
)

var (
	allResFiles  map[string]bool
	patchFilesMap map[string]bool
	distFileMap map[string]bool
	addedFilesMap map[string]bool
)

//Entry point of the validate command
func Validate(patchLocation, distributionLocation string, logsEnabled bool) {

	if (!logsEnabled) {
		log.SetOutput(ioutil.Discard)
	} else {
		log.Println("Logs enabled")
	}
	log.Println("validate command called")

	//Initialize variables
	initialize()

	//patch location should be a zip file
	log.Println("Patch Loc: " + patchLocation)
	if !isAZipFile(patchLocation) {
		color.Set(color.FgRed)
		fmt.Println("Patch file should be a zip file.")
		color.Unset()
		os.Exit(1)
	}
	//Check whether the patch location exists
	patchLocationExists := checkFile(patchLocation)
	if patchLocationExists {
		log.Println("Patch location exists.")
	} else {
		color.Set(color.FgRed)
		fmt.Println("Patch location does not exist. Enter a valid file location.")
		color.Unset()
		os.Exit(1)
	}

	log.Println("Reading patch zip...")
	readPatchZip(patchLocation, logsEnabled)
	log.Println("Patch zip successfully read.")
	log.Println("Entries in patch zip: ", patchFilesMap)

	log.Println("Distribution Loc: " + distributionLocation)
	//Check whether the distribution is a zip or a directory
	if isAZipFile(distributionLocation) {
		//Check whether the distribution zip exists
		zipFileExists := checkFile(distributionLocation)
		if zipFileExists {
			log.Println("Distribution location exists.")
			readDistZip(distributionLocation, logsEnabled)
		} else {
			color.Set(color.FgRed)
			fmt.Println("Distribution zip does not exist. Enter a valid location.")
			color.Unset()
			os.Exit(1)
		}
	} else {
		//Check whether the distribution location exists
		distributionLocationExists := checkDir(distributionLocation)
		if distributionLocationExists {
			log.Println("Distribution location exists.")
			readDistDir(distributionLocation, logsEnabled)
		} else {
			color.Set(color.FgRed)
			fmt.Println("Distribution location does not exist. Enter a valid location.")
			color.Unset()
			os.Exit(1)
		}
	}
	//Validate files
	validate()
}

//This initializes the variables
func initialize() {
	allResFiles = make(map[string]bool)
	allResFiles[_LICENSE_FILE_NAME] = true
	allResFiles[_NOT_A_CONTRIBUTION_FILE_NAME] = true
	allResFiles[_README_FILE_NAME] = true
	allResFiles[_UPDATE_DESCRIPTOR_FILE_NAME] = true
	allResFiles[_INSTRUCTIONS_FILE_NAME] = true

	patchFilesMap = make(map[string]bool)
	distFileMap = make(map[string]bool)
	addedFilesMap = make(map[string]bool)
}

//This method validates the files
func validate() {
	//Iterate through all the files in the patch. All files should be in the distribution unless they are newly
	// added files
	for patchLoc, _ := range patchFilesMap {
		log.Println("Checking location: ", patchLoc)
		//Check whether the distribution has a file with the same name
		_, found := distFileMap[patchLoc]
		//If there is a file
		if found {
			log.Println(patchLoc, "found in distFileMap")
		} else {
			//If there is no file
			log.Println(patchLoc, "not found in distFileMap")
			//Check whether it is a newly added file
			_, found := addedFilesMap[patchLoc]
			//if it is a newly added file
			if found {
				log.Println(patchLoc, "found in addedFilesMap")
			} else {
				//If it is not a newly added file, print an error
				log.Println(patchLoc, "not found in addedFilesMap")
				log.Println("addedFilesMap: ", addedFilesMap)
				color.Set(color.FgRed)
				fmt.Println(patchLoc, "not found in distribution and it is not a newly added file.")
				fmt.Println("If it is a new file, please add an entry in", _UPDATE_DESCRIPTOR_FILE_NAME,
					"file")
				fmt.Println("\nValidation FAILED\n")
				color.Unset()
				os.Exit(1)
			}
		}
	}
	fmt.Println("\nValidation SUCCESSFUL\n")
}

//This function reads the files of the given patch zip
func readPatchZip(zipLocation string, logsEnabled bool) {
	log.Println("Zip file reading started: ", zipLocation)

	patchName := strings.TrimSuffix(zipLocation, ".zip")
	if lastIndex := strings.LastIndex(patchName, string(os.PathSeparator)); lastIndex > -1 {
		patchName = patchName[lastIndex + 1:]
	}
	log.Println("Patch name: ", patchName)

	//Check whether the patch name has the required prefix
	if !strings.HasPrefix(patchName, _UPDATE_NAME_PREFIX) {
		color.Set(color.FgRed)
		fmt.Println("Patch file does not have", _UPDATE_NAME_PREFIX, "prefix")
		color.Unset()
		os.Exit(1)
	} else {
		log.Println("Patch file does have", _UPDATE_NAME_PREFIX, "prefix")
	}

	// Create a reader out of the zip archive
	zipReader, err := zip.OpenReader(zipLocation)
	if err != nil {
		color.Set(color.FgRed)
		fmt.Println("Error occurred while reading zip:", err)
		color.Unset()
		log.Fatal(err)
	}
	defer zipReader.Close()

	totalFiles := len(zipReader.Reader.File)
	log.Println("File count in zip: ", totalFiles)

	fileCount := 0
	//Create a new writer to show the progress
	writer := uilive.New()
	//start listening for updates and render
	writer.Start()

	// Iterate through each file/dir found in the zip
	for _, file := range zipReader.Reader.File {
		fileCount++
		if (!logsEnabled) {
			fmt.Fprintf(writer, "Reading files from patch zip: (%d/%d)\n", fileCount, totalFiles)
			time.Sleep(time.Millisecond * 5)
		}

		log.Println("Checking file: ", file.Name)

		//Every file should be in a root folder. Check for the os.PathSeparator character to identify this
		index := strings.Index(file.Name, string(os.PathSeparator))
		if index == -1 {
			color.Set(color.FgRed)
			fmt.Println("Patch zip file should have a root folder called", patchName)
			color.Unset()
			os.Exit(1)
		} else {
			rootFolder := file.Name[:index]
			log.Println("RootFolder:", rootFolder)
			if rootFolder != patchName {
				color.Set(color.FgRed)
				fmt.Println(file.Name, "should be in", patchName, "root folder. But it is in ",
					rootFolder, "folder")
				color.Unset()
				os.Exit(1)
			}
		}
		//Check whether the file is a resource file
		_, found := allResFiles[file.FileInfo().Name()]
		//If it is not a resource file
		if !found {
			//It should be in a carbon.home folder
			containsCarbonHome := strings.Contains(file.Name, _CARBON_HOME)
			if (!containsCarbonHome) {
				color.Set(color.FgRed)
				fmt.Println("Does not have a '" + _CARBON_HOME + "' folder in the root folder '" +
				patchName + "'")
				color.Unset()
				os.Exit(1)
			}
			log.Println("Have a", _CARBON_HOME, "folder")

			temp := strings.TrimPrefix(file.Name, patchName + string(os.PathSeparator) + _CARBON_HOME)
			log.Println("Entry: ", temp)
			patchFilesMap[temp] = true
		} else {
			//If the file is a resource file, delete the entry from allResFiles. This map is later used
			// to track missing resource files
			log.Println(file.FileInfo().Name(), "was found in resource map")
			delete(allResFiles, file.FileInfo().Name())
			log.Println(file.FileInfo().Name(), "was removed from the map")
			//If the file is update-descriptor.yaml file, we need to read the newly added files.
			// Otherwise there will be no match for these files and validation will be failed
			if file.FileInfo().Name() == _UPDATE_DESCRIPTOR_FILE_NAME {
				//Open the file
				yamlFile, err := file.Open()
				if err != nil {
					color.Set(color.FgRed)
					fmt.Println("Error occurred while reading the", _UPDATE_DESCRIPTOR_FILE_NAME, "file:",
						err)
					color.Unset()
				}
				//Get the byte array
				data, err := ioutil.ReadAll(yamlFile)
				if err != nil {
					log.Fatal(err)
				}
				descriptor := update_descriptor{}
				//Get the values
				err = yaml.Unmarshal(data, &descriptor)
				if err != nil {
					color.Set(color.FgRed)
					fmt.Println("Error occurred while unmarshalling the yaml:", err)
					color.Unset()
				}
				log.Println("descriptor:", descriptor)
				//Add all files to addedFilesMap
				for _, addedFile := range descriptor.File_changes.Added_files {
					addedFilesMap[addedFile] = true
				}
			}
		}
	}
	//Stop the writer
	writer.Stop()

	//Delete instructions.txt file if it is left in the map because it is optional
	delete(allResFiles, _INSTRUCTIONS_FILE_NAME)
	log.Println("Resource map:", allResFiles)
	log.Println(patchFilesMap)

	//At this point, the size of the allResFiles should be zero. If one or more files are not found, that means
	// that some required files are missing
	if (len(allResFiles) != 0) {
		//Print the missing files
		color.Set(color.FgRed)
		fmt.Println("Following resource file(s) were not found in the patch: ")
		for key, _ := range allResFiles {
			fmt.Println("\t", "-", key)
		}
		color.Unset()
		os.Exit(1)
	}

	//Check whether all files are read
	log.Println("Zip file reading finished")
	log.Println("Total files read: ", fileCount)
	if totalFiles == fileCount {
		log.Println("All files read")
	} else {
		color.Set(color.FgRed)
		fmt.Println("All files not read from zip file")
		color.Unset()
		os.Exit(1)
	}
}

//This function reads the files of the given distribution zip
func readDistZip(zipLocation string, logsEnabled bool) (bool, error) {
	log.Println("Zip file reading started: ", zipLocation)

	//Get the distribution name
	distName := strings.TrimSuffix(zipLocation, ".zip")
	if lastIndex := strings.LastIndex(distName, string(os.PathSeparator)); lastIndex > -1 {
		distName = distName[lastIndex + 1:]
	}

	// Create a reader out of the zip archive
	zipReader, err := zip.OpenReader(zipLocation)
	if err != nil {
		color.Set(color.FgRed)
		fmt.Println("Error occurred while reading zip:", err)
		color.Unset()
		log.Fatal(err)
	}
	defer zipReader.Close()

	totalFiles := len(zipReader.Reader.File)
	log.Println("File count in zip: ", totalFiles)

	fileCount := 0
	//Create a writer to show the progress
	writer := uilive.New()
	//start listening for updates and render
	writer.Start()

	// Iterate through each file/dir found in
	for _, file := range zipReader.Reader.File {
		fileCount++
		if (!logsEnabled) {
			fmt.Fprintf(writer, "Reading files from distribution zip: (%d/%d)\n", fileCount, totalFiles)
			time.Sleep(time.Millisecond * 2)
		}

		log.Println("Checking file: ", file.Name)

		//Get the relative path in the zip
		temp := strings.TrimPrefix(file.Name, distName)
		log.Println("Entry: ", temp)
		//Add to the map
		distFileMap[temp] = true

	}
	//Stop the writer
	writer.Stop()

	//Check whether all files are read
	log.Println("Zip file reading finished")
	log.Println("Total files read: ", fileCount)
	if totalFiles == fileCount {
		log.Println("All files read")
		return true, nil
	} else {
		color.Set(color.FgRed)
		fmt.Println("All files not read from zip file")
		color.Unset()
		os.Exit(1)
		return false, nil
	}
}

func readDistDir(distributionLocation string, logsEnabled bool) {
	//Create a writer to show the progress
	writer := uilive.New()
	//start listening for updates and render
	writer.Start()
	fileCount := 0
	//Start the walk
	err := filepath.Walk(distributionLocation, func(path string, fileInfo os.FileInfo, err error) error {
		fileCount++;
		if (!logsEnabled) {
			fmt.Fprintf(writer, "Reading files from distribution directory: %d files read\n", fileCount)
			time.Sleep(time.Millisecond * 2)
		}
		log.Println("Walking: ", path)
		//Get the relative path
		temp := strings.TrimPrefix(path, distributionLocation)
		log.Println("Entry: ", temp)
		//Add to the map
		distFileMap[temp] = true
		return nil
	})
	if err != nil {
		color.Set(color.FgRed)
		fmt.Println("Error occurred while reading the zip file: ", err)
		os.Exit(1)
		color.Unset()
	}
	log.Println("Total files read:", fileCount)
	//stop the writer
	writer.Stop()
}
<|endoftext|>


<commit_before>// Copyright 2018-2019 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// cpu - connection to CPU server over SSH protocol
//
// Synopsis:
//     cpu [OPTIONS]
//
// Description:
//     On local machine useful flags are all save -remote
//     On remote machines, all save dbg9p, key, hostkey.
//
//     CPU is an ssh client that starts up a shell on a remote machine,
//     as usual; but, further, makes a namespace of the local machine
//     available in a private mount rooted at /tmp/cpu.
//     Wherever you go, there your files are.
//     You can, in the ssh session, do something like this:
//     chroot /tmp/cpu /bin/bash
//     and at that point, you are running a bash, imported from your local
//     machine, on the remote machine; it will use your .profile and
//     all your files are available. You can also do something like
//     cat /tmp/cpu/etc/hosts
//     if your host file is lacking; or
//     cp /etc/hosts /tmp/cpu/tmp
//     to get the /etc/hosts on the remote machine to your local machine.
//
//     The cpu client makes this work by starting a cpu command on the
//     remote machine with a -remote switch and several other arguments.
//     The local cpu starts a 9p server and, using ssh port forwarding,
//     makes that server available to the remote. On the remote side, the
//     cpu command establishes a private, unshared mount of tmpfs on /tmp;
//     creates /tmp/cpu; and mounts the 9p server on that directory.
//
//     CPU has many options, as shown above; most you need not worry about.
//     The most common invocation is
//     cpu -h hostname
//     which will start a shell and mount the 9p server in /tmp/cpu.
//     Note this mount proceeds over the ssh session, and further
//     it mounts in a private /tmp; there is little to see when
//     it is running from outside the ssh session
//
// Options:
//     -bin string
//           path of cpu binary
//     -d    enable debug prints
//     -dbg9p
//           show 9p io
//     -hostkey string
//           host key file
//     -key string
//           key file (default "$HOME/.ssh/cpu_rsa")
//     -network string
//           network to use (default "tcp")
//     -p string
//           port to use (default "22")
//     -port9p string
//           port9p # on remote machine for 9p mount
//     -remote
//           Indicates we are the remote side of the cpu session
//     -srv string
//           what server to run (default "unpfs")
// Examples
// In these examples, cpu runs with warning messages enabled.
// The first message is a warning that cpu could not use overlayfs to build a
// a reasonable union mount. The next are showing you what it is mounting, the
// union mount having failed.
// These mounts are the best we could do for a reasonable compromise of
// wanting local resources visible (e.g. /dev) and using resources from the
// remote machine (e.g. /etc, /lib, /usr and so on).
// u-root doesn't really need /lib and /usr, and u-root's /etc is minimal by design,
// so this works.
// Also note that the user's 9p server running on the local machine is mounted at /tmp/cpu.
// We can turn these off at some point but for now, in early days, we may want them.
// Note that these messages come from the remote side.
// cpu to a machine with bash as your shell and run a command
//   cpu -sp 23 date
//     2019/05/17 16:53:22 Overlayfs mount failed: invalid argument. Proceeding with selective mounts from /tmp/cpu into /
//     2019/05/17 16:53:22 Mounted /tmp/cpu/lib on /lib
//     2019/05/17 16:53:22 Mounted /tmp/cpu/lib64 on /lib64
//     2019/05/17 16:53:22 Warning: mounting /tmp/cpu/lib32 on /lib32 failed: no such file or directory
//     2019/05/17 16:53:22 Mounted /tmp/cpu/usr on /usr
//     2019/05/17 16:53:22 Mounted /tmp/cpu/bin on /bin
//     2019/05/17 16:53:22 Mounted /tmp/cpu/etc on /etc
//     Fri May 17 16:53:23 UTC 2019
// cpu to a machine and run $SHELL (since no arguments were given)
// NOTE: $SHELL is NOT installed on the remote machine! It (and all its .so's and . files)
// come from the local machine.
// cpu sp -23
//    2019/05/17 16:58:04 Overlayfs mount failed: invalid argument. Proceeding with selective mounts from /tmp/cpu into /
//    2019/05/17 16:58:04 Mounted /tmp/cpu/lib on /lib
//    2019/05/17 16:58:04 Mounted /tmp/cpu/lib64 on /lib64
//    2019/05/17 16:58:04 Warning: mounting /tmp/cpu/lib32 on /lib32 failed: no such file or directory
//    2019/05/17 16:58:04 Mounted /tmp/cpu/usr on /usr
//    2019/05/17 16:58:04 Mounted /tmp/cpu/bin on /bin
//    2019/05/17 16:58:05 Mounted /tmp/cpu/etc on /etc
//    root@(none):/# echo ~
//    /tmp/cpu/home/rminnich
//    root@(none):/# ls ~
//    IDAPROPASSWORD  go      ida-7.2  projects            salishan2019random~
//    bin             gopath  papers   salishan2019random  snap
//    root@(none):/#
//    # Now that we are on the node, modprobe something
//    root@(none):/# depmod
//    depmod: ERROR: could not open directory /lib/modules/5.0.0-rc3+: No such file or directory
//    depmod: FATAL: could not search modules: No such file or directory
//    root@(none):/#
//    # Note that, if we had the right modules on our LOCAL machine for this remote machine, we could
//    # insert them. This further means you can build a modular kernel in FLASH and insmod needed modules
//    # later (as long as your core kernel has networking, that is!). Modules could include, e.g., an AHCI
//    # driver.
//    # run the lspci command but redirect output to ~
//    # note it is not installed on the remote machine; it comes from our local machine.
//    root@(none):/# lspci
//    00:00.0 Host bridge: Intel Corporation 82G33/G31/P35/P31 Express DRAM Controller
//    00:01.0 VGA compatible controller: Device 1234:1111 (rev 02)
//    00:02.0 Unclassified device [00ff]: Red Hat, Inc. Virtio RNG
//    00:03.0 Ethernet controller: Intel Corporation 82540EM Gigabit Ethernet Controller (rev 03)
//    00:1f.0 ISA bridge: Intel Corporation 82801IB (ICH9) LPC Interface Controller (rev 02)
//    00:1f.2 SATA controller: Intel Corporation 82801IR/IO/IH (ICH9R/DO/DH) 6 port SATA Controller [AHCI mode] (rev 02)
//    00:1f.3 SMBus: Intel Corporation 82801I (ICH9 Family) SMBus Controller (rev 02)
//    root@(none):/# lspci > ~/xyz
//    root@(none):/# exit
//    # exit and notice that file is on my local machine now:
//    exit
//    rminnich@xcpu:~/gopath/src/github.com/u-root/u-root/xcmds/cpu$ ls -l ~/xyz
//    -rw-r--r-- 1 rminnich rminnich 577 May 17 17:06 /home/rminnich/xyz
//    rminnich@xcpu:~/gopath/src/github.com/u-root/u-root/xcmds/cpu$
package main
<commit_msg>Tweak documentation<commit_after>// Copyright 2018-2019 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// cpu - connection to CPU server over SSH protocol
//
// Synopsis:
//     cpu [OPTIONS]
//
// Description:
//     On local machine useful flags are all save -remote
//     On remote machines, all save dbg9p, key, hostkey.
//
//     CPU is an ssh client that starts up a shell on a remote machine,
//     as usual; but, further, makes a namespace of the local machine
//     available in a private mount rooted at /tmp/cpu.
//     Wherever you go, there your files are.
//     You can, in the ssh session, do something like this:
//     chroot /tmp/cpu /bin/bash
//     and at that point, you are running a bash, imported from your local
//     machine, on the remote machine; it will use your .profile and
//     all your files are available. You can also do something like
//     cat /tmp/cpu/etc/hosts
//     if your host file is lacking; or
//     cp /etc/hosts /tmp/cpu/tmp
//     to get the /etc/hosts on the remote machine to your local machine.
//
//     The cpu client makes this work by starting a cpu command on the
//     remote machine with a -remote switch and several other arguments.
//     The local cpu starts a 9p server and, using ssh port forwarding,
//     makes that server available to the remote. On the remote side, the
//     cpu command establishes a private, unshared mount of tmpfs on /tmp;
//     creates /tmp/cpu; and mounts the 9p server on that directory.
//
//     CPU has many options, as shown above; most you need not worry about.
//     The most common invocation is
//     cpu -h hostname
//     which will start a shell and mount the 9p server in /tmp/cpu.
//     Note this mount proceeds over the ssh session, and further
//     it mounts in a private /tmp; there is little to see when
//     it is running from outside the ssh session
//
// Options:
//     -bin string
//           path of cpu binary
//     -d    enable debug prints
//     -dbg9p
//           show 9p io
//     -hostkey string
//           host key file
//     -key string
//           key file (default "$HOME/.ssh/cpu_rsa")
//     -network string
//           network to use (default "tcp")
//     -p string
//           port to use (default "22")
//     -port9p string
//           port9p # on remote machine for 9p mount
//     -remote
//           Indicates we are the remote side of the cpu session
//     -srv string
//           what server to run (default none; use internal)
// Examples
// In these examples, cpu runs with warning messages enabled.
// The first message is a warning that cpu could not use overlayfs to build a
// a reasonable union mount. The next are showing you what it is mounting, the
// union mount having failed.
// These mounts are the best we could do for a reasonable compromise of
// wanting local resources visible (e.g. /dev) and using resources from the
// remote machine (e.g. /etc, /lib, /usr and so on).
// u-root doesn't really need /lib and /usr, and u-root's /etc is minimal by design,
// so this works.
// Also note that the user's 9p server running on the local machine is mounted at /tmp/cpu.
// We can turn these off at some point but for now, in early days, we may want them.
// Note that these messages come from the remote side.
// cpu to a machine with bash as your shell and run a command
//   cpu -sp 23 date
//     2019/05/17 16:53:22 Overlayfs mount failed: invalid argument. Proceeding with selective mounts from /tmp/cpu into /
//     2019/05/17 16:53:22 Mounted /tmp/cpu/lib on /lib
//     2019/05/17 16:53:22 Mounted /tmp/cpu/lib64 on /lib64
//     2019/05/17 16:53:22 Warning: mounting /tmp/cpu/lib32 on /lib32 failed: no such file or directory
//     2019/05/17 16:53:22 Mounted /tmp/cpu/usr on /usr
//     2019/05/17 16:53:22 Mounted /tmp/cpu/bin on /bin
//     2019/05/17 16:53:22 Mounted /tmp/cpu/etc on /etc
//     Fri May 17 16:53:23 UTC 2019
// cpu to a machine and run $SHELL (since no arguments were given)
// NOTE: $SHELL is NOT installed on the remote machine! It (and all its .so's and . files)
// come from the local machine.
// cpu sp -23
//    2019/05/17 16:58:04 Overlayfs mount failed: invalid argument. Proceeding with selective mounts from /tmp/cpu into /
//    2019/05/17 16:58:04 Mounted /tmp/cpu/lib on /lib
//    2019/05/17 16:58:04 Mounted /tmp/cpu/lib64 on /lib64
//    2019/05/17 16:58:04 Warning: mounting /tmp/cpu/lib32 on /lib32 failed: no such file or directory
//    2019/05/17 16:58:04 Mounted /tmp/cpu/usr on /usr
//    2019/05/17 16:58:04 Mounted /tmp/cpu/bin on /bin
//    2019/05/17 16:58:05 Mounted /tmp/cpu/etc on /etc
//    root@(none):/# echo ~
//    /tmp/cpu/home/rminnich
//    root@(none):/# ls ~
//    IDAPROPASSWORD  go      ida-7.2  projects            salishan2019random~
//    bin             gopath  papers   salishan2019random  snap
//    root@(none):/#
//    # Now that we are on the node, modprobe something
//    root@(none):/# depmod
//    depmod: ERROR: could not open directory /lib/modules/5.0.0-rc3+: No such file or directory
//    depmod: FATAL: could not search modules: No such file or directory
//    root@(none):/#
//    # Note that, if we had the right modules on our LOCAL machine for this remote machine, we could
//    # insert them. This further means you can build a modular kernel in FLASH and insmod needed modules
//    # later (as long as your core kernel has networking, that is!). Modules could include, e.g., an AHCI
//    # driver.
//    # run the lspci command but redirect output to ~
//    # note it is not installed on the remote machine; it comes from our local machine.
//    root@(none):/# lspci
//    00:00.0 Host bridge: Intel Corporation 82G33/G31/P35/P31 Express DRAM Controller
//    00:01.0 VGA compatible controller: Device 1234:1111 (rev 02)
//    00:02.0 Unclassified device [00ff]: Red Hat, Inc. Virtio RNG
//    00:03.0 Ethernet controller: Intel Corporation 82540EM Gigabit Ethernet Controller (rev 03)
//    00:1f.0 ISA bridge: Intel Corporation 82801IB (ICH9) LPC Interface Controller (rev 02)
//    00:1f.2 SATA controller: Intel Corporation 82801IR/IO/IH (ICH9R/DO/DH) 6 port SATA Controller [AHCI mode] (rev 02)
//    00:1f.3 SMBus: Intel Corporation 82801I (ICH9 Family) SMBus Controller (rev 02)
//    root@(none):/# lspci > ~/xyz
//    root@(none):/# exit
//    # exit and notice that file is on my local machine now:
//    exit
//    rminnich@xcpu:~/gopath/src/github.com/u-root/u-root/xcmds/cpu$ ls -l ~/xyz
//    -rw-r--r-- 1 rminnich rminnich 577 May 17 17:06 /home/rminnich/xyz
//    rminnich@xcpu:~/gopath/src/github.com/u-root/u-root/xcmds/cpu$
package main
<|endoftext|>


<commit_before>// Copyright 2017 The LUCI 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 notify

import (
	"bytes"
	"compress/gzip"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"

	"github.com/golang/protobuf/proto"

	"go.chromium.org/gae/service/datastore"
	"go.chromium.org/gae/service/info"
	"go.chromium.org/gae/service/mail"
	"go.chromium.org/luci/appengine/tq"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/common/data/stringset"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	gitpb "go.chromium.org/luci/common/proto/git"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/server/auth"

	"go.chromium.org/luci/luci_notify/api/config"
	notifypb "go.chromium.org/luci/luci_notify/api/config"
	"go.chromium.org/luci/luci_notify/internal"
	"go.chromium.org/luci/luci_notify/mailtmpl"
)

// createEmailTasks constructs EmailTasks to be dispatched onto the task queue.
func createEmailTasks(c context.Context, recipients []EmailNotify, input *config.TemplateInput) ([]*tq.Task, error) {
	// Get templates.
	bundle, err := getBundle(c, input.Build.Builder.Project)
	if err != nil {
		return nil, errors.Annotate(err, "failed to get a bundle of email templates").Err()
	}

	// Generate emails.
	// An EmailTask with subject and body per template name.
	// They will be used as templates for actual tasks.
	taskTemplates := map[string]*internal.EmailTask{}
	for _, r := range recipients {
		name := r.Template
		if name == "" {
			name = mailtmpl.DefaultTemplateName
		}

		if _, ok := taskTemplates[name]; ok {
			continue
		}

		subject, body := bundle.GenerateEmail(name, input)

		// Note: this buffer should not be reused.
		buf := &bytes.Buffer{}
		gz := gzip.NewWriter(buf)
		io.WriteString(gz, body)
		if err := gz.Close(); err != nil {
			panic("failed to gzip HTML body in memory")
		}
		taskTemplates[name] = &internal.EmailTask{
			Subject:  subject,
			BodyGzip: buf.Bytes(),
		}
	}

	// Create a task per recipient.
	// Do not bundle multiple recipients into one task because we don't use BCC.
	tasks := make([]*tq.Task, 0, len(recipients))
	seen := stringset.New(len(recipients))
	for _, r := range recipients {
		name := r.Template
		if name == "" {
			name = mailtmpl.DefaultTemplateName
		}

		emailKey := fmt.Sprintf("%d-%s-%s", input.Build.Id, name, r.Email)
		if seen.Has(emailKey) {
			continue
		}
		seen.Add(emailKey)

		task := *taskTemplates[name] // copy
		task.Recipients = []string{r.Email}
		tasks = append(tasks, &tq.Task{
			DeduplicationKey: emailKey,
			Payload:          &task,
		})
	}
	return tasks, nil
}

// isRecipientAllowed returns true if the given recipient is allowed to be notified about the given build.
func isRecipientAllowed(c context.Context, recipient string, build *buildbucketpb.Build) bool {
	// TODO(mknyszek): Do a real ACL check here.
	if strings.HasSuffix(recipient, "@google.com") || strings.HasSuffix(recipient, "@chromium.org") {
		return true
	}
	logging.Warningf(c, "Address %q is not allowed to be notified of build %d", recipient, build.Id)
	return false
}

// BlamelistRepoWhiteset computes the aggregate repository whitelist for all
// blamelist notification configurations in a given set of notifications.
func BlamelistRepoWhiteset(notifications notifypb.Notifications) stringset.Set {
	whiteset := stringset.New(0)
	for _, notification := range notifications.GetNotifications() {
		blamelistInfo := notification.GetNotifyBlamelist()
		for _, repo := range blamelistInfo.GetRepositoryWhitelist() {
			whiteset.Add(repo)
		}
	}
	return whiteset
}

// ComputeRecipients computes the set of recipients given a set of
// notifications, and potentially "input" and "output" blamelists.
//
// An "input" blamelist is computed from the input commit to a build, while an
// "output" blamelist is derived from output commits.
func ComputeRecipients(c context.Context, notifications notifypb.Notifications, inputBlame []*gitpb.Commit, outputBlame Logs) []EmailNotify {
	return computeRecipientsInternal(c, notifications, inputBlame, outputBlame,
		func(c context.Context, url string) ([]byte, error) {
			transport, err := auth.GetRPCTransport(c, auth.AsSelf)
			if err != nil {
				return nil, err
			}

			req, err := http.NewRequestWithContext(c, "GET", url, nil)
			if err != nil {
				return nil, err
			}

			response, err := (&http.Client{Transport: transport}).Do(req)
			if err != nil {
				return nil, errors.Annotate(err, "failed to get data from %q", url).Err()
			}

			defer response.Body.Close()
			bytes, err := ioutil.ReadAll(response.Body)
			if err != nil {
				return nil, errors.Annotate(err, "failed to read response body from %q", url).Err()
			}

			return bytes, nil
		})
}

// computeRecipientsInternal also takes fetchFunc, so http requests can be
// mocked out for testing.
func computeRecipientsInternal(c context.Context, notifications notifypb.Notifications, inputBlame []*gitpb.Commit, outputBlame Logs, fetchFunc func(context.Context, string) ([]byte, error)) []EmailNotify {
	recipients := make([]EmailNotify, 0)
	for _, notification := range notifications.GetNotifications() {
		// Aggregate the static list of recipients from the Notifications.
		for _, recipient := range notification.GetEmail().GetRecipients() {
			recipients = append(recipients, EmailNotify{
				Email:    recipient,
				Template: notification.Template,
			})
		}

		// Don't bother dealing with anything blamelist related if there's no config for it.
		if notification.NotifyBlamelist == nil {
			continue
		}

		// If the whitelist is empty, use the static blamelist.
		whitelist := notification.NotifyBlamelist.GetRepositoryWhitelist()
		if len(whitelist) == 0 {
			recipients = append(recipients, commitsBlamelist(inputBlame, notification.Template)...)
			continue
		}

		// If the whitelist is non-empty, use the dynamic blamelist.
		whiteset := stringset.NewFromSlice(whitelist...)
		recipients = append(recipients, outputBlame.Filter(whiteset).Blamelist(notification.Template)...)
	}

	// Acquired before appending to "recipients" from the tasks below.
	mRecipients := sync.Mutex{}
	err := parallel.WorkPool(8, func(ch chan<- func() error) {
		for _, notification := range notifications.GetNotifications() {
			template := notification.Template
			for _, rotation := range notification.GetEmail().GetRotaNgRotations() {
				rotation := rotation
				ch <- func() error {
					return fetchRotaOncallers(c, rotation, template, fetchFunc, &recipients, &mRecipients)
				}
			}
		}
	})

	if err != nil {
		// Just log the error and continue. Nothing much else we can do, and it's possible that we only failed
		// to fetch some of the recipients, so we can at least return the ones we were able to compute.
		logging.Errorf(c, "failed to fetch some or all Rota-NG oncallers: %s", err)
	}

	return recipients
}

func fetchRotaOncallers(c context.Context, rotation, template string, fetchFunc func(context.Context, string) ([]byte, error), recipients *[]EmailNotify, mRecipients *sync.Mutex) error {
	address := "https://rota-ng.appspot.com/legacy/" + url.PathEscape(rotation) + ".json"
	resp, err := fetchFunc(c, address)
	if err != nil {
		err = errors.Annotate(err, "failed to fetch address %q", address).Err()
		return err
	}

	var oncallEmails struct {
		Emails []string
	}
	if err = json.Unmarshal(resp, &oncallEmails); err != nil {
		return errors.Annotate(err, "failed to unmarshal JSON").Err()
	}

	mRecipients.Lock()
	defer mRecipients.Unlock()
	for _, email := range oncallEmails.Emails {
		*recipients = append(*recipients, EmailNotify{
			Email:    email,
			Template: template,
		})
	}

	return nil
}

// Notify discovers, consolidates and filters recipients from a Builder's notifications,
// and 'email_notify' properties, then dispatches notifications if necessary.
// Does not dispatch a notification for same email, template and build more than
// once. Ignores current transaction in c, if any.
func Notify(c context.Context, d *tq.Dispatcher, recipients []EmailNotify, templateParams *config.TemplateInput) error {
	c = datastore.WithoutTransaction(c)

	// Remove unallowed recipients.
	allRecipients := recipients
	recipients = recipients[:0]
	for _, r := range allRecipients {
		if isRecipientAllowed(c, r.Email, templateParams.Build) {
			recipients = append(recipients, r)
		}
	}

	if len(recipients) == 0 {
		logging.Infof(c, "Nobody to notify...")
		return nil
	}
	tasks, err := createEmailTasks(c, recipients, templateParams)
	if err != nil {
		return errors.Annotate(err, "failed to create email tasks").Err()
	}
	return d.AddTask(c, tasks...)
}

// InitDispatcher registers the send email task with the given dispatcher.
func InitDispatcher(d *tq.Dispatcher) {
	d.RegisterTask(&internal.EmailTask{}, SendEmail, "email", nil)
}

// SendEmail is a push queue handler that attempts to send an email.
func SendEmail(c context.Context, task proto.Message) error {
	appID := info.AppID(c)
	sender := fmt.Sprintf("%s <noreply@%s.appspotmail.com>", appID, appID)

	// TODO(mknyszek): Query Milo for additional build information.
	emailTask := task.(*internal.EmailTask)

	body := emailTask.Body
	if len(emailTask.BodyGzip) > 0 {
		r, err := gzip.NewReader(bytes.NewReader(emailTask.BodyGzip))
		if err != nil {
			return err
		}
		buf, err := ioutil.ReadAll(r)
		if err != nil {
			return err
		}
		body = string(buf)
	}

	return mail.Send(c, &mail.Message{
		Sender:   sender,
		To:       emailTask.Recipients,
		Subject:  emailTask.Subject,
		HTMLBody: body,
		ReplyTo:  emailTask.Recipients[0],
	})
}
<commit_msg>Don't use http.NewRequestWithContext.<commit_after>// Copyright 2017 The LUCI 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 notify

import (
	"bytes"
	"compress/gzip"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"

	"github.com/golang/protobuf/proto"

	"go.chromium.org/gae/service/datastore"
	"go.chromium.org/gae/service/info"
	"go.chromium.org/gae/service/mail"
	"go.chromium.org/luci/appengine/tq"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/common/data/stringset"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	gitpb "go.chromium.org/luci/common/proto/git"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/server/auth"

	"go.chromium.org/luci/luci_notify/api/config"
	notifypb "go.chromium.org/luci/luci_notify/api/config"
	"go.chromium.org/luci/luci_notify/internal"
	"go.chromium.org/luci/luci_notify/mailtmpl"
)

// createEmailTasks constructs EmailTasks to be dispatched onto the task queue.
func createEmailTasks(c context.Context, recipients []EmailNotify, input *config.TemplateInput) ([]*tq.Task, error) {
	// Get templates.
	bundle, err := getBundle(c, input.Build.Builder.Project)
	if err != nil {
		return nil, errors.Annotate(err, "failed to get a bundle of email templates").Err()
	}

	// Generate emails.
	// An EmailTask with subject and body per template name.
	// They will be used as templates for actual tasks.
	taskTemplates := map[string]*internal.EmailTask{}
	for _, r := range recipients {
		name := r.Template
		if name == "" {
			name = mailtmpl.DefaultTemplateName
		}

		if _, ok := taskTemplates[name]; ok {
			continue
		}

		subject, body := bundle.GenerateEmail(name, input)

		// Note: this buffer should not be reused.
		buf := &bytes.Buffer{}
		gz := gzip.NewWriter(buf)
		io.WriteString(gz, body)
		if err := gz.Close(); err != nil {
			panic("failed to gzip HTML body in memory")
		}
		taskTemplates[name] = &internal.EmailTask{
			Subject:  subject,
			BodyGzip: buf.Bytes(),
		}
	}

	// Create a task per recipient.
	// Do not bundle multiple recipients into one task because we don't use BCC.
	tasks := make([]*tq.Task, 0, len(recipients))
	seen := stringset.New(len(recipients))
	for _, r := range recipients {
		name := r.Template
		if name == "" {
			name = mailtmpl.DefaultTemplateName
		}

		emailKey := fmt.Sprintf("%d-%s-%s", input.Build.Id, name, r.Email)
		if seen.Has(emailKey) {
			continue
		}
		seen.Add(emailKey)

		task := *taskTemplates[name] // copy
		task.Recipients = []string{r.Email}
		tasks = append(tasks, &tq.Task{
			DeduplicationKey: emailKey,
			Payload:          &task,
		})
	}
	return tasks, nil
}

// isRecipientAllowed returns true if the given recipient is allowed to be notified about the given build.
func isRecipientAllowed(c context.Context, recipient string, build *buildbucketpb.Build) bool {
	// TODO(mknyszek): Do a real ACL check here.
	if strings.HasSuffix(recipient, "@google.com") || strings.HasSuffix(recipient, "@chromium.org") {
		return true
	}
	logging.Warningf(c, "Address %q is not allowed to be notified of build %d", recipient, build.Id)
	return false
}

// BlamelistRepoWhiteset computes the aggregate repository whitelist for all
// blamelist notification configurations in a given set of notifications.
func BlamelistRepoWhiteset(notifications notifypb.Notifications) stringset.Set {
	whiteset := stringset.New(0)
	for _, notification := range notifications.GetNotifications() {
		blamelistInfo := notification.GetNotifyBlamelist()
		for _, repo := range blamelistInfo.GetRepositoryWhitelist() {
			whiteset.Add(repo)
		}
	}
	return whiteset
}

// ComputeRecipients computes the set of recipients given a set of
// notifications, and potentially "input" and "output" blamelists.
//
// An "input" blamelist is computed from the input commit to a build, while an
// "output" blamelist is derived from output commits.
func ComputeRecipients(c context.Context, notifications notifypb.Notifications, inputBlame []*gitpb.Commit, outputBlame Logs) []EmailNotify {
	return computeRecipientsInternal(c, notifications, inputBlame, outputBlame,
		func(c context.Context, url string) ([]byte, error) {
			transport, err := auth.GetRPCTransport(c, auth.AsSelf)
			if err != nil {
				return nil, err
			}

			req, err := http.NewRequest("GET", url, nil)
			if err != nil {
				return nil, err
			}
			req = req.WithContext(c)

			response, err := (&http.Client{Transport: transport}).Do(req)
			if err != nil {
				return nil, errors.Annotate(err, "failed to get data from %q", url).Err()
			}

			defer response.Body.Close()
			bytes, err := ioutil.ReadAll(response.Body)
			if err != nil {
				return nil, errors.Annotate(err, "failed to read response body from %q", url).Err()
			}

			return bytes, nil
		})
}

// computeRecipientsInternal also takes fetchFunc, so http requests can be
// mocked out for testing.
func computeRecipientsInternal(c context.Context, notifications notifypb.Notifications, inputBlame []*gitpb.Commit, outputBlame Logs, fetchFunc func(context.Context, string) ([]byte, error)) []EmailNotify {
	recipients := make([]EmailNotify, 0)
	for _, notification := range notifications.GetNotifications() {
		// Aggregate the static list of recipients from the Notifications.
		for _, recipient := range notification.GetEmail().GetRecipients() {
			recipients = append(recipients, EmailNotify{
				Email:    recipient,
				Template: notification.Template,
			})
		}

		// Don't bother dealing with anything blamelist related if there's no config for it.
		if notification.NotifyBlamelist == nil {
			continue
		}

		// If the whitelist is empty, use the static blamelist.
		whitelist := notification.NotifyBlamelist.GetRepositoryWhitelist()
		if len(whitelist) == 0 {
			recipients = append(recipients, commitsBlamelist(inputBlame, notification.Template)...)
			continue
		}

		// If the whitelist is non-empty, use the dynamic blamelist.
		whiteset := stringset.NewFromSlice(whitelist...)
		recipients = append(recipients, outputBlame.Filter(whiteset).Blamelist(notification.Template)...)
	}

	// Acquired before appending to "recipients" from the tasks below.
	mRecipients := sync.Mutex{}
	err := parallel.WorkPool(8, func(ch chan<- func() error) {
		for _, notification := range notifications.GetNotifications() {
			template := notification.Template
			for _, rotation := range notification.GetEmail().GetRotaNgRotations() {
				rotation := rotation
				ch <- func() error {
					return fetchRotaOncallers(c, rotation, template, fetchFunc, &recipients, &mRecipients)
				}
			}
		}
	})

	if err != nil {
		// Just log the error and continue. Nothing much else we can do, and it's possible that we only failed
		// to fetch some of the recipients, so we can at least return the ones we were able to compute.
		logging.Errorf(c, "failed to fetch some or all Rota-NG oncallers: %s", err)
	}

	return recipients
}

func fetchRotaOncallers(c context.Context, rotation, template string, fetchFunc func(context.Context, string) ([]byte, error), recipients *[]EmailNotify, mRecipients *sync.Mutex) error {
	address := "https://rota-ng.appspot.com/legacy/" + url.PathEscape(rotation) + ".json"
	resp, err := fetchFunc(c, address)
	if err != nil {
		err = errors.Annotate(err, "failed to fetch address %q", address).Err()
		return err
	}

	var oncallEmails struct {
		Emails []string
	}
	if err = json.Unmarshal(resp, &oncallEmails); err != nil {
		return errors.Annotate(err, "failed to unmarshal JSON").Err()
	}

	mRecipients.Lock()
	defer mRecipients.Unlock()
	for _, email := range oncallEmails.Emails {
		*recipients = append(*recipients, EmailNotify{
			Email:    email,
			Template: template,
		})
	}

	return nil
}

// Notify discovers, consolidates and filters recipients from a Builder's notifications,
// and 'email_notify' properties, then dispatches notifications if necessary.
// Does not dispatch a notification for same email, template and build more than
// once. Ignores current transaction in c, if any.
func Notify(c context.Context, d *tq.Dispatcher, recipients []EmailNotify, templateParams *config.TemplateInput) error {
	c = datastore.WithoutTransaction(c)

	// Remove unallowed recipients.
	allRecipients := recipients
	recipients = recipients[:0]
	for _, r := range allRecipients {
		if isRecipientAllowed(c, r.Email, templateParams.Build) {
			recipients = append(recipients, r)
		}
	}

	if len(recipients) == 0 {
		logging.Infof(c, "Nobody to notify...")
		return nil
	}
	tasks, err := createEmailTasks(c, recipients, templateParams)
	if err != nil {
		return errors.Annotate(err, "failed to create email tasks").Err()
	}
	return d.AddTask(c, tasks...)
}

// InitDispatcher registers the send email task with the given dispatcher.
func InitDispatcher(d *tq.Dispatcher) {
	d.RegisterTask(&internal.EmailTask{}, SendEmail, "email", nil)
}

// SendEmail is a push queue handler that attempts to send an email.
func SendEmail(c context.Context, task proto.Message) error {
	appID := info.AppID(c)
	sender := fmt.Sprintf("%s <noreply@%s.appspotmail.com>", appID, appID)

	// TODO(mknyszek): Query Milo for additional build information.
	emailTask := task.(*internal.EmailTask)

	body := emailTask.Body
	if len(emailTask.BodyGzip) > 0 {
		r, err := gzip.NewReader(bytes.NewReader(emailTask.BodyGzip))
		if err != nil {
			return err
		}
		buf, err := ioutil.ReadAll(r)
		if err != nil {
			return err
		}
		body = string(buf)
	}

	return mail.Send(c, &mail.Message{
		Sender:   sender,
		To:       emailTask.Recipients,
		Subject:  emailTask.Subject,
		HTMLBody: body,
		ReplyTo:  emailTask.Recipients[0],
	})
}
<|endoftext|>


<commit_before>// Copyright 2017 The LUCI 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 notify

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

	"github.com/golang/protobuf/ptypes"
	"golang.org/x/net/context"
	"google.golang.org/genproto/protobuf/field_mask"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"go.chromium.org/gae/service/datastore"

	"go.chromium.org/luci/appengine/tq"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	bbv1 "go.chromium.org/luci/common/api/buildbucket/buildbucket/v1"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	gitpb "go.chromium.org/luci/common/proto/git"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/grpc/prpc"
	"go.chromium.org/luci/server/auth"
	"go.chromium.org/luci/server/router"

	notifyConfig "go.chromium.org/luci/luci_notify/config"
)

var (
	errBuilderDeleted = fmt.Errorf("builder deleted between datastore.Get calls")
)

func getBuilderID(b *Build) string {
	return "buildbucket/" + b.Builder.IDString()
}

func commitIndex(commits []*gitpb.Commit, revision string) int {
	for i, commit := range commits {
		if commit.Id == revision {
			return i
		}
	}
	return -1
}

// EmailNotify contains information for delivery and personalization of notification emails.
type EmailNotify struct {
	Email    string `json:"email"`
	Template string `json:"template"`
}

func extractEmailNotifyValues(parametersJSON string) ([]EmailNotify, error) {
	if parametersJSON == "" {
		return nil, nil
	}
	// json equivalent: {"email_notify": [{"email": "<address>"}, ...]}
	var output struct {
		EmailNotify []EmailNotify `json:"email_notify"`
	}

	if err := json.NewDecoder(strings.NewReader(parametersJSON)).Decode(&output); err != nil {
		return nil, errors.Annotate(err, "invalid msg.ParametersJson").Err()
	}
	return output.EmailNotify, nil
}

// handleBuild processes a Build and sends appropriate notifications.
//
// This function should serve as documentation of the process of going from
// a Build to sent notifications. It also should explicitly handle ACLs and
// stop the process of handling notifications early to avoid wasting compute time.
//
// history is a function that contacts gitiles to obtain the git history for
// revision ordering purposes. It's passed in as a parameter in order to mock it
// for testing.
func handleBuild(c context.Context, d *tq.Dispatcher, build *Build, history HistoryFunc) error {
	switch ex, err := datastore.Exists(c, &notifyConfig.Project{Name: build.Builder.Project}); {
	case err != nil:
		return err
	case !ex.All():
		return nil // This project is not tracked by luci-notify
	}

	builderID := getBuilderID(build)
	logging.Infof(c, "Finding config for %q, %s", builderID, build.Status)
	notifiers, err := notifyConfig.LookupNotifiers(c, build.Builder.Project, builderID)
	if err != nil {
		return errors.Annotate(err, "looking up notifiers").Err()
	}
	if len(notifiers) == 0 {
		// No configurations were found, but we might need to generate notifications
		// based on properties. Don't fall through to avoid storing build status for
		// builds without configurations.
		return Notify(c, d, notifiers, StatusUnknown, build)
	}

	if len(build.Input.GetGitilesCommits()) == 0 {
		logging.Infof(c, "No revision information found for this build, ignoring...")
		return nil
	}
	gCommit := build.Input.GitilesCommits[0]

	// Get the Builder for the first time, and initialize if there's nothing there.
	builder := &Builder{ID: builderID}
	switch err := datastore.Get(c, builder); {
	case err == datastore.ErrNoSuchEntity:
		// If the Builder isn't found, start a transaction and try to store
		// the builder for the first time.
		keepGoing := false
		err = datastore.RunInTransaction(c, func(c context.Context) error {
			switch err := datastore.Get(c, builder); {
			case err == datastore.ErrNoSuchEntity:
				// Notify, but avoid on_change notification by setting Status to StatusUnknown.
				builder.Status = StatusUnknown
				if err := Notify(c, d, notifiers, builder.Status, build); err != nil {
					return err
				}
				// Initialize the Builder in the datastore.
				if err := datastore.Put(c, NewBuilder(builderID, gCommit.Id, &build.Build)); err != nil {
					return err
				}
				return nil
			case err != nil:
				return err
			}
			// If we actually managed to get a Builder, we should continue with the
			// full code path.
			keepGoing = true
			return nil
		}, nil)
		if err != nil || !keepGoing {
			return err
		}
	case err != nil:
		return err
	}

	// commits contains a list of commits from builder.StatusRevision..gCommit.Revision (oldest..newest).
	commits, err := history(c, gCommit.Host, gCommit.Project, builder.StatusRevision, gCommit.Id)
	if err != nil {
		return err
	}
	if len(commits) == 0 {
		logging.Debugf(c, "Found build with old commit, ignoring...")
		// Notify about the build, but ignore on_change by creating a builder with an unknown status.
		return Notify(c, d, notifiers, StatusUnknown, build)
	}

	// Update `builder`, and check if we need to store a newer version, then store it.
	return datastore.RunInTransaction(c, func(c context.Context) error {
		switch err := datastore.Get(c, builder); {
		case err == datastore.ErrNoSuchEntity:
			return errBuilderDeleted
		case err != nil:
			return err
		}

		// If there's been no change in status, don't bother updating.
		if build.Status == builder.Status {
			return nil
		}

		index := commitIndex(commits, builder.StatusRevision)
		outOfOrder := false
		createTime, _ := ptypes.Timestamp(build.CreateTime)
		switch {
		// If the revision is not found, we can conclude that the Builder has
		// advanced beyond gCommit.Revision. This is because:
		//   1) builder.StatusRevision only ever moves forward.
		//   2) commits contains the git history up to gCommit.Revision.
		case index < 0:
			logging.Debugf(c, "Found build with old commit during transaction.")
			outOfOrder = true

		// If the revision is current, check build creation time.
		case index == 0 && builder.StatusBuildTime.After(createTime):
			logging.Debugf(c, "Found build with the same commit but an old time.")
			outOfOrder = true
		}

		// If the build is out-of-order, we want to ignore only on_change notifications.
		// Since on_change only occurs when Builder Status != StatusUnknown, set builder
		// to a new builder with exactly this property.
		if outOfOrder {
			return Notify(c, d, notifiers, StatusUnknown, build)
		}
		if err := Notify(c, d, notifiers, builder.Status, build); err != nil {
			return err
		}
		return datastore.Put(c, NewBuilder(builderID, gCommit.Id, &build.Build))
	}, nil)
}

func newBuildsClient(c context.Context, host string) (buildbucketpb.BuildsClient, error) {
	t, err := auth.GetRPCTransport(c, auth.AsSelf)
	if err != nil {
		return nil, err
	}
	return buildbucketpb.NewBuildsPRPCClient(&prpc.Client{
		C:    &http.Client{Transport: t},
		Host: host,
	}), nil
}

// BuildbucketPubSubHandler is the main entrypoint for a new update from buildbucket's pubsub.
//
// This handler delegates the actual processing of the build to handleBuild.
// Its primary purpose is to unwrap context boilerplate and deal with progress-stopping errors.
func BuildbucketPubSubHandler(ctx *router.Context, d *tq.Dispatcher) error {
	c := ctx.Context
	build, err := extractBuild(c, ctx.Request)
	switch {
	case err != nil:
		return errors.Annotate(err, "failed to extract build").Err()

	case build == nil:
		// Ignore.
		return nil

	default:
		if err := handleBuild(c, d, build, gitilesHistory); err != nil {
			return errors.Annotate(err, "failed to notify").Tag(transient.Tag).Err()
		}
		return nil
	}
}

// Build is buildbucketpb.Build along with the parsed 'email_notify' values.
type Build struct {
	buildbucketpb.Build
	EmailNotify []EmailNotify
}

// extractBuild constructs a Build from the PubSub HTTP request.
func extractBuild(c context.Context, r *http.Request) (*Build, error) {
	// sent by pubsub.
	// This struct is just convenient for unwrapping the json message
	var msg struct {
		Message struct {
			Data []byte
		}
		Attributes map[string]interface{}
	}
	if err := json.NewDecoder(r.Body).Decode(&msg); err != nil {
		return nil, errors.Annotate(err, "could not decode message").Err()
	}

	if v, ok := msg.Attributes["version"].(string); ok && v != "v1" {
		// Ignore v2 pubsub messages. TODO(nodir): use v2.
		return nil, nil
	}
	var message struct {
		Build    bbv1.ApiCommonBuildMessage
		Hostname string
	}
	switch err := json.Unmarshal(msg.Message.Data, &message); {
	case err != nil:
		return nil, errors.Annotate(err, "could not parse pubsub message data").Err()
	case !strings.HasPrefix(message.Build.Bucket, "luci."):
		logging.Infof(c, "Received build that isn't part of LUCI, ignoring...")
		return nil, nil
	case message.Build.Status != bbv1.StatusCompleted:
		logging.Infof(c, "Received build that hasn't completed yet, ignoring...")
		return nil, nil
	}

	buildsClient, err := newBuildsClient(c, message.Hostname)
	if err != nil {
		return nil, err
	}
	res, err := buildsClient.GetBuild(c, &buildbucketpb.GetBuildRequest{
		Id: message.Build.Id,
		Fields: &field_mask.FieldMask{
			Paths: []string{"steps"},
		},
	})
	switch {
	case status.Code(err) == codes.NotFound:
		logging.Warningf(c, "no access to build %d", message.Build.Id)
		return nil, nil
	case err != nil:
		return nil, errors.Annotate(err, "could not fetch buildbucket build %d", message.Build.Id).Err()
	}

	emails, err := extractEmailNotifyValues(message.Build.ParametersJson)
	if err != nil {
		return nil, errors.Annotate(err, "could not decode email_notify").Err()
	}

	return &Build{
		Build:       *res,
		EmailNotify: emails,
	}, nil
}
<commit_msg>[luci-notify] Fetch all build fields.<commit_after>// Copyright 2017 The LUCI 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 notify

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

	"github.com/golang/protobuf/ptypes"
	"golang.org/x/net/context"
	"google.golang.org/genproto/protobuf/field_mask"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"go.chromium.org/gae/service/datastore"

	"go.chromium.org/luci/appengine/tq"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	bbv1 "go.chromium.org/luci/common/api/buildbucket/buildbucket/v1"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	gitpb "go.chromium.org/luci/common/proto/git"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/grpc/prpc"
	"go.chromium.org/luci/server/auth"
	"go.chromium.org/luci/server/router"

	notifyConfig "go.chromium.org/luci/luci_notify/config"
)

var (
	errBuilderDeleted = fmt.Errorf("builder deleted between datastore.Get calls")

	// buildFieldMask defines which buildbucketpb.Build fields to fetch and
	// make available to email templates.
	// If a message field is specified here without periods, e.g. "steps", all
	// of its subfields are included too.
	buildFieldMask = &field_mask.FieldMask{
		Paths: []string{
			"id",
			"builder",
			"number",
			"created_by",
			"view_url",
			"create_time",
			"start_time",
			"end_time",
			"update_time",
			"status",
			"input",
			"output",
			"steps",
			"infra",
			"tags",
		},
	}
)

func getBuilderID(b *Build) string {
	return "buildbucket/" + b.Builder.IDString()
}

func commitIndex(commits []*gitpb.Commit, revision string) int {
	for i, commit := range commits {
		if commit.Id == revision {
			return i
		}
	}
	return -1
}

// EmailNotify contains information for delivery and personalization of notification emails.
type EmailNotify struct {
	Email    string `json:"email"`
	Template string `json:"template"`
}

func extractEmailNotifyValues(parametersJSON string) ([]EmailNotify, error) {
	if parametersJSON == "" {
		return nil, nil
	}
	// json equivalent: {"email_notify": [{"email": "<address>"}, ...]}
	var output struct {
		EmailNotify []EmailNotify `json:"email_notify"`
	}

	if err := json.NewDecoder(strings.NewReader(parametersJSON)).Decode(&output); err != nil {
		return nil, errors.Annotate(err, "invalid msg.ParametersJson").Err()
	}
	return output.EmailNotify, nil
}

// handleBuild processes a Build and sends appropriate notifications.
//
// This function should serve as documentation of the process of going from
// a Build to sent notifications. It also should explicitly handle ACLs and
// stop the process of handling notifications early to avoid wasting compute time.
//
// history is a function that contacts gitiles to obtain the git history for
// revision ordering purposes. It's passed in as a parameter in order to mock it
// for testing.
func handleBuild(c context.Context, d *tq.Dispatcher, build *Build, history HistoryFunc) error {
	switch ex, err := datastore.Exists(c, &notifyConfig.Project{Name: build.Builder.Project}); {
	case err != nil:
		return err
	case !ex.All():
		return nil // This project is not tracked by luci-notify
	}

	builderID := getBuilderID(build)
	logging.Infof(c, "Finding config for %q, %s", builderID, build.Status)
	notifiers, err := notifyConfig.LookupNotifiers(c, build.Builder.Project, builderID)
	if err != nil {
		return errors.Annotate(err, "looking up notifiers").Err()
	}
	if len(notifiers) == 0 {
		// No configurations were found, but we might need to generate notifications
		// based on properties. Don't fall through to avoid storing build status for
		// builds without configurations.
		return Notify(c, d, notifiers, StatusUnknown, build)
	}

	if len(build.Input.GetGitilesCommits()) == 0 {
		logging.Infof(c, "No revision information found for this build, ignoring...")
		return nil
	}
	gCommit := build.Input.GitilesCommits[0]

	// Get the Builder for the first time, and initialize if there's nothing there.
	builder := &Builder{ID: builderID}
	switch err := datastore.Get(c, builder); {
	case err == datastore.ErrNoSuchEntity:
		// If the Builder isn't found, start a transaction and try to store
		// the builder for the first time.
		keepGoing := false
		err = datastore.RunInTransaction(c, func(c context.Context) error {
			switch err := datastore.Get(c, builder); {
			case err == datastore.ErrNoSuchEntity:
				// Notify, but avoid on_change notification by setting Status to StatusUnknown.
				builder.Status = StatusUnknown
				if err := Notify(c, d, notifiers, builder.Status, build); err != nil {
					return err
				}
				// Initialize the Builder in the datastore.
				if err := datastore.Put(c, NewBuilder(builderID, gCommit.Id, &build.Build)); err != nil {
					return err
				}
				return nil
			case err != nil:
				return err
			}
			// If we actually managed to get a Builder, we should continue with the
			// full code path.
			keepGoing = true
			return nil
		}, nil)
		if err != nil || !keepGoing {
			return err
		}
	case err != nil:
		return err
	}

	// commits contains a list of commits from builder.StatusRevision..gCommit.Revision (oldest..newest).
	commits, err := history(c, gCommit.Host, gCommit.Project, builder.StatusRevision, gCommit.Id)
	if err != nil {
		return err
	}
	if len(commits) == 0 {
		logging.Debugf(c, "Found build with old commit, ignoring...")
		// Notify about the build, but ignore on_change by creating a builder with an unknown status.
		return Notify(c, d, notifiers, StatusUnknown, build)
	}

	// Update `builder`, and check if we need to store a newer version, then store it.
	return datastore.RunInTransaction(c, func(c context.Context) error {
		switch err := datastore.Get(c, builder); {
		case err == datastore.ErrNoSuchEntity:
			return errBuilderDeleted
		case err != nil:
			return err
		}

		// If there's been no change in status, don't bother updating.
		if build.Status == builder.Status {
			return nil
		}

		index := commitIndex(commits, builder.StatusRevision)
		outOfOrder := false
		createTime, _ := ptypes.Timestamp(build.CreateTime)
		switch {
		// If the revision is not found, we can conclude that the Builder has
		// advanced beyond gCommit.Revision. This is because:
		//   1) builder.StatusRevision only ever moves forward.
		//   2) commits contains the git history up to gCommit.Revision.
		case index < 0:
			logging.Debugf(c, "Found build with old commit during transaction.")
			outOfOrder = true

		// If the revision is current, check build creation time.
		case index == 0 && builder.StatusBuildTime.After(createTime):
			logging.Debugf(c, "Found build with the same commit but an old time.")
			outOfOrder = true
		}

		// If the build is out-of-order, we want to ignore only on_change notifications.
		// Since on_change only occurs when Builder Status != StatusUnknown, set builder
		// to a new builder with exactly this property.
		if outOfOrder {
			return Notify(c, d, notifiers, StatusUnknown, build)
		}
		if err := Notify(c, d, notifiers, builder.Status, build); err != nil {
			return err
		}
		return datastore.Put(c, NewBuilder(builderID, gCommit.Id, &build.Build))
	}, nil)
}

func newBuildsClient(c context.Context, host string) (buildbucketpb.BuildsClient, error) {
	t, err := auth.GetRPCTransport(c, auth.AsSelf)
	if err != nil {
		return nil, err
	}
	return buildbucketpb.NewBuildsPRPCClient(&prpc.Client{
		C:    &http.Client{Transport: t},
		Host: host,
	}), nil
}

// BuildbucketPubSubHandler is the main entrypoint for a new update from buildbucket's pubsub.
//
// This handler delegates the actual processing of the build to handleBuild.
// Its primary purpose is to unwrap context boilerplate and deal with progress-stopping errors.
func BuildbucketPubSubHandler(ctx *router.Context, d *tq.Dispatcher) error {
	c := ctx.Context
	build, err := extractBuild(c, ctx.Request)
	switch {
	case err != nil:
		return errors.Annotate(err, "failed to extract build").Err()

	case build == nil:
		// Ignore.
		return nil

	default:
		if err := handleBuild(c, d, build, gitilesHistory); err != nil {
			return errors.Annotate(err, "failed to notify").Tag(transient.Tag).Err()
		}
		return nil
	}
}

// Build is buildbucketpb.Build along with the parsed 'email_notify' values.
type Build struct {
	buildbucketpb.Build
	EmailNotify []EmailNotify
}

// extractBuild constructs a Build from the PubSub HTTP request.
func extractBuild(c context.Context, r *http.Request) (*Build, error) {
	// sent by pubsub.
	// This struct is just convenient for unwrapping the json message
	var msg struct {
		Message struct {
			Data []byte
		}
		Attributes map[string]interface{}
	}
	if err := json.NewDecoder(r.Body).Decode(&msg); err != nil {
		return nil, errors.Annotate(err, "could not decode message").Err()
	}

	if v, ok := msg.Attributes["version"].(string); ok && v != "v1" {
		// Ignore v2 pubsub messages. TODO(nodir): use v2.
		return nil, nil
	}
	var message struct {
		Build    bbv1.ApiCommonBuildMessage
		Hostname string
	}
	switch err := json.Unmarshal(msg.Message.Data, &message); {
	case err != nil:
		return nil, errors.Annotate(err, "could not parse pubsub message data").Err()
	case !strings.HasPrefix(message.Build.Bucket, "luci."):
		logging.Infof(c, "Received build that isn't part of LUCI, ignoring...")
		return nil, nil
	case message.Build.Status != bbv1.StatusCompleted:
		logging.Infof(c, "Received build that hasn't completed yet, ignoring...")
		return nil, nil
	}

	buildsClient, err := newBuildsClient(c, message.Hostname)
	if err != nil {
		return nil, err
	}
	res, err := buildsClient.GetBuild(c, &buildbucketpb.GetBuildRequest{
		Id:     message.Build.Id,
		Fields: buildFieldMask,
	})
	switch {
	case status.Code(err) == codes.NotFound:
		logging.Warningf(c, "no access to build %d", message.Build.Id)
		return nil, nil
	case err != nil:
		return nil, errors.Annotate(err, "could not fetch buildbucket build %d", message.Build.Id).Err()
	}

	emails, err := extractEmailNotifyValues(message.Build.ParametersJson)
	if err != nil {
		return nil, errors.Annotate(err, "could not decode email_notify").Err()
	}

	return &Build{
		Build:       *res,
		EmailNotify: emails,
	}, nil
}
<|endoftext|>