<commit_before>package agent

import (
	"bytes"
	"encoding/json"
	"github.com/boltdb/bolt"
	resourced_config "github.com/resourced/resourced/config"
	"github.com/resourced/resourced/libprocess"
	"github.com/resourced/resourced/libstring"
	"github.com/resourced/resourced/libtime"
	resourced_readers "github.com/resourced/resourced/readers"
	resourced_writers "github.com/resourced/resourced/writers"
	"os"
	"reflect"
	"strings"
	"time"
)

// NewAgent is the constructor fot Agent struct.
func NewAgent() (*Agent, error) {
	agent := &Agent{}

	agent.setTags()

	err := agent.setConfigStorage()
	if err != nil {
		return nil, err
	}

	err = agent.setDb()
	if err != nil {
		return nil, err
	}

	return agent, err
}

// Agent struct carries most of the functionality of ResourceD.
// It collects information through readers and serve them up as HTTP+JSON.
type Agent struct {
	ConfigStorage *resourced_config.ConfigStorage
	DbPath        string
	Db            *bolt.DB
	Tags          []string
}

func (a *Agent) setTags() {
	a.Tags = make([]string, 0)

	tags := os.Getenv("RESOURCED_TAGS")
	if tags != "" {
		tagsSlice := strings.Split(tags, ",")
		a.Tags = make([]string, len(tagsSlice))

		for i, tag := range tagsSlice {
			a.Tags[i] = strings.TrimSpace(tag)
		}
	}
}

// setDb configures the local storage.
func (a *Agent) setDb() error {
	var err error

	dbPath := os.Getenv("RESOURCED_DB")
	if dbPath == "" {
		dbPath = "~/resourced/db"

		err = os.MkdirAll(libstring.ExpandTilde("~/resourced"), 0755)
		if err != nil {
			return err
		}
	}

	a.DbPath = libstring.ExpandTilde(dbPath)

	a.Db, err = bolt.Open(a.DbPath, 0644, nil)
	if err != nil {
		return err
	}

	// Create "resources" bucket
	a.Db.Update(func(tx *bolt.Tx) error {
		tx.CreateBucket([]byte("resources"))
		return nil
	})

	return err
}

// dbBucket returns the boltdb bucket.
func (a *Agent) dbBucket(tx *bolt.Tx) *bolt.Bucket {
	return tx.Bucket([]byte("resources"))
}

// pathWithPrefix prepends the short version of config.Kind to path.
func (a *Agent) pathWithPrefix(config resourced_config.Config) string {
	if config.Kind == "reader" {
		return a.pathWithReaderPrefix(config)
	} else if config.Kind == "writer" {
		return a.pathWithWriterPrefix(config)
	}
	return config.Path
}

func (a *Agent) pathWithReaderPrefix(input interface{}) string {
	switch v := input.(type) {
	case resourced_config.Config:
		return "/r" + v.Path
	case string:
		if strings.HasPrefix(v, "/r/") {
			return v
		} else {
			return "/r" + v
		}
	}
	return ""
}

func (a *Agent) pathWithWriterPrefix(input interface{}) string {
	switch v := input.(type) {
	case resourced_config.Config:
		return "/w" + v.Path
	case string:
		if strings.HasPrefix(v, "/w/") {
			return v
		} else {
			return "/w" + v
		}
	}
	return ""
}

// Run executes a reader/writer config.
// Run will save reader data as JSON in local db.
func (a *Agent) Run(config resourced_config.Config) (output []byte, err error) {
	if config.Command != "" {
		output, err = a.runCommand(config)
	} else if config.GoStruct != "" && config.Kind == "reader" {
		output, err = a.runGoStructReader(config)
	} else if config.GoStruct != "" && config.Kind == "writer" {
		output, err = a.runGoStructWriter(config)
	}
	if err != nil {
		return output, err
	}

	err = a.saveRun(config, output)

	return output, err
}

// runCommand shells out external program and returns the output.
func (a *Agent) runCommand(config resourced_config.Config) ([]byte, error) {
	cmd := libprocess.NewCmd(config.Command)

	if config.Kind == "writer" {
		// Get readers data.
		readersData := make(map[string]interface{})

		for _, readerPath := range config.ReaderPaths {
			readerJsonBytes, err := a.GetRunByPath(a.pathWithReaderPrefix(readerPath))

			if err == nil {
				var data interface{}
				err := json.Unmarshal(readerJsonBytes, &data)
				if err == nil {
					readersData[readerPath] = data
				}
			}
		}

		readersDataJsonBytes, err := json.Marshal(readersData)
		if err != nil {
			return nil, err
		}

		cmd.Stdin = bytes.NewReader(readersDataJsonBytes)
	}

	return cmd.Output()
}

// initGoStructReader initialize and return IReader.
func (a *Agent) initGoStructReader(config resourced_config.Config) (resourced_readers.IReader, error) {
	// Initialize IReader
	reader, err := resourced_readers.NewGoStruct(config.GoStruct)
	if err != nil {
		return nil, err
	}

	// Populate IReader fields dynamically
	if len(config.GoStructFields) > 0 {
		for structFieldInString, value := range config.GoStructFields {
			goStructField := reflect.ValueOf(reader).Elem().FieldByName(structFieldInString)

			if goStructField.IsValid() && goStructField.CanSet() {
				valueOfValue := reflect.ValueOf(value)
				goStructField.Set(valueOfValue)
			}
		}
	}

	return reader, err
}

// initGoStructWriter initialize and return IWriter.
func (a *Agent) initGoStructWriter(config resourced_config.Config) (resourced_writers.IWriter, error) {
	// Initialize IWriter
	writer, err := resourced_writers.NewGoStruct(config.GoStruct)
	if err != nil {
		return nil, err
	}

	// Populate IWriter fields dynamically
	if len(config.GoStructFields) > 0 {
		for structFieldInString, value := range config.GoStructFields {
			goStructField := reflect.ValueOf(writer).Elem().FieldByName(structFieldInString)

			if goStructField.IsValid() && goStructField.CanSet() {
				valueOfValue := reflect.ValueOf(value)
				goStructField.Set(valueOfValue)
			}
		}
	}

	// Get readers data.
	readersData := make(map[string][]byte)

	for _, readerPath := range config.ReaderPaths {
		readerJsonBytes, err := a.GetRunByPath(a.pathWithReaderPrefix(readerPath))
		if err == nil {
			readersData[readerPath] = readerJsonBytes
		}
	}

	writer.SetReadersData(readersData)

	return writer, err
}

// runGoStruct executes IReader/IWriter and returns the output.
// Note that IWriter also implements IReader
func (a *Agent) runGoStruct(readerOrWriter resourced_readers.IReader) ([]byte, error) {
	err := readerOrWriter.Run()
	if err != nil {
		errData := make(map[string]string)
		errData["Error"] = err.Error()
		return json.Marshal(errData)
	}

	return readerOrWriter.ToJson()
}

// runGoStructReader executes IReader and returns the output.
func (a *Agent) runGoStructReader(config resourced_config.Config) ([]byte, error) {
	// Initialize IReader
	reader, err := a.initGoStructReader(config)
	if err != nil {
		return nil, err
	}

	return a.runGoStruct(reader)
}

// runGoStructWriter executes IWriter and returns error if exists.
func (a *Agent) runGoStructWriter(config resourced_config.Config) ([]byte, error) {
	// Initialize IWriter
	writer, err := a.initGoStructWriter(config)
	if err != nil {
		return nil, err
	}

	return a.runGoStruct(writer)
}

// saveRun gathers default basic information and saves output into local storage.
func (a *Agent) saveRun(config resourced_config.Config, output []byte) error {
	// Do not perform save if config.Path is empty.
	if config.Path == "" {
		return nil
	}

	record := make(map[string]interface{})
	record["UnixNano"] = time.Now().UnixNano()
	record["Path"] = config.Path
	record["Interval"] = config.Interval
	record["Tags"] = a.Tags

	if config.Command != "" {
		record["Command"] = config.Command
	}

	if config.GoStruct != "" {
		record["GoStruct"] = config.GoStruct
	}

	hostname, err := os.Hostname()
	if err != nil {
		return err
	}
	record["Hostname"] = hostname

	runData := make(map[string]interface{})
	err = json.Unmarshal(output, &runData)
	if err != nil {
		return err
	}

	record["Data"] = runData

	recordInJson, err := json.Marshal(record)
	if err != nil {
		return err
	}

	err = a.Db.Update(func(tx *bolt.Tx) error {
		return a.dbBucket(tx).Put([]byte(a.pathWithPrefix(config)), recordInJson)
	})

	return err
}

// GetRun returns the JSON data stored in local storage given Config struct.
func (a *Agent) GetRun(config resourced_config.Config) ([]byte, error) {
	return a.GetRunByPath(a.pathWithPrefix(config))
}

// GetRunByPath returns JSON data stored in local storage given path string.
func (a *Agent) GetRunByPath(path string) ([]byte, error) {
	var data []byte

	a.Db.View(func(tx *bolt.Tx) error {
		data = a.dbBucket(tx).Get([]byte(path))
		return nil
	})

	return data, nil
}

// RunForever executes Run() in an infinite loop with a sleep of config.Interval.
func (a *Agent) RunForever(config resourced_config.Config) {
	go func(a *Agent, config resourced_config.Config) {
		for {
			a.Run(config)
			libtime.SleepString(config.Interval)
		}
	}(a, config)
}

// RunAllForever executes all readers & writers in an infinite loop.
func (a *Agent) RunAllForever() {
	for _, config := range a.ConfigStorage.Readers {
		a.RunForever(config)
	}
	for _, config := range a.ConfigStorage.Writers {
		a.RunForever(config)
	}
}
<commit_msg>Able to provide net/interfaces data on every reader.<commit_after>package agent

import (
	"bytes"
	"encoding/json"
	"github.com/boltdb/bolt"
	resourced_config "github.com/resourced/resourced/config"
	"github.com/resourced/resourced/libprocess"
	"github.com/resourced/resourced/libstring"
	"github.com/resourced/resourced/libtime"
	resourced_readers "github.com/resourced/resourced/readers"
	resourced_writers "github.com/resourced/resourced/writers"
	"os"
	"reflect"
	"strings"
	"time"
)

// NewAgent is the constructor fot Agent struct.
func NewAgent() (*Agent, error) {
	agent := &Agent{}

	agent.setTags()

	err := agent.setConfigStorage()
	if err != nil {
		return nil, err
	}

	err = agent.setDb()
	if err != nil {
		return nil, err
	}

	return agent, err
}

// Agent struct carries most of the functionality of ResourceD.
// It collects information through readers and serve them up as HTTP+JSON.
type Agent struct {
	ConfigStorage *resourced_config.ConfigStorage
	DbPath        string
	Db            *bolt.DB
	Tags          []string
}

func (a *Agent) setTags() {
	a.Tags = make([]string, 0)

	tags := os.Getenv("RESOURCED_TAGS")
	if tags != "" {
		tagsSlice := strings.Split(tags, ",")
		a.Tags = make([]string, len(tagsSlice))

		for i, tag := range tagsSlice {
			a.Tags[i] = strings.TrimSpace(tag)
		}
	}
}

// setDb configures the local storage.
func (a *Agent) setDb() error {
	var err error

	dbPath := os.Getenv("RESOURCED_DB")
	if dbPath == "" {
		dbPath = "~/resourced/db"

		err = os.MkdirAll(libstring.ExpandTilde("~/resourced"), 0755)
		if err != nil {
			return err
		}
	}

	a.DbPath = libstring.ExpandTilde(dbPath)

	a.Db, err = bolt.Open(a.DbPath, 0644, nil)
	if err != nil {
		return err
	}

	// Create "resources" bucket
	a.Db.Update(func(tx *bolt.Tx) error {
		tx.CreateBucket([]byte("resources"))
		return nil
	})

	return err
}

// dbBucket returns the boltdb bucket.
func (a *Agent) dbBucket(tx *bolt.Tx) *bolt.Bucket {
	return tx.Bucket([]byte("resources"))
}

// pathWithPrefix prepends the short version of config.Kind to path.
func (a *Agent) pathWithPrefix(config resourced_config.Config) string {
	if config.Kind == "reader" {
		return a.pathWithReaderPrefix(config)
	} else if config.Kind == "writer" {
		return a.pathWithWriterPrefix(config)
	}
	return config.Path
}

func (a *Agent) pathWithReaderPrefix(input interface{}) string {
	switch v := input.(type) {
	case resourced_config.Config:
		return "/r" + v.Path
	case string:
		if strings.HasPrefix(v, "/r/") {
			return v
		} else {
			return "/r" + v
		}
	}
	return ""
}

func (a *Agent) pathWithWriterPrefix(input interface{}) string {
	switch v := input.(type) {
	case resourced_config.Config:
		return "/w" + v.Path
	case string:
		if strings.HasPrefix(v, "/w/") {
			return v
		} else {
			return "/w" + v
		}
	}
	return ""
}

// Run executes a reader/writer config.
// Run will save reader data as JSON in local db.
func (a *Agent) Run(config resourced_config.Config) (output []byte, err error) {
	if config.Command != "" {
		output, err = a.runCommand(config)
	} else if config.GoStruct != "" && config.Kind == "reader" {
		output, err = a.runGoStructReader(config)
	} else if config.GoStruct != "" && config.Kind == "writer" {
		output, err = a.runGoStructWriter(config)
	}
	if err != nil {
		return output, err
	}

	err = a.saveRun(config, output)

	return output, err
}

// runCommand shells out external program and returns the output.
func (a *Agent) runCommand(config resourced_config.Config) ([]byte, error) {
	cmd := libprocess.NewCmd(config.Command)

	if config.Kind == "writer" {
		// Get readers data.
		readersData := make(map[string]interface{})

		for _, readerPath := range config.ReaderPaths {
			readerJsonBytes, err := a.GetRunByPath(a.pathWithReaderPrefix(readerPath))

			if err == nil {
				var data interface{}
				err := json.Unmarshal(readerJsonBytes, &data)
				if err == nil {
					readersData[readerPath] = data
				}
			}
		}

		readersDataJsonBytes, err := json.Marshal(readersData)
		if err != nil {
			return nil, err
		}

		cmd.Stdin = bytes.NewReader(readersDataJsonBytes)
	}

	return cmd.Output()
}

// initGoStructReader initialize and return IReader.
func (a *Agent) initGoStructReader(config resourced_config.Config) (resourced_readers.IReader, error) {
	// Initialize IReader
	reader, err := resourced_readers.NewGoStruct(config.GoStruct)
	if err != nil {
		return nil, err
	}

	// Populate IReader fields dynamically
	if len(config.GoStructFields) > 0 {
		for structFieldInString, value := range config.GoStructFields {
			goStructField := reflect.ValueOf(reader).Elem().FieldByName(structFieldInString)

			if goStructField.IsValid() && goStructField.CanSet() {
				valueOfValue := reflect.ValueOf(value)
				goStructField.Set(valueOfValue)
			}
		}
	}

	return reader, err
}

// initGoStructWriter initialize and return IWriter.
func (a *Agent) initGoStructWriter(config resourced_config.Config) (resourced_writers.IWriter, error) {
	// Initialize IWriter
	writer, err := resourced_writers.NewGoStruct(config.GoStruct)
	if err != nil {
		return nil, err
	}

	// Populate IWriter fields dynamically
	if len(config.GoStructFields) > 0 {
		for structFieldInString, value := range config.GoStructFields {
			goStructField := reflect.ValueOf(writer).Elem().FieldByName(structFieldInString)

			if goStructField.IsValid() && goStructField.CanSet() {
				valueOfValue := reflect.ValueOf(value)
				goStructField.Set(valueOfValue)
			}
		}
	}

	// Get readers data.
	readersData := make(map[string][]byte)

	for _, readerPath := range config.ReaderPaths {
		readerJsonBytes, err := a.GetRunByPath(a.pathWithReaderPrefix(readerPath))
		if err == nil {
			readersData[readerPath] = readerJsonBytes
		}
	}

	writer.SetReadersData(readersData)

	return writer, err
}

// runGoStruct executes IReader/IWriter and returns the output.
// Note that IWriter also implements IReader
func (a *Agent) runGoStruct(readerOrWriter resourced_readers.IReader) ([]byte, error) {
	err := readerOrWriter.Run()
	if err != nil {
		errData := make(map[string]string)
		errData["Error"] = err.Error()
		return json.Marshal(errData)
	}

	return readerOrWriter.ToJson()
}

// runGoStructReader executes IReader and returns the output.
func (a *Agent) runGoStructReader(config resourced_config.Config) ([]byte, error) {
	// Initialize IReader
	reader, err := a.initGoStructReader(config)
	if err != nil {
		return nil, err
	}

	return a.runGoStruct(reader)
}

// runGoStructWriter executes IWriter and returns error if exists.
func (a *Agent) runGoStructWriter(config resourced_config.Config) ([]byte, error) {
	// Initialize IWriter
	writer, err := a.initGoStructWriter(config)
	if err != nil {
		return nil, err
	}

	return a.runGoStruct(writer)
}

// saveRun gathers default basic information and saves output into local storage.
func (a *Agent) saveRun(config resourced_config.Config, output []byte) error {
	// Do not perform save if config.Path is empty.
	if config.Path == "" {
		return nil
	}

	record := make(map[string]interface{})
	record["UnixNano"] = time.Now().UnixNano()
	record["Path"] = config.Path
	record["Interval"] = config.Interval
	record["Tags"] = a.Tags

	if config.Command != "" {
		record["Command"] = config.Command
	}

	if config.GoStruct != "" {
		record["GoStruct"] = config.GoStruct
	}

	hostname, err := os.Hostname()
	if err != nil {
		return err
	}
	record["Hostname"] = hostname

	// net/interfaces data
	interfacesReader := resourced_readers.NewNetInterfaces()
	if interfacesReader.Run() == nil {
		record["NetworkInterfaces"] = make(map[string]map[string]interface{})

		for iface, stats := range interfacesReader.Data {
			netInterfaces := record["NetworkInterfaces"].(map[string]map[string]interface{})

			netInterfaces[iface] = make(map[string]interface{})
			netInterfaces[iface]["HardwareAddress"] = stats.HardwareAddr
			netInterfaces[iface]["IPAddresses"] = make([]string, len(stats.Addrs))

			for i, addr := range stats.Addrs {
				ipAddresses := netInterfaces[iface]["IPAddresses"].([]string)
				ipAddresses[i] = addr.Addr
			}
		}
	}

	runData := make(map[string]interface{})
	err = json.Unmarshal(output, &runData)
	if err != nil {
		return err
	}

	record["Data"] = runData

	recordInJson, err := json.Marshal(record)
	if err != nil {
		return err
	}

	err = a.Db.Update(func(tx *bolt.Tx) error {
		return a.dbBucket(tx).Put([]byte(a.pathWithPrefix(config)), recordInJson)
	})

	return err
}

// GetRun returns the JSON data stored in local storage given Config struct.
func (a *Agent) GetRun(config resourced_config.Config) ([]byte, error) {
	return a.GetRunByPath(a.pathWithPrefix(config))
}

// GetRunByPath returns JSON data stored in local storage given path string.
func (a *Agent) GetRunByPath(path string) ([]byte, error) {
	var data []byte

	a.Db.View(func(tx *bolt.Tx) error {
		data = a.dbBucket(tx).Get([]byte(path))
		return nil
	})

	return data, nil
}

// RunForever executes Run() in an infinite loop with a sleep of config.Interval.
func (a *Agent) RunForever(config resourced_config.Config) {
	go func(a *Agent, config resourced_config.Config) {
		for {
			a.Run(config)
			libtime.SleepString(config.Interval)
		}
	}(a, config)
}

// RunAllForever executes all readers & writers in an infinite loop.
func (a *Agent) RunAllForever() {
	for _, config := range a.ConfigStorage.Readers {
		a.RunForever(config)
	}
	for _, config := range a.ConfigStorage.Writers {
		a.RunForever(config)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 Google Inc. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Command agent forwards requests from an inverting proxy to a backend server.
//
// To build, run:
//
//    $ make
//
// And to use, run:
//
//    $ $(GOPATH)/bin/proxy-forwarding-agent -proxy <proxy-url> -backend <backend-ID>

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/http/cookiejar"
	"net/http/httputil"
	"net/url"
	"strings"
	"time"

	"github.com/golang/groupcache/lru"
	"golang.org/x/net/publicsuffix"
	"golang.org/x/oauth2/google"
	compute "google.golang.org/api/compute/v1"

	"github.com/google/inverting-proxy/agent/banner"
	"github.com/google/inverting-proxy/agent/sessions"
	"github.com/google/inverting-proxy/agent/utils"
	"github.com/google/inverting-proxy/agent/websockets"
)

const (
	requestCacheLimit = 1000
	emailScope        = "email"
)

var (
	proxy                = flag.String("proxy", "", "URL (including scheme) of the inverting proxy")
	proxyTimeout         = flag.Duration("proxy-timeout", 60*time.Second, "Client timeout when sending requests to the inverting proxy")
	host                 = flag.String("host", "localhost:8080", "Hostname (including port) of the backend server")
	backendID            = flag.String("backend", "", "Unique ID for this backend.")
	debug                = flag.Bool("debug", false, "Whether or not to print debug log messages")
	disableSSLForTest    = flag.Bool("disable-ssl-for-test", false, "Disable requirements for SSL when running tests locally")
	favIconURL           = flag.String("favicon-url", "", "URL of the favicon")
	forwardUserID        = flag.Bool("forward-user-id", false, "Whether or not to include the ID (email address) of the end user in requests to the backend")
	injectBanner         = flag.String("inject-banner", "", "HTML snippet to inject in served webpages")
	bannerHeight         = flag.String("banner-height", "40px", "Height of the injected banner. This is ignored if no banner is set.")
	shimWebsockets       = flag.Bool("shim-websockets", false, "Whether or not to replace websockets with a shim")
	shimPath             = flag.String("shim-path", "", "Path under which to handle websocket shim requests")
	healthCheckPath      = flag.String("health-check-path", "/", "Path on backend host to issue health checks against.  Defaults to the root.")
	healthCheckFreq      = flag.Int("health-check-interval-seconds", 0, "Wait time in seconds between health checks.  Set to zero to disable health checks.  Checks disabled by default.")
	healthCheckUnhealthy = flag.Int("health-check-unhealthy-threshold", 2, "A so-far healthy backend will be marked unhealthy after this many consecutive failures. The minimum value is 1.")
	disableGCEVM         = flag.Bool("disable-gce-vm-header", false, "Disable the agent from adding a GCE VM header.")

	sessionCookieName       = flag.String("session-cookie-name", "", "Name of the session cookie; an empty value disables agent-based session tracking")
	sessionCookieTimeout    = flag.Duration("session-cookie-timeout", 12*time.Hour, "Expiration flag for the session cookie")
	sessionCookieCacheLimit = flag.Int("session-cookie-cache-limit", 1000, "Upper bound on the number of concurrent sessions that can be tracked by the agent")
	rewriteWebsocketHost    = flag.Bool("rewrite-websocket-host", false, "Whether to rewrite the Host header to the original request when shimming a websocket connection")

	sessionLRU *sessions.Cache
)

func hostProxy(ctx context.Context, host, shimPath string, injectShimCode bool) (http.Handler, error) {
	hostProxy := httputil.NewSingleHostReverseProxy(&url.URL{
		Scheme: "http",
		Host:   host,
	})
	hostProxy.FlushInterval = 100 * time.Millisecond
	var h http.Handler = hostProxy
	h = sessionLRU.SessionHandler(h)
	if shimPath != "" {
		var err error
		// Note that we pass in the sessionHandler to the websocket proxy twice (h and sessionLRU.SessionHandler)
		// h is the wrapped handler that will handle all non-websocket-shim requests
		// sessionLRU.SessionHandler will be called for websocket open requests
		// This is necessary to handle an edge case in session handling. Websocket open requests arrive with a path
		// of `/$shimPath/open`. The original target URL and path are encoded in the request body, which is
		// restored in the websocket handler. This means that attempting to restore session cookies that are
		// restricted to a path prefix not equal to "/" will fail for websocket open requests. Passing in the
		// sessionHandler twice allows the websocket handler to ensure that cookies are applied based on the
		// correct, restored path.
		h, err = websockets.Proxy(ctx, h, host, shimPath, *rewriteWebsocketHost, sessionLRU.SessionHandler)
		if injectShimCode {
			shimFunc, err := websockets.ShimBody(shimPath)
			if err != nil {
				return nil, err
			}
			hostProxy.ModifyResponse = shimFunc
		}
		if err != nil {
			return nil, err
		}
	}
	if *injectBanner == "" {
		return h, nil
	}
	return banner.Proxy(ctx, h, *injectBanner, *bannerHeight, *favIconURL)
}

// forwardRequest forwards the given request from the proxy to
// the backend server and reports the response back to the proxy.
func forwardRequest(client *http.Client, hostProxy http.Handler, request *utils.ForwardedRequest) error {
	httpRequest := request.Contents
	if *forwardUserID {
		httpRequest.Header.Add(utils.HeaderUserID, request.User)
	}
	responseForwarder, err := utils.NewResponseForwarder(client, *proxy, request.BackendID, request.RequestID)
	if err != nil {
		return fmt.Errorf("failed to create the response forwarder: %v", err)
	}
	hostProxy.ServeHTTP(responseForwarder, httpRequest)
	if *debug {
		log.Printf("Backend latency for request %s: %s\n", request.RequestID, time.Since(request.StartTime).String())
	}
	if err := responseForwarder.Close(); err != nil {
		return fmt.Errorf("failed to close the response forwarder: %v", err)
	}
	return nil
}

// healthCheck issues a health check against the backend server
// and returns the result.
func healthCheck() error {
	resp, err := http.Get("http://" + *host + *healthCheckPath)
	if err != nil {
		log.Printf("Health Check request failed: %s", err.Error())
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		log.Printf("Health Check request had non-200 status code: %d", resp.StatusCode)
		return fmt.Errorf("Bad Health Check Response Code: %s", resp.Status)
	}
	return nil
}

// processOneRequest reads a single request from the proxy and forwards it to the backend server.
func processOneRequest(client *http.Client, hostProxy http.Handler, backendID string, requestID string) {
	requestForwarder := func(client *http.Client, request *utils.ForwardedRequest) error {
		if err := forwardRequest(client, hostProxy, request); err != nil {
			log.Printf("Failure forwarding a request: [%s] %q\n", requestID, err.Error())
			return fmt.Errorf("failed to forward the request %q: %v", requestID, err)
		}
		return nil
	}
	if err := utils.ReadRequest(client, *proxy, backendID, requestID, requestForwarder); err != nil {
		log.Printf("Failed to forward a request: [%s] %q\n", requestID, err.Error())
	}
}

// pollForNewRequests repeatedly reaches out to the proxy server to ask if any pending are available, and then
// processes any newly-seen ones.
func pollForNewRequests(client *http.Client, hostProxy http.Handler, backendID string) {
	previouslySeenRequests := lru.New(requestCacheLimit)
	var retryCount uint
	for {
		if requests, err := utils.ListPendingRequests(client, *proxy, backendID); err != nil {
			log.Printf("Failed to read pending requests: %q\n", err.Error())
			time.Sleep(utils.ExponentialBackoffDuration(retryCount))
			retryCount++
		} else {
			retryCount = 0
			for _, requestID := range requests {
				if _, ok := previouslySeenRequests.Get(requestID); !ok {
					previouslySeenRequests.Add(requestID, requestID)
					go processOneRequest(client, hostProxy, backendID, requestID)
				}
			}
		}
	}
}

func getGoogleClient(ctx context.Context) (*http.Client, error) {
	sdkConfig, err := google.NewSDKConfig("")
	if err == nil {
		return sdkConfig.Client(ctx), nil
	}

	client, err := google.DefaultClient(ctx, compute.CloudPlatformScope, emailScope)
	if err != nil {
		return nil, err
	}
	client.Transport = utils.RoundTripperWithVMIdentity(ctx, client.Transport, *proxy, *disableGCEVM)
	client.Jar, err = cookiejar.New(&cookiejar.Options{
		PublicSuffixList: publicsuffix.List,
	})
	if err != nil {
		return nil, err
	}
	return client, nil
}

// waitForHealthy runs health checks against the backend and returns
// the first time it sees a healthy check.
func waitForHealthy() {
	if *healthCheckFreq <= 0 {
		return
	}
	if healthCheck() == nil {
		return
	}
	ticker := time.NewTicker(time.Duration(*healthCheckFreq) * time.Second)
	for _ = range ticker.C {
		if healthCheck() == nil {
			ticker.Stop()
			return
		}
	}
}

// runHealthChecks runs health checks against the backend and shuts down
// the proxy if the backend is unhealthy.
func runHealthChecks() {
	if *healthCheckFreq <= 0 {
		return
	}
	if *healthCheckUnhealthy < 1 {
		*healthCheckUnhealthy = 1
	}
	// Always start in the unhealthy state, but only require a single positive
	// health check to become healthy for the first time, and do the first check
	// immediately.
	ticker := time.NewTicker(time.Duration(*healthCheckFreq) * time.Second)
	badHealthChecks := 0
	for _ = range ticker.C {
		if healthCheck() != nil {
			badHealthChecks++
		} else {
			badHealthChecks = 0
		}
		if badHealthChecks >= *healthCheckUnhealthy {
			ticker.Stop()
			log.Fatal("Too many unhealthy checks")
		}
	}
}

// runAdapter sets up the HTTP client for the agent to use (including OAuth credentials),
// and then does the actual work of forwarding requests and responses.
func runAdapter() error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client, err := getGoogleClient(ctx)
	if err != nil {
		return err
	}
	client.Timeout = *proxyTimeout

	hostProxy, err := hostProxy(ctx, *host, *shimPath, *shimWebsockets)
	if err != nil {
		return err
	}
	pollForNewRequests(client, hostProxy, *backendID)
	return nil
}

func main() {
	flag.Parse()

	if *proxy == "" {
		log.Fatal("You must specify the address of the proxy")
	}
	if *backendID == "" {
		log.Fatal("You must specify a backend ID")
	}
	if !strings.HasPrefix(*healthCheckPath, "/") {
		*healthCheckPath = "/" + *healthCheckPath
	}
	if *sessionCookieName != "" {
		sessionLRU = sessions.NewCache(*sessionCookieName, *sessionCookieTimeout, *sessionCookieCacheLimit, *disableSSLForTest)
	}

	waitForHealthy()
	go runHealthChecks()

	if err := runAdapter(); err != nil {
		log.Fatal(err.Error())
	}
}
<commit_msg>Support stripping the Authorization header from requests<commit_after>/*
Copyright 2016 Google Inc. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Command agent forwards requests from an inverting proxy to a backend server.
//
// To build, run:
//
//    $ make
//
// And to use, run:
//
//    $ $(GOPATH)/bin/proxy-forwarding-agent -proxy <proxy-url> -backend <backend-ID>

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/http/cookiejar"
	"net/http/httputil"
	"net/url"
	"strings"
	"time"

	"github.com/golang/groupcache/lru"
	"golang.org/x/net/publicsuffix"
	"golang.org/x/oauth2/google"
	compute "google.golang.org/api/compute/v1"

	"github.com/google/inverting-proxy/agent/banner"
	"github.com/google/inverting-proxy/agent/sessions"
	"github.com/google/inverting-proxy/agent/utils"
	"github.com/google/inverting-proxy/agent/websockets"
)

const (
	requestCacheLimit   = 1000
	emailScope          = "email"
	headerAuthorization = "Authorization"
)

var (
	proxy                = flag.String("proxy", "", "URL (including scheme) of the inverting proxy")
	proxyTimeout         = flag.Duration("proxy-timeout", 60*time.Second, "Client timeout when sending requests to the inverting proxy")
	host                 = flag.String("host", "localhost:8080", "Hostname (including port) of the backend server")
	backendID            = flag.String("backend", "", "Unique ID for this backend.")
	debug                = flag.Bool("debug", false, "Whether or not to print debug log messages")
	disableSSLForTest    = flag.Bool("disable-ssl-for-test", false, "Disable requirements for SSL when running tests locally")
	favIconURL           = flag.String("favicon-url", "", "URL of the favicon")
	forwardUserID        = flag.Bool("forward-user-id", false, "Whether or not to include the ID (email address) of the end user in requests to the backend")
	injectBanner         = flag.String("inject-banner", "", "HTML snippet to inject in served webpages")
	bannerHeight         = flag.String("banner-height", "40px", "Height of the injected banner. This is ignored if no banner is set.")
	shimWebsockets       = flag.Bool("shim-websockets", false, "Whether or not to replace websockets with a shim")
	shimPath             = flag.String("shim-path", "", "Path under which to handle websocket shim requests")
	healthCheckPath      = flag.String("health-check-path", "/", "Path on backend host to issue health checks against.  Defaults to the root.")
	healthCheckFreq      = flag.Int("health-check-interval-seconds", 0, "Wait time in seconds between health checks.  Set to zero to disable health checks.  Checks disabled by default.")
	healthCheckUnhealthy = flag.Int("health-check-unhealthy-threshold", 2, "A so-far healthy backend will be marked unhealthy after this many consecutive failures. The minimum value is 1.")
	disableGCEVM         = flag.Bool("disable-gce-vm-header", false, "Disable the agent from adding a GCE VM header.")

	sessionCookieName       = flag.String("session-cookie-name", "", "Name of the session cookie; an empty value disables agent-based session tracking")
	sessionCookieTimeout    = flag.Duration("session-cookie-timeout", 12*time.Hour, "Expiration flag for the session cookie")
	sessionCookieCacheLimit = flag.Int("session-cookie-cache-limit", 1000, "Upper bound on the number of concurrent sessions that can be tracked by the agent")
	rewriteWebsocketHost    = flag.Bool("rewrite-websocket-host", false, "Whether to rewrite the Host header to the original request when shimming a websocket connection")
	stripCredentials        = flag.Bool("strip-credentials", false, "Whether to strip the Authorization header from all requests.")

	sessionLRU *sessions.Cache
)

func hostProxy(ctx context.Context, host, shimPath string, injectShimCode bool) (http.Handler, error) {
	hostProxy := httputil.NewSingleHostReverseProxy(&url.URL{
		Scheme: "http",
		Host:   host,
	})
	hostProxy.FlushInterval = 100 * time.Millisecond
	var h http.Handler = hostProxy
	h = sessionLRU.SessionHandler(h)
	if shimPath != "" {
		var err error
		// Note that we pass in the sessionHandler to the websocket proxy twice (h and sessionLRU.SessionHandler)
		// h is the wrapped handler that will handle all non-websocket-shim requests
		// sessionLRU.SessionHandler will be called for websocket open requests
		// This is necessary to handle an edge case in session handling. Websocket open requests arrive with a path
		// of `/$shimPath/open`. The original target URL and path are encoded in the request body, which is
		// restored in the websocket handler. This means that attempting to restore session cookies that are
		// restricted to a path prefix not equal to "/" will fail for websocket open requests. Passing in the
		// sessionHandler twice allows the websocket handler to ensure that cookies are applied based on the
		// correct, restored path.
		h, err = websockets.Proxy(ctx, h, host, shimPath, *rewriteWebsocketHost, sessionLRU.SessionHandler)
		if injectShimCode {
			shimFunc, err := websockets.ShimBody(shimPath)
			if err != nil {
				return nil, err
			}
			hostProxy.ModifyResponse = shimFunc
		}
		if err != nil {
			return nil, err
		}
	}
	if *injectBanner == "" {
		return h, nil
	}
	return banner.Proxy(ctx, h, *injectBanner, *bannerHeight, *favIconURL)
}

// forwardRequest forwards the given request from the proxy to
// the backend server and reports the response back to the proxy.
func forwardRequest(client *http.Client, hostProxy http.Handler, request *utils.ForwardedRequest) error {
	httpRequest := request.Contents
	if *forwardUserID {
		httpRequest.Header.Add(utils.HeaderUserID, request.User)
	}
	if *stripCredentials {
		httpRequest.Header.Del(headerAuthorization)
	}
	responseForwarder, err := utils.NewResponseForwarder(client, *proxy, request.BackendID, request.RequestID)
	if err != nil {
		return fmt.Errorf("failed to create the response forwarder: %v", err)
	}
	hostProxy.ServeHTTP(responseForwarder, httpRequest)
	if *debug {
		log.Printf("Backend latency for request %s: %s\n", request.RequestID, time.Since(request.StartTime).String())
	}
	if err := responseForwarder.Close(); err != nil {
		return fmt.Errorf("failed to close the response forwarder: %v", err)
	}
	return nil
}

// healthCheck issues a health check against the backend server
// and returns the result.
func healthCheck() error {
	resp, err := http.Get("http://" + *host + *healthCheckPath)
	if err != nil {
		log.Printf("Health Check request failed: %s", err.Error())
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		log.Printf("Health Check request had non-200 status code: %d", resp.StatusCode)
		return fmt.Errorf("Bad Health Check Response Code: %s", resp.Status)
	}
	return nil
}

// processOneRequest reads a single request from the proxy and forwards it to the backend server.
func processOneRequest(client *http.Client, hostProxy http.Handler, backendID string, requestID string) {
	requestForwarder := func(client *http.Client, request *utils.ForwardedRequest) error {
		if err := forwardRequest(client, hostProxy, request); err != nil {
			log.Printf("Failure forwarding a request: [%s] %q\n", requestID, err.Error())
			return fmt.Errorf("failed to forward the request %q: %v", requestID, err)
		}
		return nil
	}
	if err := utils.ReadRequest(client, *proxy, backendID, requestID, requestForwarder); err != nil {
		log.Printf("Failed to forward a request: [%s] %q\n", requestID, err.Error())
	}
}

// pollForNewRequests repeatedly reaches out to the proxy server to ask if any pending are available, and then
// processes any newly-seen ones.
func pollForNewRequests(client *http.Client, hostProxy http.Handler, backendID string) {
	previouslySeenRequests := lru.New(requestCacheLimit)
	var retryCount uint
	for {
		if requests, err := utils.ListPendingRequests(client, *proxy, backendID); err != nil {
			log.Printf("Failed to read pending requests: %q\n", err.Error())
			time.Sleep(utils.ExponentialBackoffDuration(retryCount))
			retryCount++
		} else {
			retryCount = 0
			for _, requestID := range requests {
				if _, ok := previouslySeenRequests.Get(requestID); !ok {
					previouslySeenRequests.Add(requestID, requestID)
					go processOneRequest(client, hostProxy, backendID, requestID)
				}
			}
		}
	}
}

func getGoogleClient(ctx context.Context) (*http.Client, error) {
	sdkConfig, err := google.NewSDKConfig("")
	if err == nil {
		return sdkConfig.Client(ctx), nil
	}

	client, err := google.DefaultClient(ctx, compute.CloudPlatformScope, emailScope)
	if err != nil {
		return nil, err
	}
	client.Transport = utils.RoundTripperWithVMIdentity(ctx, client.Transport, *proxy, *disableGCEVM)
	client.Jar, err = cookiejar.New(&cookiejar.Options{
		PublicSuffixList: publicsuffix.List,
	})
	if err != nil {
		return nil, err
	}
	return client, nil
}

// waitForHealthy runs health checks against the backend and returns
// the first time it sees a healthy check.
func waitForHealthy() {
	if *healthCheckFreq <= 0 {
		return
	}
	if healthCheck() == nil {
		return
	}
	ticker := time.NewTicker(time.Duration(*healthCheckFreq) * time.Second)
	for _ = range ticker.C {
		if healthCheck() == nil {
			ticker.Stop()
			return
		}
	}
}

// runHealthChecks runs health checks against the backend and shuts down
// the proxy if the backend is unhealthy.
func runHealthChecks() {
	if *healthCheckFreq <= 0 {
		return
	}
	if *healthCheckUnhealthy < 1 {
		*healthCheckUnhealthy = 1
	}
	// Always start in the unhealthy state, but only require a single positive
	// health check to become healthy for the first time, and do the first check
	// immediately.
	ticker := time.NewTicker(time.Duration(*healthCheckFreq) * time.Second)
	badHealthChecks := 0
	for _ = range ticker.C {
		if healthCheck() != nil {
			badHealthChecks++
		} else {
			badHealthChecks = 0
		}
		if badHealthChecks >= *healthCheckUnhealthy {
			ticker.Stop()
			log.Fatal("Too many unhealthy checks")
		}
	}
}

// runAdapter sets up the HTTP client for the agent to use (including OAuth credentials),
// and then does the actual work of forwarding requests and responses.
func runAdapter() error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client, err := getGoogleClient(ctx)
	if err != nil {
		return err
	}
	client.Timeout = *proxyTimeout

	hostProxy, err := hostProxy(ctx, *host, *shimPath, *shimWebsockets)
	if err != nil {
		return err
	}
	pollForNewRequests(client, hostProxy, *backendID)
	return nil
}

func main() {
	flag.Parse()

	if *proxy == "" {
		log.Fatal("You must specify the address of the proxy")
	}
	if *backendID == "" {
		log.Fatal("You must specify a backend ID")
	}
	if !strings.HasPrefix(*healthCheckPath, "/") {
		*healthCheckPath = "/" + *healthCheckPath
	}
	if *sessionCookieName != "" {
		sessionLRU = sessions.NewCache(*sessionCookieName, *sessionCookieTimeout, *sessionCookieCacheLimit, *disableSSLForTest)
	}

	waitForHealthy()
	go runHealthChecks()

	if err := runAdapter(); err != nil {
		log.Fatal(err.Error())
	}
}
<|endoftext|>


<commit_before>package home

import (
    "fmt"
    "encoding/json"
    "net/http"
    "html/template"
    "errors"
    "database/sql"

    _ "github.com/go-sql-driver/mysql"
    "git-go-logictree/app/common"
)

type Condition struct {
    Text, Type, Field, Operator, Value string
}

type treeNode struct {
    Parent *treeNode
    Children []*treeNode
    Node Condition
    Left, Right int
}

type conditionSqlRow struct {
    Field, Operator, Value, Type string
    Left, Right int
}

func GetHomePage(rw http.ResponseWriter, req *http.Request) {
    type Page struct {
        Title string
        Conditions []Condition
    }

    p := Page{
        Title: "home",
        Conditions: getConditions(),
    }

    common.Templates = template.Must(template.ParseFiles("templates/home/home.html", common.LayoutPath))
    err := common.Templates.ExecuteTemplate(rw, "base", p)
    common.CheckError(err, 2)
}

func UpdateConditions(rw http.ResponseWriter, req *http.Request) {
    conditions := req.FormValue("conditions");

    parsedConditions, _ := parseJSON(conditions);
    fmt.Println(parsedConditions)
    treeRoot, err := unserializeFormattedTree(parsedConditions)

    equalityStr, logicStr, err := treeRoot.toMysql()
    if err != nil {
        fmt.Println(err)
        return
    }

    updateDatabase(equalityStr, logicStr)

    GetHomePage(rw, req)
}

func Truncate(rw http.ResponseWriter, req *http.Request) {
    _, err := common.DB.Query("TRUNCATE TABLE logictree.equality")
    common.CheckError(err, 2)

    _, err = common.DB.Query("TRUNCATE TABLE logictree.logic")
    common.CheckError(err, 2)

    GetHomePage(rw, req)
}

func getConditions() []Condition {
    conditions := make([]Condition, 0)

    rows, err := common.DB.Query("SELECT field, operator, value FROM logictree.equality")
    common.CheckError(err, 2)

    var field, operator, value string

    i := 0

    for rows.Next() {
        rows.Scan(&field, &operator, &value)
        common.CheckError(err, 2)

        if i != 0 {
            conditions = append(conditions, Condition{
                Text: "AND",
                Operator: "AND",
                Type: "logic",
            })
        }

        conditions = append(conditions, Condition{
            Text: fmt.Sprintf("%s %s %s", field, operator, value),
            Type: "equality",
            Field: field,
            Operator: operator,
            Value: value,
        })

        i++
    }

    return conditions
}

func parseJSON(conditionsString string) ([]Condition, error) {
    var conditionsSlice []Condition

    err := json.Unmarshal([]byte(conditionsString), &conditionsSlice)
    if err != nil {
        return nil, err
    }

    return conditionsSlice, nil
}

/** Treat conditions like a queue. Rules:
 * If you reach a (, pop the condition, drop down a depth and assign results to root's children
 * If you reach a ), pop the condition, pop back up a depth with the root
 * If you reach a logical condition, pop the condition, assign it as the root's node
 * If you reach an equality condition, pop the condition, assign it as one of the children of the root
 * At the end of the loop, return the root's first child (since we have parans around all conditions we're going to be one level too deep)
**/
func unserializeFormattedTree(conditions []Condition) (*treeNode, error) {
    var root treeNode
    var emptyNode, condition Condition

    key := 0

    for key < len(conditions) {
        // Pop the front item from the slice
        condition = conditions[0]
        conditions = append(conditions[:key], conditions[key+1:]...)

        switch condition.Type {
        case "scope":
            if condition.Operator == "(" {
                children, _ := unserializeFormattedTree(conditions)

                if len(root.Children) == 0 {
                    root.Children = []*treeNode{children}
                } else {
                    root.Children = append(root.Children, children)
                }
            }

            if condition.Operator == ")" {
                if root.Node == emptyNode {
                    return root.Children[0], nil
                } else {
                    return &root, nil
                }
            }
        case "logic":
            root.Node = condition
        case "equality":
            root.Children = append(root.Children, &treeNode{Parent: &root, Node: condition})
        }
    }

    if len(root.Children) > 0 && !root.Node.matches(emptyNode) {
        return &root, nil
    }

    return root.Children[0], nil
}

/**
 * Assumption: data should be ordered by LEFT
 * Recursive steps:
 * If current node is a branch:
 *      Create node with it
 *      Pop item from conditions
 *      Add children recursively
 * If node is a leaf:
 *      Create node with it
 *      Pop item from conditions
 * Return node
**/
func unserializeRawTree(conditions []conditionSqlRow) *treeNode {
    return nil
}

func serializeTree(node *treeNode) ([]Condition, error) {
    if node.Children == nil || len(node.Children) == 0 {
        // Has no children - should be equality

        if node.Node.Type != "equality" {
            return nil, errors.New("ERROR: This tree has a logic condition as a leaf. Quitting.")
        }

        return []Condition{node.Node}, nil
    } else {
        // Has children - should be logic

        if node.Node.Type != "logic" {
            return nil, errors.New("ERROR: This tree has an equality condition as a branch. Quitting.")
        }
    }

    linearConditions := []Condition{Condition{Text: "(", Type: "scope", Operator: "("}}

    for key, child := range node.Children {
        if key != 0 {
            linearConditions = append(linearConditions, node.Node)
        }

        serializedChild, err := serializeTree(child)

        if err != nil {
            return nil, err
        }

        linearConditions = append(linearConditions, serializedChild...)
    }

    linearConditions = append(linearConditions, Condition{Text: ")", Type: "scope", Operator: ")"})

    return linearConditions, nil
}

func (t *treeNode) attachLeftsAndRights() {
    indexStart := 0
    t.attachLeftsAndRightsRecursively(&indexStart)
}

func (t *treeNode) attachLeftsAndRightsRecursively(index *int) {
    *index++
    t.Left = *index

    for _, child := range t.Children {
        child.attachLeftsAndRightsRecursively(index)
    }

    *index++
    t.Right = *index
}

func (t *treeNode) toMysql() (equalityStr, logicStr string, err error) {
    t.attachLeftsAndRights()

    equalityStr, logicStr = t.toMysqlRecursively()

    if len(equalityStr) == 0 {
        return "", "", errors.New(fmt.Sprintf("Error: equality string was empty, which was unexpected. EqualityStr: %s :: LogicStr: %s", equalityStr, logicStr))
    }

    if len(logicStr) == 0 {
        return "", "", errors.New(fmt.Sprintf("Error: logic string was empty, which was unexpected. EqualityStr: %s :: LogicStr: %s", equalityStr, logicStr))
    }

    equalityStr = equalityStr[:(len(equalityStr)-1)]
    logicStr = logicStr[:(len(logicStr)-1)]

    return equalityStr, logicStr, nil
}

func (t *treeNode) toMysqlRecursively() (equalityStr, logicStr string) {
    var equalityTemp, logicTemp string
    for _, child := range t.Children {
        equalityTemp, logicTemp = child.toMysqlRecursively()
        equalityStr += equalityTemp
        logicStr += logicTemp
    }

    switch (t.Node.Type) {
    case "equality":
        equalityStr += fmt.Sprintf("('%s', '%s', '%s', 'equality', %d, %d),", t.Node.Field, t.Node.Operator, t.Node.Value, t.Left, t.Right)
    case "logic":
        logicStr += fmt.Sprintf("('%s', 'logic', %d, %d),", t.Node.Operator, t.Left, t.Right)
    }

    return equalityStr, logicStr
}

func (t *treeNode) print() string {
    var s string

    for _, child := range t.Children {
        s += child.print()
    }

    return s + " :: " + fmt.Sprintf("%v, %d, %d", t.Node, t.Left, t.Right)
}

func simplifyConditions(conditions []Condition) string {
    var t string

    for k, c := range conditions {
        if k != 0 {
            t += " "
        }

        t += c.Text
    }

    return t
}

func (t *treeNode) getChildrenConditions() []Condition {
    var children []Condition

    for _, child := range t.Children {
        children = append(children, child.Node)
    }

    return children
}

func updateDatabase(equalityStr, logicStr string) {
    db, _ := sql.Open("mysql", "root:@/")
    defer db.Close()

    _, err := db.Query("TRUNCATE TABLE logictree.conditions")
    common.CheckError(err, 2)

    _, err = db.Query("INSERT INTO logictree.conditions (field, operator, value, type, lt, rt) VALUES "+equalityStr)
    common.CheckError(err, 2)
    _, err = db.Query("INSERT INTO logictree.conditions (operator, type, lt, rt) VALUES "+logicStr)
    common.CheckError(err, 2)
}

func (conditionA Condition) matches(conditionB Condition) bool {
    if conditionA.Text != conditionB.Text {
        return false
    }

    if conditionA.Type != conditionB.Type {
        return false
    }

    if conditionA.Field != conditionB.Field {
        return false
    }

    if conditionA.Operator != conditionB.Operator {
        return false
    }

    if conditionA.Value != conditionB.Value {
        return false
    }

    return true
}

<commit_msg>wip - recursive unserialize raw is generatige huge amounts of children<commit_after>package home

import (
    "fmt"
    "encoding/json"
    "net/http"
    "html/template"
    "errors"
    "database/sql"

    _ "github.com/go-sql-driver/mysql"
    "git-go-logictree/app/common"
)

type Condition struct {
    Text, Type, Field, Operator, Value string
}

type treeNode struct {
    Parent *treeNode
    Children []*treeNode
    Node Condition
    Left, Right int
}

type conditionSqlRow struct {
    Field, Operator, Value, Type string
    Left, Right int
}

func GetHomePage(rw http.ResponseWriter, req *http.Request) {
    type Page struct {
        Title string
        Conditions []Condition
    }

    p := Page{
        Title: "home",
        Conditions: getConditions(),
    }

    common.Templates = template.Must(template.ParseFiles("templates/home/home.html", common.LayoutPath))
    err := common.Templates.ExecuteTemplate(rw, "base", p)
    common.CheckError(err, 2)
}

func UpdateConditions(rw http.ResponseWriter, req *http.Request) {
    conditions := req.FormValue("conditions");

    parsedConditions, _ := parseJSON(conditions);
    fmt.Println(parsedConditions)
    treeRoot, err := unserializeFormattedTree(parsedConditions)

    equalityStr, logicStr, err := treeRoot.toMysql()
    if err != nil {
        fmt.Println(err)
        return
    }

    updateDatabase(equalityStr, logicStr)

    GetHomePage(rw, req)
}

func Truncate(rw http.ResponseWriter, req *http.Request) {
    _, err := common.DB.Query("TRUNCATE TABLE logictree.equality")
    common.CheckError(err, 2)

    _, err = common.DB.Query("TRUNCATE TABLE logictree.logic")
    common.CheckError(err, 2)

    GetHomePage(rw, req)
}

func getConditions() []Condition {
    conditions := make([]Condition, 0)

    rows, err := common.DB.Query("SELECT field, operator, value FROM logictree.equality")
    common.CheckError(err, 2)

    var field, operator, value string

    i := 0

    for rows.Next() {
        rows.Scan(&field, &operator, &value)
        common.CheckError(err, 2)

        if i != 0 {
            conditions = append(conditions, Condition{
                Text: "AND",
                Operator: "AND",
                Type: "logic",
            })
        }

        conditions = append(conditions, Condition{
            Text: fmt.Sprintf("%s %s %s", field, operator, value),
            Type: "equality",
            Field: field,
            Operator: operator,
            Value: value,
        })

        i++
    }

    return conditions
}

func parseJSON(conditionsString string) ([]Condition, error) {
    var conditionsSlice []Condition

    err := json.Unmarshal([]byte(conditionsString), &conditionsSlice)
    if err != nil {
        return nil, err
    }

    return conditionsSlice, nil
}

/** Treat conditions like a queue. Rules:
 * If you reach a (, pop the condition, drop down a depth and assign results to root's children
 * If you reach a ), pop the condition, pop back up a depth with the root
 * If you reach a logical condition, pop the condition, assign it as the root's node
 * If you reach an equality condition, pop the condition, assign it as one of the children of the root
 * At the end of the loop, return the root's first child (since we have parans around all conditions we're going to be one level too deep)
**/
func unserializeFormattedTree(conditions []Condition) (*treeNode, error) {
    var root treeNode
    var emptyNode, condition Condition

    key := 0

    for key < len(conditions) {
        // Pop the front item from the slice
        condition = conditions[0]
        conditions = append(conditions[:key], conditions[key+1:]...)

        switch condition.Type {
        case "scope":
            if condition.Operator == "(" {
                children, _ := unserializeFormattedTree(conditions)

                if len(root.Children) == 0 {
                    root.Children = []*treeNode{children}
                } else {
                    root.Children = append(root.Children, children)
                }
            }

            if condition.Operator == ")" {
                if root.Node == emptyNode {
                    return root.Children[0], nil
                } else {
                    return &root, nil
                }
            }
        case "logic":
            root.Node = condition
        case "equality":
            root.Children = append(root.Children, &treeNode{Parent: &root, Node: condition})
        }
    }

    if len(root.Children) > 0 && !root.Node.matches(emptyNode) {
        return &root, nil
    }

    return root.Children[0], nil
}

/**
 * Assumption: data should be ordered by LEFT
 * Recursive steps:
 * If current node is a branch:
 *      Create node with it
 *      Pop item from conditions
 *      Add children recursively
 * If node is a leaf:
 *      Create node with it
 *      Pop item from conditions
 * Return node
**/
func unserializeRawTree(conditions []conditionSqlRow) *treeNode {
    var condition conditionSqlRow
    root := &treeNode{}
    key := 0

    for key < len(conditions) {
        // Pop the front item from the slice
        condition = conditions[0]
        conditions = append(conditions[:key], conditions[key+1:]...)

        (*root).Node = condition.conv()

        if condition.Left == condition.Right-1 {
            (*root).Children = append((*root).Children, unserializeRawTree(conditions))
        }
    }

    return root
}

func serializeTree(node *treeNode) ([]Condition, error) {
    if node.Children == nil || len(node.Children) == 0 {
        // Has no children - should be equality

        if node.Node.Type != "equality" {
            return nil, errors.New("ERROR: This tree has a logic condition as a leaf. Quitting.")
        }

        return []Condition{node.Node}, nil
    } else {
        // Has children - should be logic

        if node.Node.Type != "logic" {
            return nil, errors.New("ERROR: This tree has an equality condition as a branch. Quitting.")
        }
    }

    linearConditions := []Condition{Condition{Text: "(", Type: "scope", Operator: "("}}

    for key, child := range node.Children {
        if key != 0 {
            linearConditions = append(linearConditions, node.Node)
        }

        serializedChild, err := serializeTree(child)

        if err != nil {
            return nil, err
        }

        linearConditions = append(linearConditions, serializedChild...)
    }

    linearConditions = append(linearConditions, Condition{Text: ")", Type: "scope", Operator: ")"})

    return linearConditions, nil
}

func (t *treeNode) attachLeftsAndRights() {
    indexStart := 0
    t.attachLeftsAndRightsRecursively(&indexStart)
}

func (t *treeNode) attachLeftsAndRightsRecursively(index *int) {
    *index++
    t.Left = *index

    for _, child := range t.Children {
        child.attachLeftsAndRightsRecursively(index)
    }

    *index++
    t.Right = *index
}

func (t *treeNode) toMysql() (equalityStr, logicStr string, err error) {
    t.attachLeftsAndRights()

    equalityStr, logicStr = t.toMysqlRecursively()

    if len(equalityStr) == 0 {
        return "", "", errors.New(fmt.Sprintf("Error: equality string was empty, which was unexpected. EqualityStr: %s :: LogicStr: %s", equalityStr, logicStr))
    }

    if len(logicStr) == 0 {
        return "", "", errors.New(fmt.Sprintf("Error: logic string was empty, which was unexpected. EqualityStr: %s :: LogicStr: %s", equalityStr, logicStr))
    }

    equalityStr = equalityStr[:(len(equalityStr)-1)]
    logicStr = logicStr[:(len(logicStr)-1)]

    return equalityStr, logicStr, nil
}

func (t *treeNode) toMysqlRecursively() (equalityStr, logicStr string) {
    var equalityTemp, logicTemp string
    for _, child := range t.Children {
        equalityTemp, logicTemp = child.toMysqlRecursively()
        equalityStr += equalityTemp
        logicStr += logicTemp
    }

    switch (t.Node.Type) {
    case "equality":
        equalityStr += fmt.Sprintf("('%s', '%s', '%s', 'equality', %d, %d),", t.Node.Field, t.Node.Operator, t.Node.Value, t.Left, t.Right)
    case "logic":
        logicStr += fmt.Sprintf("('%s', 'logic', %d, %d),", t.Node.Operator, t.Left, t.Right)
    }

    return equalityStr, logicStr
}

func (t *treeNode) print() string {
    var s string

    if t == nil {
        return ""
    }

    for _, child := range t.Children {
        s += child.print()
    }

    return s + " :: " + fmt.Sprintf("%v, %d, %d", t.Node, t.Left, t.Right)
}

func simplifyConditions(conditions []Condition) string {
    var t string

    for k, c := range conditions {
        if k != 0 {
            t += " "
        }

        t += c.Text
    }

    return t
}

func (t *treeNode) getChildrenConditions() []Condition {
    var children []Condition

    for _, child := range t.Children {
        children = append(children, child.Node)
    }

    return children
}

func updateDatabase(equalityStr, logicStr string) {
    db, _ := sql.Open("mysql", "root:@/")
    defer db.Close()

    _, err := db.Query("TRUNCATE TABLE logictree.conditions")
    common.CheckError(err, 2)

    _, err = db.Query("INSERT INTO logictree.conditions (field, operator, value, type, lt, rt) VALUES "+equalityStr)
    common.CheckError(err, 2)
    _, err = db.Query("INSERT INTO logictree.conditions (operator, type, lt, rt) VALUES "+logicStr)
    common.CheckError(err, 2)
}

func (conditionA Condition) matches(conditionB Condition) bool {
    if conditionA.Text != conditionB.Text {
        return false
    }

    if conditionA.Type != conditionB.Type {
        return false
    }

    if conditionA.Field != conditionB.Field {
        return false
    }

    if conditionA.Operator != conditionB.Operator {
        return false
    }

    if conditionA.Value != conditionB.Value {
        return false
    }

    return true
}

func (condition conditionSqlRow) conv() Condition {
    return Condition{Text: condition.Field + " " + condition.Operator + " " + condition.Value, Type: condition.Type, Field: condition.Field, Operator: condition.Operator, Value: condition.Value}
}

<|endoftext|>


<commit_before>// Copyright 2015 Claudemiro Alves Feitosa Neto. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package gochip8

import (
	"testing"

	"github.com/bmizerany/assert"
)

// dumbSound - Only for testing purpuses
type dumbSound struct{}

func (s dumbSound) Beep() {
	// play beep
}

// newCpuAt creates a new cpu and write the ins to the pc memory location
// Only for testing purpuse
func newCpuAt(ins uint16) cpu {
	s := dumbSound{}
	c := newCpu(s)
	c.memory.WriteWord(uint16(c.pc), ins)

	return c
}

// 00E0 - CLS
func Test00E0(t *testing.T) {

	c := newCpuAt(0x00E0)
	c.gfx.SetPixel(0, 0)
	c.Step()

	for i := 0; i < 64; i++ {
		for j := 0; j < 32; j++ {
			assert.Equal(t, byte(COLOR_BLACK), byte(c.gfx.GetPixel(i, j)))
		}
	}
}

// 00EE - RET
func Test00EE(t *testing.T) {
	c := newCpuAt(0x00EE)

	c.stack.Push(pc(1))
	c.Step()

	assert.Equal(t, pc(1), c.pc)
}

// 1nnn - JP addr
func Test1000(t *testing.T) {
	c := newCpuAt(0x1FFF)
	c.Step()

	assert.Equal(t, pc(0x0FFF), c.pc)
}

// 2nnn - CALL addr
func Test2000(t *testing.T) {
	c := newCpuAt(0x2FFF)
	oldPC := c.pc + 2 // The increment happen before the Step

	c.Step()

	assert.Equal(t, oldPC, c.stack.Pop())
	assert.Equal(t, pc(0x0FFF), c.pc)
}

// 3xkk - SE Vx, byte
func Test3000Equal(t *testing.T) {
	c := newCpuAt(0x31FF)
	c.regs[1] = 0xFF

	oldPC := c.pc + 4
	c.Step()

	assert.Equal(t, oldPC, c.pc)
}

// 3xkk - SE Vx, byte
func Test3000NotEqual(t *testing.T) {
	c := newCpuAt(0x31FF)
	c.regs[1] = 0xAA

	oldPC := c.pc + 2
	c.Step()

	assert.Equal(t, oldPC, c.pc)
}

// 4xkk - SNE Vx, byte
func Test4000Equal(t *testing.T) {
	c := newCpuAt(0x41FF)
	c.regs[1] = 0xFF

	oldPC := c.pc + 2
	c.Step()

	assert.Equal(t, oldPC, c.pc)
}

// 4xkk - SNE Vx, byte
func Test4000NotEqual(t *testing.T) {
	c := newCpuAt(0x41FF)
	c.regs[1] = 0xAA

	oldPC := c.pc + 4
	c.Step()

	assert.Equal(t, oldPC, c.pc)
}

// 5xy0 - SE Vx, Vy
func Test5000Equal(t *testing.T) {
	c := newCpuAt(0x512F)
	c.regs[1] = 0xFF
	c.regs[2] = 0xFF

	oldPC := c.pc + 4
	c.Step()

	assert.Equal(t, oldPC, c.pc)
}

// 5xy0 - SE Vx, Vy
func Test5000NotEqual(t *testing.T) {
	c := newCpuAt(0x512F)
	c.regs[1] = 0xFF
	c.regs[2] = 0xAA

	oldPC := c.pc + 2
	c.Step()

	assert.Equal(t, oldPC, c.pc)
}

// 6xkk - LD Vx, byte
func Test6000(t *testing.T) {
	c := newCpuAt(0x61FF)
	c.Step()
	assert.Equal(t, c.regs[1], byte(0xFF))
}

// 7xkk - ADD Vx, byte
func Test7000(t *testing.T) {
	c := newCpuAt(0x7101)
	c.regs[1] = 1
	c.Step()
	assert.Equal(t, c.regs[1], byte(2))
}
<commit_msg>Added subtract timers test<commit_after>// Copyright 2015 Claudemiro Alves Feitosa Neto. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package gochip8

import (
	"testing"

	"github.com/bmizerany/assert"
)

// dumbSound - Only for testing purpuses
type dumbSound struct {
	BeepCalled bool
}

func (s *dumbSound) Beep() {
	(*s).BeepCalled = true
}

// newCpuAt creates a new cpu and write the ins to the pc memory location
// Only for testing purpuse
func newCpuAt(ins uint16) cpu {
	s := &dumbSound{}
	c := newCpu(s)
	c.memory.WriteWord(uint16(c.pc), ins)

	return c
}

func TestTimers(t *testing.T) {
	s := &dumbSound{}
	c := newCpu(s)
	c.memory.WriteWord(uint16(c.pc), 0x0E0)

	c.dt = 1
	c.st = 1

	assert.Equal(t, s.BeepCalled, false)

	c.Step()

	assert.Equal(t, c.dt, byte(0))
	assert.Equal(t, c.st, byte(0))
	assert.Equal(t, s.BeepCalled, true)
}

// 00E0 - CLS
func Test00E0(t *testing.T) {

	c := newCpuAt(0x00E0)
	c.gfx.SetPixel(0, 0)
	c.Step()

	for i := 0; i < 64; i++ {
		for j := 0; j < 32; j++ {
			assert.Equal(t, byte(COLOR_BLACK), byte(c.gfx.GetPixel(i, j)))
		}
	}
}

// 00EE - RET
func Test00EE(t *testing.T) {
	c := newCpuAt(0x00EE)

	c.stack.Push(pc(1))
	c.Step()

	assert.Equal(t, pc(1), c.pc)
}

// 1nnn - JP addr
func Test1000(t *testing.T) {
	c := newCpuAt(0x1FFF)
	c.Step()

	assert.Equal(t, pc(0x0FFF), c.pc)
}

// 2nnn - CALL addr
func Test2000(t *testing.T) {
	c := newCpuAt(0x2FFF)
	oldPC := c.pc + 2 // The increment happen before the Step

	c.Step()

	assert.Equal(t, oldPC, c.stack.Pop())
	assert.Equal(t, pc(0x0FFF), c.pc)
}

// 3xkk - SE Vx, byte
func Test3000Equal(t *testing.T) {
	c := newCpuAt(0x31FF)
	c.regs[1] = 0xFF

	oldPC := c.pc + 4
	c.Step()

	assert.Equal(t, oldPC, c.pc)
}

// 3xkk - SE Vx, byte
func Test3000NotEqual(t *testing.T) {
	c := newCpuAt(0x31FF)
	c.regs[1] = 0xAA

	oldPC := c.pc + 2
	c.Step()

	assert.Equal(t, oldPC, c.pc)
}

// 4xkk - SNE Vx, byte
func Test4000Equal(t *testing.T) {
	c := newCpuAt(0x41FF)
	c.regs[1] = 0xFF

	oldPC := c.pc + 2
	c.Step()

	assert.Equal(t, oldPC, c.pc)
}

// 4xkk - SNE Vx, byte
func Test4000NotEqual(t *testing.T) {
	c := newCpuAt(0x41FF)
	c.regs[1] = 0xAA

	oldPC := c.pc + 4
	c.Step()

	assert.Equal(t, oldPC, c.pc)
}

// 5xy0 - SE Vx, Vy
func Test5000Equal(t *testing.T) {
	c := newCpuAt(0x512F)
	c.regs[1] = 0xFF
	c.regs[2] = 0xFF

	oldPC := c.pc + 4
	c.Step()

	assert.Equal(t, oldPC, c.pc)
}

// 5xy0 - SE Vx, Vy
func Test5000NotEqual(t *testing.T) {
	c := newCpuAt(0x512F)
	c.regs[1] = 0xFF
	c.regs[2] = 0xAA

	oldPC := c.pc + 2
	c.Step()

	assert.Equal(t, oldPC, c.pc)
}

// 6xkk - LD Vx, byte
func Test6000(t *testing.T) {
	c := newCpuAt(0x61FF)
	c.Step()
	assert.Equal(t, c.regs[1], byte(0xFF))
}

// 7xkk - ADD Vx, byte
func Test7000(t *testing.T) {
	c := newCpuAt(0x7101)
	c.regs[1] = 1
	c.Step()
	assert.Equal(t, c.regs[1], byte(2))
}
<|endoftext|>


<commit_before>// Application to wrap Python scripts downloaded from Gerrit.
//
// The Python scripts run benchmarks and emit Perf compatible JSON files, which
// are then uploaded to Google Cloud Storage.
//
// Every Python script must support the following flags:
//
//     -b BROWSER, --browser=BROWSER
//                           The browser to use to run MotionMark in.
//     -s SUITE, --suite=SUITE
//                           Run only the specified suite of tests.
//     -e EXECUTABLE, --executable-path=EXECUTABLE
//                           Path to the executable to the driver binary.
//     -a ARGUMENTS, --arguments=ARGUMENTS
//                           Extra arguments to pass to the browser.
//     -g GITHASH, --githash=GITHASH
//                           A git-hash associated with this run.
//     -o OUTPUT, --output=OUTPUT
//                           Path to the output json file.
//
package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"

	"cloud.google.com/go/storage"
	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/executil"
	"go.skia.org/infra/go/gcs"
	"go.skia.org/infra/go/gcs/gcsclient"
	"go.skia.org/infra/go/git"
	"go.skia.org/infra/go/gitauth"
	"go.skia.org/infra/go/now"
	"go.skia.org/infra/go/skerr"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
	"golang.org/x/oauth2"
	"google.golang.org/api/option"
)

const (
	// The Google Cloud Storage Bucket to write the results to.
	bucket = "chrome-comp-ui-perf-skia"

	// The path in the bucket where Perf results should be written.
	bucketPath = "ingest"

	// The repo that has commits associated with runs of the cron job.
	repo = "https://skia.googlesource.com/perf-compui"

	python = "/Library/Frameworks/Python.framework/Versions/3.9/bin/python3"

	benchmarkTimeout = 2 * time.Hour
)

var (
	// Key can be changed via -ldflags.
	Key = "base64 encoded service account key JSON goes here."

	// Version can be changed via -ldflags.
	Version = "unsupplied"
)

// benchmark represents a single benchmark configuration.
type benchmark struct {
	// The checkout URL of the git repo that contains the scripts to run
	repoURL string

	// The directories in the git repo that need to be checked out.
	checkoutPaths []string

	// The full name of the script to run in the git repo relative to the root
	// of the checkout.
	scriptName string

	// Flags to pass to the Python script.
	flags []string
}

// All the various benchmarks we run.
var benchmarks = map[string]benchmark{
	// We always run the canary to validate that the whole pipeline works even
	// if the "real" benchmark scripts start to fail.
	"canary": {
		repoURL:       "https://skia.googlesource.com/buildbot",
		checkoutPaths: []string{"comp-ui"},
		scriptName:    "comp-ui/benchmark-mock.py",
		flags: []string{
			"--browser", "mock",
		},
	},
	"chrome-motionmark": {
		repoURL:       "https://chromium.googlesource.com/chromium/src",
		checkoutPaths: []string{"tools/browserbench-webdriver"},
		scriptName:    "tools/browserbench-webdriver/motionmark.py",
		flags: []string{
			"--browser", "chrome",
			"--executable-path", filepath.Join(os.Getenv("HOME"), "chromedriver"),
		},
	},
	"chrome-jetstream": {
		repoURL:       "https://chromium.googlesource.com/chromium/src",
		checkoutPaths: []string{"tools/browserbench-webdriver"},
		scriptName:    "tools/browserbench-webdriver/jetstream.py",
		flags: []string{
			"--browser", "chrome",
			"--executable-path", filepath.Join(os.Getenv("HOME"), "chromedriver"),
		},
	},
	"chrome-speedometer": {
		repoURL:       "https://chromium.googlesource.com/chromium/src",
		checkoutPaths: []string{"tools/browserbench-webdriver"},
		scriptName:    "tools/browserbench-webdriver/speedometer.py",
		flags: []string{
			"--browser", "chrome",
			"--executable-path", filepath.Join(os.Getenv("HOME"), "chromedriver"),
		},
	},
}

func main() {
	flagSet := flag.NewFlagSet("", flag.ExitOnError)
	flagSet.Usage = func() {
		fmt.Printf("Usage: %s <flags>\n\n", os.Args[0])
		flagSet.PrintDefaults()
	}

	local := flagSet.Bool("local", false, "Running locally if true. As opposed to in production.")

	common.InitWithMust(
		"comp-ui-cron-job",
		common.CloudLogging(local, "skia-public"),
		common.FlagSetOpt(flagSet),
	)
	sklog.Infof("Version: %s", Version)

	ctx := context.Background()

	ts, err := auth.NewTokenSourceFromKeyString(ctx, *local, Key, storage.ScopeFullControl, auth.ScopeUserinfoEmail, auth.ScopeGerrit)
	if err != nil {
		sklog.Fatal(err)
	}

	gcsClient, err := getGCSClient(ctx, ts)
	if err != nil {
		sklog.Fatal(err)
	}

	workDir, err := os.MkdirTemp("", "comp-ui-cron-job")
	if err != nil {
		sklog.Fatal(err)
	}
	defer func() {
		err := os.RemoveAll(workDir)
		if err != nil {
			sklog.Error(err)
		}
	}()

	// We presume that if running locally that you've already authenticated to
	// Gerrit, otherwise write out a git cookie that enables R/W access to the
	// git repo.
	//
	// Authenticate to Gerrit since the perf-compui repo is private.
	if !*local {
		sklog.Info("Configuring git auth.")
		if _, err := gitauth.New(ts, "/tmp/git-cookie", true, ""); err != nil {
			sklog.Fatal(err)
		}
	}

	sklog.Info("Getting githash.")
	gitHash, err := getGitHash(ctx, workDir)
	if err != nil {
		sklog.Fatal(err)
	}

	for benchmarkName, config := range benchmarks {
		outputFilename, err := runSingleBenchmark(ctx, benchmarkName, config, gitHash, workDir)
		if err != nil {
			sklog.Errorf("Failed to run benchmark %q: %s", benchmarkName, err)
			continue
		}

		err = uploadResultsFile(ctx, gcsClient, benchmarkName, outputFilename)
		if err != nil {
			sklog.Errorf("Failed to upload benchmark results %q: %s", benchmarkName, err)
		}
	}
	sklog.Flush()
}

// getGitHash returns the git hash of the last commit to the perf-compui repo,
// which only gets a single commit per day.
func getGitHash(ctx context.Context, workDir string) (string, error) {
	// Find the githash for 'today' from https://skia.googlesource.com/perf-compui.
	g, err := git.NewRepo(ctx, repo, filepath.Join(workDir, "perf-compui"))
	if err != nil {
		return "", skerr.Wrap(err)
	}

	hashes, err := g.RevList(ctx, "HEAD", "-n1")
	if err != nil {
		return "", skerr.Wrap(err)
	}
	return hashes[0], nil
}

func getGCSClient(ctx context.Context, ts oauth2.TokenSource) (*gcsclient.StorageClient, error) {
	storageClient, err := storage.NewClient(ctx, option.WithTokenSource(ts))
	if err != nil {
		return nil, skerr.Wrap(err)
	}

	gcsClient := gcsclient.New(storageClient, bucket)
	return gcsClient, nil

}

func runSingleBenchmark(ctx context.Context, benchmarkName string, config benchmark, gitHash string, workDir string) (string, error) {
	sklog.Infof("runSingleBenchMark - benchmarkName: %q  url: %q  gitHash: %q workDir: %q", benchmarkName, config.scriptName, gitHash, workDir)

	gitCheckoutDir, err := checkoutPythonScript(ctx, config, workDir, benchmarkName)
	if err != nil {
		return "", skerr.Wrap(err)
	}

	// Compute the filenames we will use.
	scriptFilename := filepath.Join(gitCheckoutDir, config.scriptName)
	outputDirectory := filepath.Join(workDir, benchmarkName)
	outputFilename := filepath.Join(outputDirectory, "results.json")

	// Create output directory.
	err = os.MkdirAll(outputDirectory, 0755)
	if err != nil {
		return "", skerr.Wrap(err)
	}

	// Compute the full set of args to pass to the python script.
	flags := append([]string{}, config.flags...)
	flags = append(flags, "--githash", gitHash)
	flags = append(flags, "--output", outputFilename)
	args := append([]string{scriptFilename}, flags...)

	// Run the script.
	err = runBenchMarkScript(ctx, args, workDir)
	if err != nil {
		return "", skerr.Wrap(err)
	}
	return outputFilename, nil
}

func runBenchMarkScript(ctx context.Context, args []string, workDir string) error {
	ctx, cancel := context.WithTimeout(ctx, benchmarkTimeout)
	defer cancel()

	return runCmdLogOutput(ctx, workDir, python, args...)
}

func uploadResultsFile(ctx context.Context, gcsClient gcs.GCSClient, benchmarkName string, outputFilename string) error {
	// GCS paths always use "/" separators.
	destinationPath := path.Join(bucketPath, computeUploadPathFromTime(ctx), benchmarkName, "results.json")
	sklog.Infof("Upload to %q", destinationPath)
	w := gcsClient.FileWriter(ctx, destinationPath, gcs.FileWriteOptions{
		ContentEncoding: "application/json",
	})

	// Copy the output file up to GCS.
	err := util.WithReadFile(outputFilename, func(f io.Reader) error {
		_, err := io.Copy(w, f)
		return skerr.Wrap(err)
	})
	if err != nil {
		return skerr.Wrap(err)
	}

	return skerr.Wrap(w.Close())
}

// computeUploadPathFromTime returns the date-time portion of the GCS path that
// Perf expects uploads to use, see
// https://skia.googlesource.com/buildbot/+/refs/heads/main/perf/FORMAT.md.
func computeUploadPathFromTime(ctx context.Context) string {
	return now.Now(ctx).UTC().Format("2006/01/02/15")
}

// checkoutPythonScript checks out a sparse checkout of the specified directories
// into workDir.
func checkoutPythonScript(ctx context.Context, config benchmark, workDir string, benchmarkName string) (string, error) {
	dest := filepath.Join(workDir, "git", benchmarkName)
	return dest, newSparseCheckout(ctx, workDir, config.repoURL, dest, config.checkoutPaths)
}

// runCmdLogOutput runs a command using executil.CommandContext and logs any output
// to sklog.Info().
func runCmdLogOutput(ctx context.Context, cwd string, cmd string, args ...string) error {
	cc := executil.CommandContext(ctx, cmd, args...)
	cc.Dir = cwd
	output, err := cc.CombinedOutput()
	lines := strings.Split(string(output), "\n")
	for _, line := range lines {
		sklog.Info(line)
	}
	if err != nil {
		return skerr.Wrap(err)
	}
	return nil
}

// newSparseCheckout does a sparse checkout of the given 'directories'.
func newSparseCheckout(ctx context.Context, workDir, repoURL, dest string, directories []string) error {
	if err := runCmdLogOutput(ctx, workDir, "git", "clone", "--depth", "1", "--filter=blob:none", "--sparse", repoURL, dest); err != nil {
		return skerr.Wrapf(err, "Failed to clone.")
	}
	if err := runCmdLogOutput(ctx, dest, "git", "sparse-checkout", "init", "--cone"); err != nil {
		return skerr.Wrapf(err, "Failed to init sparse checkout.")
	}

	args := []string{"sparse-checkout", "set"}
	args = append(args, directories...)
	if err := runCmdLogOutput(ctx, dest, "git", args...); err != nil {
		return skerr.Wrapf(err, "Failed to do a sparse checkout.")
	}

	return nil
}
<commit_msg>Add Safari speedometer stable tests.<commit_after>// Application to wrap Python scripts downloaded from Gerrit.
//
// The Python scripts run benchmarks and emit Perf compatible JSON files, which
// are then uploaded to Google Cloud Storage.
//
// Every Python script must support the following flags:
//
//     -b BROWSER, --browser=BROWSER
//                           The browser to use to run MotionMark in.
//     -s SUITE, --suite=SUITE
//                           Run only the specified suite of tests.
//     -e EXECUTABLE, --executable-path=EXECUTABLE
//                           Path to the executable to the driver binary.
//     -a ARGUMENTS, --arguments=ARGUMENTS
//                           Extra arguments to pass to the browser.
//     -g GITHASH, --githash=GITHASH
//                           A git-hash associated with this run.
//     -o OUTPUT, --output=OUTPUT
//                           Path to the output json file.
//     --extra-keys=KEY_VALUE_PAIRS
//                           A comma separated list of key,value pairs.
//
package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"

	"cloud.google.com/go/storage"
	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/executil"
	"go.skia.org/infra/go/gcs"
	"go.skia.org/infra/go/gcs/gcsclient"
	"go.skia.org/infra/go/git"
	"go.skia.org/infra/go/gitauth"
	"go.skia.org/infra/go/now"
	"go.skia.org/infra/go/skerr"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
	"golang.org/x/oauth2"
	"google.golang.org/api/option"
)

const (
	// The Google Cloud Storage Bucket to write the results to.
	bucket = "chrome-comp-ui-perf-skia"

	// The path in the bucket where Perf results should be written.
	bucketPath = "ingest"

	// The repo that has commits associated with runs of the cron job.
	repo = "https://skia.googlesource.com/perf-compui"

	python = "/Library/Frameworks/Python.framework/Versions/3.9/bin/python3"

	benchmarkTimeout = 2 * time.Hour
)

var (
	// Key can be changed via -ldflags.
	Key = "base64 encoded service account key JSON goes here."

	// Version can be changed via -ldflags.
	Version = "unsupplied"
)

// benchmark represents a single benchmark configuration.
type benchmark struct {
	// The checkout URL of the git repo that contains the scripts to run
	repoURL string

	// The directories in the git repo that need to be checked out.
	checkoutPaths []string

	// The full name of the script to run in the git repo relative to the root
	// of the checkout.
	scriptName string

	// Flags to pass to the Python script.
	flags []string
}

// All the various benchmarks we run.
var benchmarks = map[string]benchmark{
	// We always run the canary to validate that the whole pipeline works even
	// if the "real" benchmark scripts start to fail.
	"canary": {
		repoURL:       "https://skia.googlesource.com/buildbot",
		checkoutPaths: []string{"comp-ui"},
		scriptName:    "comp-ui/benchmark-mock.py",
		flags: []string{
			"--browser", "mock",
		},
	},
	"chrome-motionmark": {
		repoURL:       "https://chromium.googlesource.com/chromium/src",
		checkoutPaths: []string{"tools/browserbench-webdriver"},
		scriptName:    "tools/browserbench-webdriver/motionmark.py",
		flags: []string{
			"--browser", "chrome",
			"--extra-keys", "channel,stable",
			"--executable-path", filepath.Join(os.Getenv("HOME"), "chromedriver"),
		},
	},
	"chrome-jetstream": {
		repoURL:       "https://chromium.googlesource.com/chromium/src",
		checkoutPaths: []string{"tools/browserbench-webdriver"},
		scriptName:    "tools/browserbench-webdriver/jetstream.py",
		flags: []string{
			"--browser", "chrome",
			"--extra-keys", "channel,stable",
			"--executable-path", filepath.Join(os.Getenv("HOME"), "chromedriver"),
		},
	},
	"chrome-speedometer": {
		repoURL:       "https://chromium.googlesource.com/chromium/src",
		checkoutPaths: []string{"tools/browserbench-webdriver"},
		scriptName:    "tools/browserbench-webdriver/speedometer.py",
		flags: []string{
			"--browser", "chrome",
			"--extra-keys", "channel,stable",
			"--executable-path", filepath.Join(os.Getenv("HOME"), "chromedriver"),
		},
	},
	"safari-speedometer": {
		repoURL:       "https://chromium.googlesource.com/chromium/src",
		checkoutPaths: []string{"tools/browserbench-webdriver"},
		scriptName:    "tools/browserbench-webdriver/speedometer.py",
		flags: []string{
			"--browser", "safari",
			"--extra-keys", "channel,stable",
		},
	},
}

func main() {
	flagSet := flag.NewFlagSet("", flag.ExitOnError)
	flagSet.Usage = func() {
		fmt.Printf("Usage: %s <flags>\n\n", os.Args[0])
		flagSet.PrintDefaults()
	}

	local := flagSet.Bool("local", false, "Running locally if true. As opposed to in production.")

	common.InitWithMust(
		"comp-ui-cron-job",
		common.CloudLogging(local, "skia-public"),
		common.FlagSetOpt(flagSet),
	)
	sklog.Infof("Version: %s", Version)

	ctx := context.Background()

	ts, err := auth.NewTokenSourceFromKeyString(ctx, *local, Key, storage.ScopeFullControl, auth.ScopeUserinfoEmail, auth.ScopeGerrit)
	if err != nil {
		sklog.Fatal(err)
	}

	gcsClient, err := getGCSClient(ctx, ts)
	if err != nil {
		sklog.Fatal(err)
	}

	workDir, err := os.MkdirTemp("", "comp-ui-cron-job")
	if err != nil {
		sklog.Fatal(err)
	}
	defer func() {
		err := os.RemoveAll(workDir)
		if err != nil {
			sklog.Error(err)
		}
	}()

	// We presume that if running locally that you've already authenticated to
	// Gerrit, otherwise write out a git cookie that enables R/W access to the
	// git repo.
	//
	// Authenticate to Gerrit since the perf-compui repo is private.
	if !*local {
		sklog.Info("Configuring git auth.")
		if _, err := gitauth.New(ts, "/tmp/git-cookie", true, ""); err != nil {
			sklog.Fatal(err)
		}
	}

	sklog.Info("Getting githash.")
	gitHash, err := getGitHash(ctx, workDir)
	if err != nil {
		sklog.Fatal(err)
	}

	for benchmarkName, config := range benchmarks {
		outputFilename, err := runSingleBenchmark(ctx, benchmarkName, config, gitHash, workDir)
		if err != nil {
			sklog.Errorf("Failed to run benchmark %q: %s", benchmarkName, err)
			continue
		}

		err = uploadResultsFile(ctx, gcsClient, benchmarkName, outputFilename)
		if err != nil {
			sklog.Errorf("Failed to upload benchmark results %q: %s", benchmarkName, err)
		}
	}
	sklog.Flush()
}

// getGitHash returns the git hash of the last commit to the perf-compui repo,
// which only gets a single commit per day.
func getGitHash(ctx context.Context, workDir string) (string, error) {
	// Find the githash for 'today' from https://skia.googlesource.com/perf-compui.
	g, err := git.NewRepo(ctx, repo, filepath.Join(workDir, "perf-compui"))
	if err != nil {
		return "", skerr.Wrap(err)
	}

	hashes, err := g.RevList(ctx, "HEAD", "-n1")
	if err != nil {
		return "", skerr.Wrap(err)
	}
	return hashes[0], nil
}

func getGCSClient(ctx context.Context, ts oauth2.TokenSource) (*gcsclient.StorageClient, error) {
	storageClient, err := storage.NewClient(ctx, option.WithTokenSource(ts))
	if err != nil {
		return nil, skerr.Wrap(err)
	}

	gcsClient := gcsclient.New(storageClient, bucket)
	return gcsClient, nil

}

func runSingleBenchmark(ctx context.Context, benchmarkName string, config benchmark, gitHash string, workDir string) (string, error) {
	sklog.Infof("runSingleBenchMark - benchmarkName: %q  url: %q  gitHash: %q workDir: %q", benchmarkName, config.scriptName, gitHash, workDir)

	gitCheckoutDir, err := checkoutPythonScript(ctx, config, workDir, benchmarkName)
	if err != nil {
		return "", skerr.Wrap(err)
	}

	// Compute the filenames we will use.
	scriptFilename := filepath.Join(gitCheckoutDir, config.scriptName)
	outputDirectory := filepath.Join(workDir, benchmarkName)
	outputFilename := filepath.Join(outputDirectory, "results.json")

	// Create output directory.
	err = os.MkdirAll(outputDirectory, 0755)
	if err != nil {
		return "", skerr.Wrap(err)
	}

	// Compute the full set of args to pass to the python script.
	flags := append([]string{}, config.flags...)
	flags = append(flags, "--githash", gitHash)
	flags = append(flags, "--output", outputFilename)
	args := append([]string{scriptFilename}, flags...)

	sklog.Infof("Running: %q", args)

	// Run the script.
	err = runBenchMarkScript(ctx, args, workDir)
	if err != nil {
		return "", skerr.Wrap(err)
	}
	return outputFilename, nil
}

func runBenchMarkScript(ctx context.Context, args []string, workDir string) error {
	ctx, cancel := context.WithTimeout(ctx, benchmarkTimeout)
	defer cancel()

	return runCmdLogOutput(ctx, workDir, python, args...)
}

func uploadResultsFile(ctx context.Context, gcsClient gcs.GCSClient, benchmarkName string, outputFilename string) error {
	// GCS paths always use "/" separators.
	destinationPath := path.Join(bucketPath, computeUploadPathFromTime(ctx), benchmarkName, "results.json")
	sklog.Infof("Upload to %q", destinationPath)
	w := gcsClient.FileWriter(ctx, destinationPath, gcs.FileWriteOptions{
		ContentEncoding: "application/json",
	})

	// Copy the output file up to GCS.
	err := util.WithReadFile(outputFilename, func(f io.Reader) error {
		_, err := io.Copy(w, f)
		return skerr.Wrap(err)
	})
	if err != nil {
		return skerr.Wrap(err)
	}

	return skerr.Wrap(w.Close())
}

// computeUploadPathFromTime returns the date-time portion of the GCS path that
// Perf expects uploads to use, see
// https://skia.googlesource.com/buildbot/+/refs/heads/main/perf/FORMAT.md.
func computeUploadPathFromTime(ctx context.Context) string {
	return now.Now(ctx).UTC().Format("2006/01/02/15")
}

// checkoutPythonScript checks out a sparse checkout of the specified directories
// into workDir.
func checkoutPythonScript(ctx context.Context, config benchmark, workDir string, benchmarkName string) (string, error) {
	dest := filepath.Join(workDir, "git", benchmarkName)
	return dest, newSparseCheckout(ctx, workDir, config.repoURL, dest, config.checkoutPaths)
}

// runCmdLogOutput runs a command using executil.CommandContext and logs any output
// to sklog.Info().
func runCmdLogOutput(ctx context.Context, cwd string, cmd string, args ...string) error {
	cc := executil.CommandContext(ctx, cmd, args...)
	cc.Dir = cwd
	output, err := cc.CombinedOutput()
	lines := strings.Split(string(output), "\n")
	for _, line := range lines {
		sklog.Info(line)
	}
	if err != nil {
		return skerr.Wrap(err)
	}
	return nil
}

// newSparseCheckout does a sparse checkout of the given 'directories'.
func newSparseCheckout(ctx context.Context, workDir, repoURL, dest string, directories []string) error {
	if err := runCmdLogOutput(ctx, workDir, "git", "clone", "--depth", "1", "--filter=blob:none", "--sparse", repoURL, dest); err != nil {
		return skerr.Wrapf(err, "Failed to clone.")
	}
	if err := runCmdLogOutput(ctx, dest, "git", "sparse-checkout", "init", "--cone"); err != nil {
		return skerr.Wrapf(err, "Failed to init sparse checkout.")
	}

	args := []string{"sparse-checkout", "set"}
	args = append(args, directories...)
	if err := runCmdLogOutput(ctx, dest, "git", args...); err != nil {
		return skerr.Wrapf(err, "Failed to do a sparse checkout.")
	}

	return nil
}
<|endoftext|>