<commit_before>// Copyright 2019 Google LLC
//
// 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
//
//     https://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.

// [START functions_firebase_reactive]

// Package upper contains a Firestore Cloud Function.
package upper

import (
	"context"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"cloud.google.com/go/firestore"
	firebase "firebase.google.com/go"
)

// FirestoreEvent is the payload of a Firestore event.
type FirestoreEvent struct {
	OldValue   FirestoreValue `json:"oldValue"`
	Value      FirestoreValue `json:"value"`
	UpdateMask struct {
		FieldPaths []string `json:"fieldPaths"`
	} `json:"updateMask"`
}

// FirestoreValue holds Firestore fields.
type FirestoreValue struct {
	CreateTime time.Time `json:"createTime"`
	// Fields is the data for this value. The type depends on the format of your
	// database. Log an interface{} value and inspect the result to see a JSON
	// representation of your database fields.
	Fields     MyData    `json:"fields"`
	Name       string    `json:"name"`
	UpdateTime time.Time `json:"updateTime"`
}

// MyData represents a value from Firestore. The type definition depends on the
// format of your database.
type MyData struct {
	Original struct {
		StringValue string `json:"stringValue"`
	} `json:"original"`
}

// GCLOUD_PROJECT is automatically set by the Cloud Functions runtime.
var projectID = os.Getenv("GCLOUD_PROJECT")

// client is a Firestore client, reused between function invocations.
var client *firestore.Client

func init() {
	// Use the application default credentials.
	conf := &firebase.Config{ProjectID: projectID}

	// Use context.Background() because the app/client should persist across
	// invocations.
	ctx := context.Background()

	app, err := firebase.NewApp(ctx, conf)
	if err != nil {
		log.Fatalf("firebase.NewApp: %v", err)
	}

	client, err = app.Firestore(ctx)
	if err != nil {
		log.Fatalf("app.Firestore: %v", err)
	}
}

// MakeUpperCase is triggered by a change to a Firestore document. It updates
// the `original` value of the document to upper case.
func MakeUpperCase(ctx context.Context, e FirestoreEvent) error {
	fullPath := strings.Split(e.Value.Name, "/documents/")[1]
	pathParts := strings.Split(fullPath, "/")
	collection := pathParts[0]
	doc := strings.Join(pathParts[1:], "/")

	curValue := e.Value.Fields.Original.StringValue
	newValue := strings.ToUpper(curValue)
	if curValue == newValue {
		log.Printf("%q is already upper case: skipping", curValue)
		return nil
	}
	log.Printf("Replacing value: %q -> %q", curValue, newValue)

	data := map[string]string{"original": newValue}
	_, err := client.Collection(collection).Doc(doc).Set(ctx, data)
	if err != nil {
		return fmt.Errorf("Set: %v", err)
	}
	return nil
}

// [END functions_firebase_reactive]
<commit_msg> functions: use GOOGLE_CLOUD_PROJECT, not GCLOUD_PROJECT (#1125)<commit_after>// Copyright 2019 Google LLC
//
// 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
//
//     https://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.

// [START functions_firebase_reactive]

// Package upper contains a Firestore Cloud Function.
package upper

import (
	"context"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"cloud.google.com/go/firestore"
	firebase "firebase.google.com/go"
)

// FirestoreEvent is the payload of a Firestore event.
type FirestoreEvent struct {
	OldValue   FirestoreValue `json:"oldValue"`
	Value      FirestoreValue `json:"value"`
	UpdateMask struct {
		FieldPaths []string `json:"fieldPaths"`
	} `json:"updateMask"`
}

// FirestoreValue holds Firestore fields.
type FirestoreValue struct {
	CreateTime time.Time `json:"createTime"`
	// Fields is the data for this value. The type depends on the format of your
	// database. Log an interface{} value and inspect the result to see a JSON
	// representation of your database fields.
	Fields     MyData    `json:"fields"`
	Name       string    `json:"name"`
	UpdateTime time.Time `json:"updateTime"`
}

// MyData represents a value from Firestore. The type definition depends on the
// format of your database.
type MyData struct {
	Original struct {
		StringValue string `json:"stringValue"`
	} `json:"original"`
}

// GOOGLE_CLOUD_PROJECT is automatically set by the Cloud Functions runtime.
var projectID = os.Getenv("GOOGLE_CLOUD_PROJECT")

// client is a Firestore client, reused between function invocations.
var client *firestore.Client

func init() {
	// Use the application default credentials.
	conf := &firebase.Config{ProjectID: projectID}

	// Use context.Background() because the app/client should persist across
	// invocations.
	ctx := context.Background()

	app, err := firebase.NewApp(ctx, conf)
	if err != nil {
		log.Fatalf("firebase.NewApp: %v", err)
	}

	client, err = app.Firestore(ctx)
	if err != nil {
		log.Fatalf("app.Firestore: %v", err)
	}
}

// MakeUpperCase is triggered by a change to a Firestore document. It updates
// the `original` value of the document to upper case.
func MakeUpperCase(ctx context.Context, e FirestoreEvent) error {
	fullPath := strings.Split(e.Value.Name, "/documents/")[1]
	pathParts := strings.Split(fullPath, "/")
	collection := pathParts[0]
	doc := strings.Join(pathParts[1:], "/")

	curValue := e.Value.Fields.Original.StringValue
	newValue := strings.ToUpper(curValue)
	if curValue == newValue {
		log.Printf("%q is already upper case: skipping", curValue)
		return nil
	}
	log.Printf("Replacing value: %q -> %q", curValue, newValue)

	data := map[string]string{"original": newValue}
	_, err := client.Collection(collection).Doc(doc).Set(ctx, data)
	if err != nil {
		return fmt.Errorf("Set: %v", err)
	}
	return nil
}

// [END functions_firebase_reactive]
<|endoftext|>


<commit_before>// Copyright 2018 gf Author(https://gitee.com/johng/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://gitee.com/johng/gf.
// pprof封装.

package ghttp

import (
    "strings"
    runpprof "runtime/pprof"
    netpprof "net/http/pprof"
    "gitee.com/johng/gf/g/os/gview"
)

// 用于pprof的对象
type utilPprof struct {}

func (p *utilPprof) Index(r *Request) {
    profiles := runpprof.Profiles()
    action   := r.Get("action")
    data     := map[string]interface{}{
        "uri"      : strings.TrimRight(r.URL.Path, "/") + "/",
        "profiles" : profiles,
    }
    if len(action) == 0 {
        buffer, _ := gview.ParseContent(`

                <title>gf ghttp pprof</title>

            {{$uri := .uri}}

                profiles:

                    {{range .profiles}}
{{.Count}}{{.Name}}{{end}}

full goroutine stack dump

            `, data)
        r.Response.Write(buffer)
        return
    }
    for _, p := range profiles {
        if p.Name() == action {
            p.WriteTo(r.Response.Writer, r.GetRequestInt("debug"))
            break
        }
    }
}

func (p *utilPprof) Cmdline(r *Request) {
    netpprof.Cmdline(r.Response.Writer, &r.Request)
}

func (p *utilPprof) Profile(r *Request) {
    netpprof.Profile(r.Response.Writer, &r.Request)
}

func (p *utilPprof) Symbol(r *Request) {
    netpprof.Symbol(r.Response.Writer, &r.Request)
}

func (p *utilPprof) Trace(r *Request) {
    netpprof.Trace(r.Response.Writer, &r.Request)
}

// 开启pprof支持
func (s *Server) EnablePprof(pattern...string) {
    p := "/debug/pprof"
    if len(pattern) > 0 {
        p = pattern[0]
    }
    up := &utilPprof{}
    _, _, uri, _ := s.parsePattern(p)
    uri = strings.TrimRight(uri, "/")
    s.BindHandler(uri + "/*action", up.Index)
    s.BindHandler(uri + "/cmdline", up.Cmdline)
    s.BindHandler(uri + "/profile", up.Profile)
    s.BindHandler(uri + "/symbol",  up.Symbol)
    s.BindHandler(uri + "/trace",   up.Trace)
}<commit_msg>ghttp updates<commit_after>// Copyright 2018 gf Author(https://gitee.com/johng/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://gitee.com/johng/gf.
// pprof封装.

package ghttp

import (
    "strings"
    runpprof "runtime/pprof"
    netpprof "net/http/pprof"
    "gitee.com/johng/gf/g/os/gview"
)

// 用于pprof的对象
type utilPprof struct {}

func (p *utilPprof) Index(r *Request) {
    profiles := runpprof.Profiles()
    action   := r.Get("action")
    data     := map[string]interface{}{
        "uri"      : strings.TrimRight(r.URL.Path, "/") + "/",
        "profiles" : profiles,
    }
    if len(action) == 0 {
        buffer, _ := gview.ParseContent(`

                <title>gf ghttp pprof</title>

            {{$uri := .uri}}

                profiles:

                    {{range .profiles}}
{{.Count}}{{.Name}}{{end}}

full goroutine stack dump

            `, data)
        r.Response.Write(buffer)
        return
    }
    for _, p := range profiles {
        if p.Name() == action {
            p.WriteTo(r.Response.Writer, r.GetRequestInt("debug"))
            break
        }
    }
}

func (p *utilPprof) Cmdline(r *Request) {
    netpprof.Cmdline(r.Response.Writer, r.Request)
}

func (p *utilPprof) Profile(r *Request) {
    netpprof.Profile(r.Response.Writer, r.Request)
}

func (p *utilPprof) Symbol(r *Request) {
    netpprof.Symbol(r.Response.Writer, r.Request)
}

func (p *utilPprof) Trace(r *Request) {
    netpprof.Trace(r.Response.Writer, r.Request)
}

// 开启pprof支持
func (s *Server) EnablePprof(pattern...string) {
    p := "/debug/pprof"
    if len(pattern) > 0 {
        p = pattern[0]
    }
    up := &utilPprof{}
    _, _, uri, _ := s.parsePattern(p)
    uri = strings.TrimRight(uri, "/")
    s.BindHandler(uri + "/*action", up.Index)
    s.BindHandler(uri + "/cmdline", up.Cmdline)
    s.BindHandler(uri + "/profile", up.Profile)
    s.BindHandler(uri + "/symbol",  up.Symbol)
    s.BindHandler(uri + "/trace",   up.Trace)
}<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	restful "github.com/emicklei/go-restful"
	"github.com/emicklei/go-restful/swagger"
	"github.com/garyburd/redigo/redis"
	"github.com/the42/ogdat/database"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"time"
)

const AppID = "5bcbfc24-8e7e-4105-99c4-dd47e7e5094a"
const watcherappid = "a6545f8f-e0c9-4917-83c7-3e47bd1e0247"

var logger *log.Logger

type analyser struct {
	dbcon analyserdb
	pool  *redis.Pool
}

func NewAnalyser(dbcon *sql.DB, pool *redis.Pool) *analyser {
	analyser := &analyser{dbcon: analyserdb{DBConn: database.DBConn{Appid: AppID, DBer: dbcon}}, pool: pool}
	return analyser
}

func isonlyweb() bool {
	boolval, err := strconv.ParseBool(os.Getenv("ONLYWEB"))
	if err != nil {
		return false
	}
	return boolval
}

func getredisconnect() string {
	const redisurl = "REDISCLOUD_URL"
	const redisdb = "ANALYSER_REDISDB"

	return os.Getenv(redisurl) + "/" + os.Getenv(redisdb)
}

func portbinding() string {
	if port := os.Getenv("PORT"); port != "" {
		return port
	}
	return "8080"
}

func discoveryhost() string {
	if host := os.Getenv("DISCOVERYURL"); host != "" {
		return host
	}
	return "http://localhost"
}

func apibasepath() string {
	return os.Getenv("APIBASEPATH")
}

func getheartbeatinterval() int {

	if i, err := strconv.Atoi(os.Getenv("HEARTBEAT_INTERVAL")); err == nil {
		return i
	}
	return 60 // Minutes
}

func heartbeat(interval int) chan bool {
	retchan := make(chan bool)
	go func() {
		for {
			dbconn, err := database.GetDatabaseConnection()
			if err != nil {
				logger.Panicln(err)
			}
			db := &database.DBConn{DBer: dbconn, Appid: AppID}
			if err := db.HeartBeat(); err != nil {
				logger.Panicln(err)
			}
			dbconn.Close()
			logger.Printf("Watchdog beating every %d minute\n", interval)
			retchan <- true
			time.Sleep(time.Duration(interval) * time.Minute)
		}
	}()
	return retchan
}

func main() {
	dbcon, err := database.GetDatabaseConnection()
	if err != nil {
		logger.Panicln(err)
	}
	defer dbcon.Close()
	analyser := NewAnalyser(dbcon, redis.NewPool(func() (redis.Conn, error) { return database.GetRedisConnection(getredisconnect()) }, 10))

	var datachange, urlchange chan []byte
	var heartbeatchannel chan bool

	populatedatasetinfo := func() {
		if err := analyser.populatedatasetinfo(); err != nil {
			logger.Panicln(err)
		}
	}

	if !isonlyweb() {
		heartbeatchannel = heartbeat(getheartbeatinterval())
		<-heartbeatchannel // Wait for the first heartbeat, so the logging in the database is properly set up

		populatedatasetinfo()

		datachange = analyser.listenredischannel(watcherappid + ":DataChange")
		urlchange = analyser.listenredischannel(watcherappid + ":UrlChange")
	}

	restful.DefaultResponseMimeType = restful.MIME_JSON
	restful.EnableContentEncoding = true
	restful.Add(NewAnalyseOGDATRESTService(analyser))

	config := swagger.Config{
		WebServicesUrl:  discoveryhost(),
		ApiPath:         "/swaggerdoc",
		SwaggerPath:     "/doc/",
		SwaggerFilePath: "swagger-ui/dist/",
		WebServices:     restful.RegisteredWebServices()} // you control what services are visible
	swagger.InstallSwaggerService(config)

	logger.Printf("analyser (%s) listening on port %s\n", AppID, portbinding())
	go func() {
		logger.Fatal(http.ListenAndServe(":"+portbinding(), nil))
	}()

	for {
		select {
		case <-urlchange:
			// TODO:naive approach here. If a URLChange or DataChange event is triggered,
			// the whole analytic database will be recreated. It would be better to trace
			// only the affected datasets and only create the relevant statistic.
			// In future, urlchange/datachange might contain a JSON-encoded []byte which contains
			// the affected IDs
			logger.Println("Received URL change notice, re-generating database analysis")
			populatedatasetinfo()
		case <-datachange:
			logger.Println("Received Data change notice, re-generating database analysis")
			populatedatasetinfo()
		case <-heartbeatchannel:
			logger.Println("Idle")
		}
	}
}

func init() {
	logger = log.New(os.Stderr, filepath.Base(os.Args[0])+": ", log.LstdFlags)
}
<commit_msg>new version of go-restful handles flat / gzip encoding in the Container<commit_after>package main

import (
	"database/sql"
	restful "github.com/emicklei/go-restful"
	"github.com/emicklei/go-restful/swagger"
	"github.com/garyburd/redigo/redis"
	"github.com/the42/ogdat/database"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"time"
)

const AppID = "5bcbfc24-8e7e-4105-99c4-dd47e7e5094a"
const watcherappid = "a6545f8f-e0c9-4917-83c7-3e47bd1e0247"

var logger *log.Logger

type analyser struct {
	dbcon analyserdb
	pool  *redis.Pool
}

func NewAnalyser(dbcon *sql.DB, pool *redis.Pool) *analyser {
	analyser := &analyser{dbcon: analyserdb{DBConn: database.DBConn{Appid: AppID, DBer: dbcon}}, pool: pool}
	return analyser
}

func isonlyweb() bool {
	boolval, err := strconv.ParseBool(os.Getenv("ONLYWEB"))
	if err != nil {
		return false
	}
	return boolval
}

func getredisconnect() string {
	const redisurl = "REDISCLOUD_URL"
	const redisdb = "ANALYSER_REDISDB"

	return os.Getenv(redisurl) + "/" + os.Getenv(redisdb)
}

func portbinding() string {
	if port := os.Getenv("PORT"); port != "" {
		return port
	}
	return "8080"
}

func discoveryhost() string {
	if host := os.Getenv("DISCOVERYURL"); host != "" {
		return host
	}
	return "http://localhost"
}

func apibasepath() string {
	return os.Getenv("APIBASEPATH")
}

func getheartbeatinterval() int {

	if i, err := strconv.Atoi(os.Getenv("HEARTBEAT_INTERVAL")); err == nil {
		return i
	}
	return 60 // Minutes
}

func heartbeat(interval int) chan bool {
	retchan := make(chan bool)
	go func() {
		for {
			dbconn, err := database.GetDatabaseConnection()
			if err != nil {
				logger.Panicln(err)
			}
			db := &database.DBConn{DBer: dbconn, Appid: AppID}
			if err := db.HeartBeat(); err != nil {
				logger.Panicln(err)
			}
			dbconn.Close()
			logger.Printf("Watchdog beating every %d minute\n", interval)
			retchan <- true
			time.Sleep(time.Duration(interval) * time.Minute)
		}
	}()
	return retchan
}

func main() {
	dbcon, err := database.GetDatabaseConnection()
	if err != nil {
		logger.Panicln(err)
	}
	defer dbcon.Close()
	analyser := NewAnalyser(dbcon, redis.NewPool(func() (redis.Conn, error) { return database.GetRedisConnection(getredisconnect()) }, 10))

	var datachange, urlchange chan []byte
	var heartbeatchannel chan bool

	populatedatasetinfo := func() {
		if err := analyser.populatedatasetinfo(); err != nil {
			logger.Panicln(err)
		}
	}

	if !isonlyweb() {
		heartbeatchannel = heartbeat(getheartbeatinterval())
		<-heartbeatchannel // Wait for the first heartbeat, so the logging in the database is properly set up

		populatedatasetinfo()

		datachange = analyser.listenredischannel(watcherappid + ":DataChange")
		urlchange = analyser.listenredischannel(watcherappid + ":UrlChange")
	}

	restful.DefaultResponseMimeType = restful.MIME_JSON
	restful.DefaultContainer.EnableContentEncoding(true)
	restful.Add(NewAnalyseOGDATRESTService(analyser))

	config := swagger.Config{
		WebServicesUrl:  discoveryhost(),
		ApiPath:         "/swaggerdoc",
		SwaggerPath:     "/doc/",
		SwaggerFilePath: "swagger-ui/dist/",
		WebServices:     restful.RegisteredWebServices()} // you control what services are visible
	swagger.InstallSwaggerService(config)

	logger.Printf("analyser (%s) listening on port %s\n", AppID, portbinding())
	go func() {
		logger.Fatal(http.ListenAndServe(":"+portbinding(), nil))
	}()

	for {
		select {
		case <-urlchange:
			// TODO:naive approach here. If a URLChange or DataChange event is triggered,
			// the whole analytic database will be recreated. It would be better to trace
			// only the affected datasets and only create the relevant statistic.
			// In future, urlchange/datachange might contain a JSON-encoded []byte which contains
			// the affected IDs
			logger.Println("Received URL change notice, re-generating database analysis")
			populatedatasetinfo()
		case <-datachange:
			logger.Println("Received Data change notice, re-generating database analysis")
			populatedatasetinfo()
		case <-heartbeatchannel:
			logger.Println("Idle")
		}
	}
}

func init() {
	logger = log.New(os.Stderr, filepath.Base(os.Args[0])+": ", log.LstdFlags)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/garyburd/redigo/redis"
	"log"
	"net/url"
	"os"
	"path/filepath"
	"strings"
)

var logger *log.Logger

func getredisconnect() string {
	const redisurl = "REDISCLOUD_URL"
	const redisdb = "ANALYSER_REDISDB"

	return os.Getenv(redisurl) + "/" + os.Getenv(redisdb)
}

// taken from https://github.com/soveran/redisurl/blob/master/redisurl.go
func connectToURL(s string) (c redis.Conn, err error) {
	redisURL, err := url.Parse(s)

	if err != nil {
		return
	}

	auth := ""

	if redisURL.User != nil {
		if password, ok := redisURL.User.Password(); ok {
			auth = password
		}
	}

	c, err = redis.Dial("tcp", redisURL.Host)

	if err != nil {
		fmt.Println(err)
		return
	}

	if len(auth) > 0 {
		_, err = c.Do("AUTH", auth)

		if err != nil {
			fmt.Println(err)
			return
		}
	}

	if dbs := strings.Split(redisURL.Path, "/"); len(dbs) > 1 && dbs[1] != "" {
		c.Do("SELECT", dbs[1])
	}

	return
}

func main() {
	rcon, err := connectToURL(getredisconnect())
	if err != nil {
		logger.Panicln(err)
	}
	defer rcon.Close()
	rcon.Do("SET", "foo", 1)
	exists, _ := redis.Bool(rcon.Do("EXISTS", "foo"))
	fmt.Printf("%#v\n", exists)
	rcon.Do("DEL", "foo")
	exists, _ = redis.Bool(rcon.Do("EXISTS", "foo"))
	fmt.Printf("%#v\n", exists)

}

func init() {
	logger = log.New(os.Stderr, filepath.Base(os.Args[0])+": ", log.LstdFlags)
}
<commit_msg>add database connection to analyser to read from Postgres<commit_after>package main

import (
	"fmt"
	"github.com/garyburd/redigo/redis"
	"github.com/the42/ogdat/database"
	"log"
	"os"
	"path/filepath"
)

const AppID = "5bcbfc24-8e7e-4105-99c4-dd47e7e5094a"

var logger *log.Logger

func getredisconnect() string {
	const redisurl = "REDISCLOUD_URL"
	const redisdb = "ANALYSER_REDISDB"

	return os.Getenv(redisurl) + "/" + os.Getenv(redisdb)
}

func main() {
	rcon, err := database.GetRedisConnection(getredisconnect())
	if err != nil {
		logger.Panicln(err)
	}
	defer rcon.Close()

	dbcon, err := database.GetDatabaseConnection()
	if err != nil {
		logger.Panicln(err)
	}
	defer dbcon.Close()

	// BEGIN
	// some random tests
	rcon.Do("SET", "foo", 1)
	exists, _ := redis.Bool(rcon.Do("EXISTS", "foo"))
	fmt.Printf("%#v\n", exists)
	rcon.Do("DEL", "foo")
	exists, _ = redis.Bool(rcon.Do("EXISTS", "foo"))
	fmt.Printf("%#v\n", exists)
	// END

}

func init() {
	logger = log.New(os.Stderr, filepath.Base(os.Args[0])+": ", log.LstdFlags)
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	restful "github.com/emicklei/go-restful"
	"github.com/emicklei/go-restful/swagger"
	"github.com/garyburd/redigo/redis"
	"github.com/the42/ogdat/database"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"time"
)

const AppID = "5bcbfc24-8e7e-4105-99c4-dd47e7e5094a"
const watcherappid = "a6545f8f-e0c9-4917-83c7-3e47bd1e0247"

var logger *log.Logger

type analyser struct {
	dbcon analyserdb
	pool  *redis.Pool
}

func NewAnalyser(dbcon *sql.DB, pool *redis.Pool) *analyser {
	analyser := &analyser{dbcon: analyserdb{DBConn: database.DBConn{Appid: AppID, DBer: dbcon}}, pool: pool}
	return analyser
}

func isonlyweb() bool {
	boolval, err := strconv.ParseBool(os.Getenv("ONLYWEB"))
	if err != nil {
		return false
	}
	return boolval
}

func getredisconnect() string {
	const redisurl = "REDISCLOUD_URL"
	const redisdb = "ANALYSER_REDISDB"

	return os.Getenv(redisurl) + "/" + os.Getenv(redisdb)
}

func portbinding() string {
	if port := os.Getenv("PORT"); port != "" {
		return port
	}
	return "8080"
}

func hostname() string {
	if host := os.Getenv("DISCOVERYURL"); host != "" {
		return host
	}
	return "http://localhost"
}

func getheartbeatinterval() int {

	if i, err := strconv.Atoi(os.Getenv("HEARTBEAT_INTERVAL")); err == nil {
		return i
	}
	return 60 // Minutes
}

func heartbeat(interval int) chan bool {
	retchan := make(chan bool)
	go func() {
		for {
			dbconn, err := database.GetDatabaseConnection()
			if err != nil {
				logger.Panicln(err)
			}
			db := &database.DBConn{DBer: dbconn, Appid: AppID}
			if err := db.HeartBeat(); err != nil {
				logger.Panicln(err)
			}
			dbconn.Close()
			logger.Printf("Watchdog beating every %d minute\n", interval)
			retchan <- true
			time.Sleep(time.Duration(interval) * time.Minute)
		}
	}()
	return retchan
}

func main() {
	dbcon, err := database.GetDatabaseConnection()
	if err != nil {
		logger.Panicln(err)
	}
	defer dbcon.Close()
	analyser := NewAnalyser(dbcon, redis.NewPool(func() (redis.Conn, error) { return database.GetRedisConnection(getredisconnect()) }, 10))

	var datachange, urlchange chan []byte
	var heartbeatchannel chan bool

	if !isonlyweb() {
		heartbeatchannel = heartbeat(getheartbeatinterval())

		<-heartbeatchannel // Wait for the first heartbeat, so the logging in the database is properly set up
		if err := analyser.populatedatasetinfo(); err != nil {
			logger.Panicln(err)
		}
		datachange = analyser.listenredischannel(watcherappid + ":DataChange")
		urlchange = analyser.listenredischannel(watcherappid + ":UrlChange")
	}

	restful.DefaultResponseMimeType = restful.MIME_JSON
	restful.EnableContentEncoding = true
	restful.Add(NewAnalyseOGDATRESTService(analyser))

	config := swagger.Config{
		WebServicesUrl:  hostname() + ":" + portbinding(),
		ApiPath:         "/swaggerdoc",
		SwaggerPath:     "/doc/",
		SwaggerFilePath: "swagger-ui/dist/",
		WebServices:     restful.RegisteredWebServices()} // you control what services are visible
	swagger.InstallSwaggerService(config)

	logger.Printf("analyser (%s) listening on port %s\n", AppID, portbinding())
	go logger.Fatal(http.ListenAndServe(":"+portbinding(), nil))

	if !isonlyweb() {
		populatedatasetinfo := func() {
			if err := analyser.populatedatasetinfo(); err != nil {
				logger.Panicln(err)
			}
		}
		for {
			select {
			case <-urlchange:
				populatedatasetinfo()
			case <-datachange:
				populatedatasetinfo()
			case <-heartbeatchannel:
				logger.Println("Idle")
			}
		}
	}
}

func init() {
	logger = log.New(os.Stderr, filepath.Base(os.Args[0])+": ", log.LstdFlags)
}
<commit_msg>for configuration, no local port required<commit_after>package main

import (
	"database/sql"
	restful "github.com/emicklei/go-restful"
	"github.com/emicklei/go-restful/swagger"
	"github.com/garyburd/redigo/redis"
	"github.com/the42/ogdat/database"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"time"
)

const AppID = "5bcbfc24-8e7e-4105-99c4-dd47e7e5094a"
const watcherappid = "a6545f8f-e0c9-4917-83c7-3e47bd1e0247"

var logger *log.Logger

type analyser struct {
	dbcon analyserdb
	pool  *redis.Pool
}

func NewAnalyser(dbcon *sql.DB, pool *redis.Pool) *analyser {
	analyser := &analyser{dbcon: analyserdb{DBConn: database.DBConn{Appid: AppID, DBer: dbcon}}, pool: pool}
	return analyser
}

func isonlyweb() bool {
	boolval, err := strconv.ParseBool(os.Getenv("ONLYWEB"))
	if err != nil {
		return false
	}
	return boolval
}

func getredisconnect() string {
	const redisurl = "REDISCLOUD_URL"
	const redisdb = "ANALYSER_REDISDB"

	return os.Getenv(redisurl) + "/" + os.Getenv(redisdb)
}

func portbinding() string {
	if port := os.Getenv("PORT"); port != "" {
		return port
	}
	return "8080"
}

func hostname() string {
	if host := os.Getenv("DISCOVERYURL"); host != "" {
		return host
	}
	return "http://localhost"
}

func getheartbeatinterval() int {

	if i, err := strconv.Atoi(os.Getenv("HEARTBEAT_INTERVAL")); err == nil {
		return i
	}
	return 60 // Minutes
}

func heartbeat(interval int) chan bool {
	retchan := make(chan bool)
	go func() {
		for {
			dbconn, err := database.GetDatabaseConnection()
			if err != nil {
				logger.Panicln(err)
			}
			db := &database.DBConn{DBer: dbconn, Appid: AppID}
			if err := db.HeartBeat(); err != nil {
				logger.Panicln(err)
			}
			dbconn.Close()
			logger.Printf("Watchdog beating every %d minute\n", interval)
			retchan <- true
			time.Sleep(time.Duration(interval) * time.Minute)
		}
	}()
	return retchan
}

func main() {
	dbcon, err := database.GetDatabaseConnection()
	if err != nil {
		logger.Panicln(err)
	}
	defer dbcon.Close()
	analyser := NewAnalyser(dbcon, redis.NewPool(func() (redis.Conn, error) { return database.GetRedisConnection(getredisconnect()) }, 10))

	var datachange, urlchange chan []byte
	var heartbeatchannel chan bool

	if !isonlyweb() {
		heartbeatchannel = heartbeat(getheartbeatinterval())

		<-heartbeatchannel // Wait for the first heartbeat, so the logging in the database is properly set up
		if err := analyser.populatedatasetinfo(); err != nil {
			logger.Panicln(err)
		}
		datachange = analyser.listenredischannel(watcherappid + ":DataChange")
		urlchange = analyser.listenredischannel(watcherappid + ":UrlChange")
	}

	restful.DefaultResponseMimeType = restful.MIME_JSON
	restful.EnableContentEncoding = true
	restful.Add(NewAnalyseOGDATRESTService(analyser))

	config := swagger.Config{
		WebServicesUrl:  hostname(),
		ApiPath:         "/swaggerdoc",
		SwaggerPath:     "/doc/",
		SwaggerFilePath: "swagger-ui/dist/",
		WebServices:     restful.RegisteredWebServices()} // you control what services are visible
	swagger.InstallSwaggerService(config)

	logger.Printf("analyser (%s) listening on port %s\n", AppID, portbinding())
	go logger.Fatal(http.ListenAndServe(":"+portbinding(), nil))

	if !isonlyweb() {
		populatedatasetinfo := func() {
			if err := analyser.populatedatasetinfo(); err != nil {
				logger.Panicln(err)
			}
		}
		for {
			select {
			case <-urlchange:
				populatedatasetinfo()
			case <-datachange:
				populatedatasetinfo()
			case <-heartbeatchannel:
				logger.Println("Idle")
			}
		}
	}
}

func init() {
	logger = log.New(os.Stderr, filepath.Base(os.Args[0])+": ", log.LstdFlags)
}
<|endoftext|>


<commit_before>package app

import (
    "os"
    "runtime"

    "github.com/ouqiang/cron-scheduler/models"
    "github.com/ouqiang/cron-scheduler/modules/ansible"
    "github.com/ouqiang/cron-scheduler/modules/crontask"
    "github.com/ouqiang/cron-scheduler/service"
    "github.com/ouqiang/cron-scheduler/modules/setting"
    "github.com/ouqiang/cron-scheduler/modules/logger"
)

var (
    AppDir       string // 应用根目录
    ConfDir      string // 配置目录
    LogDir       string // 日志目录
    DataDir      string // 数据目录，存放session文件等
    AppConfig    string // 应用配置文件
    AnsibleHosts string // ansible hosts文件
    Installed    bool   // 应用是否安装过
    IsWindows    bool   // 是否是在windows上运行
)

func InitEnv() {
    logger.InitLogger()
    wd, err := os.Getwd()
    if err != nil {
        logger.Fatal(err)
    }
    AppDir = wd
    ConfDir = AppDir + "/conf"
    LogDir = AppDir + "/log"
    DataDir = AppDir + "/data"
    AppConfig = ConfDir + "/app.ini"
    AnsibleHosts = ConfDir + "/ansible_hosts.ini"
    IsWindows = runtime.GOOS == "windows"
    checkDirExists(ConfDir, LogDir, DataDir)
    // ansible配置文件目录
    os.Setenv("ANSIBLE_CONFIG", ConfDir)
    Installed = IsInstalled()
    if Installed {
        InitDb()
        InitResource()
    }
}

// 判断应用是否安装过
func IsInstalled() bool {
    _, err := os.Stat(ConfDir + "/install.lock")
    if os.IsNotExist(err) {
        return false
    }

    return true
}

// 创建安装锁文件
func CreateInstallLock() error {
    _, err := os.Create(ConfDir + "/install.lock")
    if err != nil {
        logger.Error("创建安装锁文件失败")
    }

    return err
}

// 初始化资源
func InitResource() {
    // 初始化ansible Hosts
    ansible.DefaultHosts = ansible.NewHosts(AnsibleHosts)
    // 初始化定时任务
    crontask.DefaultCronTask = crontask.NewCronTask()
    serviceTask := new(service.Task)
    serviceTask.Initialize()
}

// 初始化DB
func InitDb()  {
    dbConfig := getDbConfig(AppConfig)
    models.Db = models.CreateDb(dbConfig)
}

// 检测目录是否存在
func checkDirExists(path ...string) {
    for _, value := range path {
        _, err := os.Stat(value)
        if os.IsNotExist(err) {
            logger.Fatal(value + "目录不存在")
        }
        if os.IsPermission(err) {
            logger.Fatal(value + "目录无权限操作")
        }
    }
}

// 获取数据库配置
func getDbConfig(configFile string) map[string]string {
    config, err := setting.Read(configFile)
    if err != nil {
        logger.Fatal(err)
    }
    section := config.Section("db")
    if err != nil {
        logger.Fatal(err)
    }
    var db map[string]string = make(map[string]string)
    db["user"] = section.Key("user").String()
    db["password"] = section.Key("password").String()
    db["host"] = section.Key("host").String()
    db["port"] = section.Key("port").String()
    db["database"] = section.Key("database").String()
    db["charset"] = section.Key("charset").String()
    db["prefix"] = section.Key("prefix").String()
    db["engine"] = section.Key("engine").String()

    return db
}
<commit_msg>不支持在windows上运行<commit_after>package app

import (
    "os"
    "runtime"

    "github.com/ouqiang/cron-scheduler/models"
    "github.com/ouqiang/cron-scheduler/modules/ansible"
    "github.com/ouqiang/cron-scheduler/modules/crontask"
    "github.com/ouqiang/cron-scheduler/service"
    "github.com/ouqiang/cron-scheduler/modules/setting"
    "github.com/ouqiang/cron-scheduler/modules/logger"
)

var (
    AppDir       string // 应用根目录
    ConfDir      string // 配置目录
    LogDir       string // 日志目录
    DataDir      string // 数据目录，存放session文件等
    AppConfig    string // 应用配置文件
    AnsibleHosts string // ansible hosts文件
    Installed    bool   // 应用是否安装过
)

func InitEnv() {
    logger.InitLogger()
    wd, err := os.Getwd()
    if err != nil {
        logger.Fatal(err)
    }
    checkEnv()
    AppDir = wd
    ConfDir = AppDir + "/conf"
    LogDir = AppDir + "/log"
    DataDir = AppDir + "/data"
    AppConfig = ConfDir + "/app.ini"
    AnsibleHosts = ConfDir + "/ansible_hosts.ini"
    checkDirExists(ConfDir, LogDir, DataDir)
    // ansible配置文件目录
    os.Setenv("ANSIBLE_CONFIG", ConfDir)
    Installed = IsInstalled()
    if Installed {
        InitDb()
        InitResource()
    }
}

func checkEnv()  {
    if runtime.GOOS == "windows" {
        logger.Fatal("不支持在windows上运行")
    }
}

// 判断应用是否安装过
func IsInstalled() bool {
    _, err := os.Stat(ConfDir + "/install.lock")
    if os.IsNotExist(err) {
        return false
    }

    return true
}

// 创建安装锁文件
func CreateInstallLock() error {
    _, err := os.Create(ConfDir + "/install.lock")
    if err != nil {
        logger.Error("创建安装锁文件失败")
    }

    return err
}

// 初始化资源
func InitResource() {
    // 初始化ansible Hosts
    ansible.DefaultHosts = ansible.NewHosts(AnsibleHosts)
    // 初始化定时任务
    crontask.DefaultCronTask = crontask.NewCronTask()
    serviceTask := new(service.Task)
    serviceTask.Initialize()
}

// 初始化DB
func InitDb()  {
    dbConfig := getDbConfig(AppConfig)
    models.Db = models.CreateDb(dbConfig)
}

// 检测目录是否存在
func checkDirExists(path ...string) {
    for _, value := range path {
        _, err := os.Stat(value)
        if os.IsNotExist(err) {
            logger.Fatal(value + "目录不存在")
        }
        if os.IsPermission(err) {
            logger.Fatal(value + "目录无权限操作")
        }
    }
}

// 获取数据库配置
func getDbConfig(configFile string) map[string]string {
    config, err := setting.Read(configFile)
    if err != nil {
        logger.Fatal(err)
    }
    section := config.Section("db")
    if err != nil {
        logger.Fatal(err)
    }
    var db map[string]string = make(map[string]string)
    db["user"] = section.Key("user").String()
    db["password"] = section.Key("password").String()
    db["host"] = section.Key("host").String()
    db["port"] = section.Key("port").String()
    db["database"] = section.Key("database").String()
    db["charset"] = section.Key("charset").String()
    db["prefix"] = section.Key("prefix").String()
    db["engine"] = section.Key("engine").String()

    return db
}
<|endoftext|>


<commit_before>// Licensed under the Apache License, Version 2.0
// Details: https://raw.githubusercontent.com/maniksurtani/quotaservice/master/LICENSE

// Package redis implements token buckets backed by Redis, inspired by the algorithms used in Guava's
// RateLimiter library - https://github.com/google/guava/blob/master/guava/src/com/google/common/util/concurrent/RateLimiter.java
package redis

import (
	"strconv"
	"time"

	"gopkg.in/redis.v5"

	"github.com/maniksurtani/quotaservice"
	"github.com/maniksurtani/quotaservice/logging"

	"sync"

	pbconfig "github.com/maniksurtani/quotaservice/protos/config"
)

// Suffixes for Redis keys
const (
	tokensNextAvblNanosSuffix = "TNA"
	accumulatedTokensSuffix   = "AT"
	flushedAtVersionKey       = "FlushedAtVersion"
)

// defaultBucket is a "const"
var defaultBucket = &quotaservice.DefaultBucket{}

// bucketFactory holds an instance of the Redis client, and constructs staticBucket and dynamicBucket instances for use
// with Redis. Contains an embedded mutex which should be used when reading or updating the reference to the Redis
// client. Also holds references to configAttributes for each namespace and refcounts of usage of commonAttributes,
// both also guarded by this mutex.
type bucketFactory struct {
	// Embedded mutex
	sync.Mutex

	// Refcounts of configAttributes instances used by dynamic buckets for each namespace, protected by the
	// embedded mutex.
	refcounts map[string]int

	// sharedAttributes are instances of configAttributes used by dynamic buckets for each namespace, protected by
	// the embedded mutex.
	sharedAttributes map[string]*configAttributes

	cfg               *pbconfig.ServiceConfig
	client            *redis.Client
	redisOpts         *redis.Options
	scriptSHA         string
	connectionRetries int
}

// NewBucketFactory creates a new bucketFactory instance
func NewBucketFactory(redisOpts *redis.Options, connectionRetries int) quotaservice.BucketFactory {
	if connectionRetries < 1 {
		connectionRetries = 1
	}

	return &bucketFactory{
		redisOpts:         redisOpts,
		connectionRetries: connectionRetries,
		sharedAttributes:  make(map[string]*configAttributes),
		refcounts:         make(map[string]int)}
}

// Init initializes a bucketFactory for use, implementing Init() on the quotaservice.BucketFactory interface
func (bf *bucketFactory) Init(cfg *pbconfig.ServiceConfig) {
	logging.Printf("Initializing redis.bucketFactory for config version %v", cfg.Version)
	bf.Lock()
	defer bf.Unlock()

	bf.cfg = cfg

	if bf.client == nil {
		start := time.Now()
		bf.connectToRedisLocked()
		logging.Printf("Re-established Redis connections in %v", time.Since(start))
	}

	// Validate contents in Redis
	start := time.Now()
	// Check if Redis has been flushed when loading the current version of the configuration by a different node in
	// the cluster. If so, this node doesn't have to.
	val := bf.client.Get(flushedAtVersionKey)
	version, err := val.Int64()
	if err != nil {
		if err.Error() != "redis: nil" {
			// An actual problem; not just a missing flushedAtVersion.
			logging.Fatalf("Error response from Redis: %v", err)
		}

		// No flushedAtVersion has been established.
		logging.Print("flushedAtVersion not set")
		bf.flush(cfg.Version)

	} else {
		logging.Printf("flushedAtVersion is %v", version)
		if int32(version) < cfg.Version {
			bf.flush(cfg.Version)
		} else {
			logging.Printf("No need to flush since Redis has alreadt been flushedAtVersion %v", version)
		}
	}
	logging.Printf("Verified Redis (including any flushes, if necessary) in %v", time.Since(start))
}

// flush flushes unused entries stored in Redis
func (bf *bucketFactory) flush(version int32) {
	start := time.Now()
	// Store flushedAtVersion to prevent multiple nodes flushing Redis unnecessarily, while the flush is in
	// progress. This may be racy, but it's a minor optimization. At worst case, we have > 1 flushes (from > 1
	// nodes in the cluster), which, while wasteful, isn't dangerous.
	bf.client.Set(flushedAtVersionKey, version, 0)

	// TODO(manik): consider a batched SCAN + DELETE if the FLUSHDB operation is slow.
	bf.client.FlushDb()

	// Re-set flushedAtVersion, since previous entry would have been removed with the flush.
	bf.client.Set(flushedAtVersionKey, version, 0)
	logging.Printf("Flushed Redis in %v", time.Since(start))
}

func (bf *bucketFactory) connectToRedisLocked() {
	// Set up connection to Redis
	bf.client = redis.NewClient(bf.redisOpts)

	t, err := bf.client.Time().Result()
	if err != nil {
		logging.Printf("Cannot connect to Redis. TIME returned %v", err)
	} else {
		logging.Printf("Connection established. Time on Redis server: %v", t)
	}

	bf.scriptSHA = loadScript(bf.client)
}

func (bf *bucketFactory) reconnectToRedis(oldClient *redis.Client) {
	bf.Lock()
	defer bf.Unlock()

	// Always close connections on errors to prevent results leaking.
	if err := bf.client.Close(); unknownCloseError(err) {
		logging.Printf("Received error on Redis client close: %+v", err)
	}

	if oldClient == bf.client {
		bf.connectToRedisLocked()
	}
}

// Client returns a reference to the underlying client instance, implementing Client() on the quotaservice.BucketFactory
// interface
func (bf *bucketFactory) Client() interface{} {
	bf.Lock()
	defer bf.Unlock()

	return bf.client
}

// NewBucket creates and returns a new instance of quotaservice.Bucket, implementing NewBucket() on the
// quotaservice.BucketFactory interface
func (bf *bucketFactory) NewBucket(namespace, bucketName string, cfg *pbconfig.BucketConfig, dyn bool) quotaservice.Bucket {
	idle := "0"
	if cfg.MaxIdleMillis > 0 {
		idle = strconv.FormatInt(int64(cfg.MaxIdleMillis), 10)
	}

	keys := []string{toRedisKey(namespace, bucketName, tokensNextAvblNanosSuffix),
		toRedisKey(namespace, bucketName, accumulatedTokensSuffix)}

	if dyn {
		var exists bool
		bf.Lock()
		defer bf.Unlock()

		var attribs *configAttributes

		if attribs, exists = bf.sharedAttributes[namespace]; !exists {
			attribs = newConfigAttributes(cfg, idle, dyn)
			bf.sharedAttributes[namespace] = attribs
			bf.refcounts[namespace] = 0
		}
		bf.refcounts[namespace]++

		// Create a dynamicBucket with a reference to the appropriate shared configAttributes instance
		return &dynamicBucket{
			&abstractBucket{
				attribs,
				cfg,
				bf,
				keys}}
	} else {
		// Create a staticBucket with its own non-shared configAttributes
		return &staticBucket{
			&abstractBucket{
				newConfigAttributes(cfg, idle, dyn),
				cfg,
				bf,
				keys}}
	}
}

func newConfigAttributes(cfg *pbconfig.BucketConfig, idle string, dyn bool) *configAttributes {
	return &configAttributes{
		strconv.FormatInt(1e9/cfg.FillRate, 10),
		strconv.FormatInt(cfg.Size, 10),
		idle,
		// Convert millis to nanos
		strconv.FormatInt(cfg.MaxDebtMillis*1e6, 10),
		defaultBucket}
}

func toRedisKey(namespace, bucketName, suffix string) string {
	return namespace + ":" + bucketName + ":" + suffix
}

func unknownCloseError(err error) bool {
	return err != nil && err.Error() != "redis: client is closed"
}

func checkScriptExists(c *redis.Client, sha string) bool {
	r := c.ScriptExists(sha)
	return r.Val()[0]
}

// loadScript loads the LUA script into Redis. The LUA script contains the token bucket algorithm
// which is executed atomically in Redis. Once the script is loaded, it is invoked using its SHA.
func loadScript(c *redis.Client) (sha string) {
	lua := `
	local tokensNextAvailableNanos = tonumber(redis.call("GET", KEYS[1]))
	if not tokensNextAvailableNanos then
		tokensNextAvailableNanos = 0
	end

	local maxTokensToAccumulate = tonumber(ARGV[3])

	local accumulatedTokens = redis.call("GET", KEYS[2])
	if not accumulatedTokens then
		accumulatedTokens = maxTokensToAccumulate
	end

	local currentTimeNanos = tonumber(ARGV[1])
	local nanosBetweenTokens = tonumber(ARGV[2])
	local requested = tonumber(ARGV[4])
	local maxWaitTime = tonumber(ARGV[5])
	local lifespan = tonumber(ARGV[6])
	local maxDebtNanos = tonumber(ARGV[7])
	local freshTokens = 0

	if currentTimeNanos > tokensNextAvailableNanos then
		freshTokens = math.floor((currentTimeNanos - tokensNextAvailableNanos) / nanosBetweenTokens)
		accumulatedTokens = math.min(maxTokensToAccumulate, accumulatedTokens + freshTokens)
		tokensNextAvailableNanos = currentTimeNanos
	end

	local waitTime = tokensNextAvailableNanos - currentTimeNanos
	local accumulatedTokensUsed = math.min(accumulatedTokens, requested)
	local tokensToWaitFor = requested - accumulatedTokensUsed
	local futureWaitNanos = tokensToWaitFor * nanosBetweenTokens

	tokensNextAvailableNanos = tokensNextAvailableNanos + futureWaitNanos
	accumulatedTokens = accumulatedTokens - accumulatedTokensUsed

	if (tokensNextAvailableNanos - currentTimeNanos > maxDebtNanos) or (waitTime > 0 and waitTime > maxWaitTime) then
    	waitTime = -1
	else
		if lifespan > 0 then
			redis.call("SET", KEYS[1], tokensNextAvailableNanos, "PX", lifespan)
			redis.call("SET", KEYS[2], math.floor(accumulatedTokens), "PX", lifespan)
		else
			redis.call("SET", KEYS[1], tokensNextAvailableNanos)
			redis.call("SET", KEYS[2], math.floor(accumulatedTokens))
		end
	end

	return waitTime
	`
	s := c.ScriptLoad(lua)
	if s.Err() != nil {
		logging.Fatalf("Unable to load LUA script into Redis; error=%v", s.Err())
		return
	}

	sha = s.Val()
	logging.Printf("Loaded LUA script into Redis; script SHA %v", sha)
	return
}
<commit_msg>Fix comments (#68)<commit_after>// Licensed under the Apache License, Version 2.0
// Details: https://raw.githubusercontent.com/maniksurtani/quotaservice/master/LICENSE

// Package redis implements token buckets backed by Redis, inspired by the algorithms used in Guava's
// RateLimiter library - https://github.com/google/guava/blob/master/guava/src/com/google/common/util/concurrent/RateLimiter.java
package redis

import (
	"strconv"
	"time"

	"gopkg.in/redis.v5"

	"github.com/maniksurtani/quotaservice"
	"github.com/maniksurtani/quotaservice/logging"

	"sync"

	pbconfig "github.com/maniksurtani/quotaservice/protos/config"
)

// Suffixes for Redis keys
const (
	tokensNextAvblNanosSuffix = "TNA"
	accumulatedTokensSuffix   = "AT"
	flushedAtVersionKey       = "FlushedAtVersion"
)

// defaultBucket is a "const"
var defaultBucket = &quotaservice.DefaultBucket{}

// bucketFactory holds an instance of the Redis client, and constructs staticBucket and dynamicBucket instances for use
// with Redis. Contains an embedded mutex which should be used when reading or updating the reference to the Redis
// client. Also holds references to configAttributes for each namespace and refcounts of usage of commonAttributes,
// both also guarded by this mutex.
type bucketFactory struct {
	// Embedded mutex
	sync.Mutex

	// Refcounts of configAttributes instances used by dynamic buckets for each namespace, protected by the
	// embedded mutex.
	refcounts map[string]int

	// sharedAttributes are instances of configAttributes used by dynamic buckets for each namespace, protected by
	// the embedded mutex.
	sharedAttributes map[string]*configAttributes

	cfg               *pbconfig.ServiceConfig
	client            *redis.Client
	redisOpts         *redis.Options
	scriptSHA         string
	connectionRetries int
}

// NewBucketFactory creates a new bucketFactory instance
func NewBucketFactory(redisOpts *redis.Options, connectionRetries int) quotaservice.BucketFactory {
	if connectionRetries < 1 {
		connectionRetries = 1
	}

	return &bucketFactory{
		redisOpts:         redisOpts,
		connectionRetries: connectionRetries,
		sharedAttributes:  make(map[string]*configAttributes),
		refcounts:         make(map[string]int)}
}

// Init initializes a bucketFactory for use, implementing Init() on the quotaservice.BucketFactory interface
func (bf *bucketFactory) Init(cfg *pbconfig.ServiceConfig) {
	logging.Printf("Initializing redis.bucketFactory for config version %v", cfg.Version)
	bf.Lock()
	defer bf.Unlock()

	bf.cfg = cfg

	if bf.client == nil {
		start := time.Now()
		bf.connectToRedisLocked()
		logging.Printf("Re-established Redis connections in %v", time.Since(start))
	}

	// Validate contents in Redis
	start := time.Now()
	// Check if Redis has been flushed when loading the current version of the configuration by a different node in
	// the cluster. If so, this node doesn't have to.
	val := bf.client.Get(flushedAtVersionKey)
	version, err := val.Int64()
	if err != nil {
		if err.Error() != "redis: nil" {
			// An actual problem; not just a missing flushedAtVersion.
			logging.Fatalf("Error response from Redis: %v", err)
		}

		// No flushedAtVersion has been established.
		logging.Print("flushedAtVersion not set")
		bf.flush(cfg.Version)

	} else {
		logging.Printf("flushedAtVersion is %v", version)
		if int32(version) < cfg.Version {
			bf.flush(cfg.Version)
		} else {
			logging.Printf("No need to flush since Redis has already been flushedAtVersion %v", version)
		}
	}
	logging.Printf("Verified Redis (including any flushes, if necessary) in %v", time.Since(start))
}

// flush flushes unused entries stored in Redis
func (bf *bucketFactory) flush(version int32) {
	start := time.Now()
	// Store flushedAtVersion to prevent multiple nodes flushing Redis unnecessarily, while the flush is in
	// progress. This may be racy, but it's a minor optimization. At worst case, we have > 1 flushes (from > 1
	// nodes in the cluster), which, while wasteful, isn't dangerous.
	bf.client.Set(flushedAtVersionKey, version, 0)

	// We could consider a batched SCAN + DELETE if the FLUSHDB operation is slow. But for the most part, this is
	// "fast enough" - to the order of a few 100s of µs.
	bf.client.FlushDb()

	// Re-set flushedAtVersion, since previous entry would have been removed with the flush.
	bf.client.Set(flushedAtVersionKey, version, 0)
	logging.Printf("Flushed Redis in %v", time.Since(start))
}

func (bf *bucketFactory) connectToRedisLocked() {
	// Set up connection to Redis
	bf.client = redis.NewClient(bf.redisOpts)

	t, err := bf.client.Time().Result()
	if err != nil {
		logging.Printf("Cannot connect to Redis. TIME returned %v", err)
	} else {
		logging.Printf("Connection established. Time on Redis server: %v", t)
	}

	bf.scriptSHA = loadScript(bf.client)
}

func (bf *bucketFactory) reconnectToRedis(oldClient *redis.Client) {
	bf.Lock()
	defer bf.Unlock()

	// Always close connections on errors to prevent results leaking.
	if err := bf.client.Close(); unknownCloseError(err) {
		logging.Printf("Received error on Redis client close: %+v", err)
	}

	if oldClient == bf.client {
		bf.connectToRedisLocked()
	}
}

// Client returns a reference to the underlying client instance, implementing Client() on the quotaservice.BucketFactory
// interface
func (bf *bucketFactory) Client() interface{} {
	bf.Lock()
	defer bf.Unlock()

	return bf.client
}

// NewBucket creates and returns a new instance of quotaservice.Bucket, implementing NewBucket() on the
// quotaservice.BucketFactory interface
func (bf *bucketFactory) NewBucket(namespace, bucketName string, cfg *pbconfig.BucketConfig, dyn bool) quotaservice.Bucket {
	idle := "0"
	if cfg.MaxIdleMillis > 0 {
		idle = strconv.FormatInt(int64(cfg.MaxIdleMillis), 10)
	}

	keys := []string{toRedisKey(namespace, bucketName, tokensNextAvblNanosSuffix),
		toRedisKey(namespace, bucketName, accumulatedTokensSuffix)}

	if dyn {
		var exists bool
		bf.Lock()
		defer bf.Unlock()

		var attribs *configAttributes

		if attribs, exists = bf.sharedAttributes[namespace]; !exists {
			attribs = newConfigAttributes(cfg, idle, dyn)
			bf.sharedAttributes[namespace] = attribs
			bf.refcounts[namespace] = 0
		}
		bf.refcounts[namespace]++

		// Create a dynamicBucket with a reference to the appropriate shared configAttributes instance
		return &dynamicBucket{
			&abstractBucket{
				attribs,
				cfg,
				bf,
				keys}}
	} else {
		// Create a staticBucket with its own non-shared configAttributes
		return &staticBucket{
			&abstractBucket{
				newConfigAttributes(cfg, idle, dyn),
				cfg,
				bf,
				keys}}
	}
}

func newConfigAttributes(cfg *pbconfig.BucketConfig, idle string, dyn bool) *configAttributes {
	return &configAttributes{
		strconv.FormatInt(1e9/cfg.FillRate, 10),
		strconv.FormatInt(cfg.Size, 10),
		idle,
		// Convert millis to nanos
		strconv.FormatInt(cfg.MaxDebtMillis*1e6, 10),
		defaultBucket}
}

func toRedisKey(namespace, bucketName, suffix string) string {
	return namespace + ":" + bucketName + ":" + suffix
}

func unknownCloseError(err error) bool {
	return err != nil && err.Error() != "redis: client is closed"
}

func checkScriptExists(c *redis.Client, sha string) bool {
	r := c.ScriptExists(sha)
	return r.Val()[0]
}

// loadScript loads the LUA script into Redis. The LUA script contains the token bucket algorithm
// which is executed atomically in Redis. Once the script is loaded, it is invoked using its SHA.
func loadScript(c *redis.Client) (sha string) {
	lua := `
	local tokensNextAvailableNanos = tonumber(redis.call("GET", KEYS[1]))
	if not tokensNextAvailableNanos then
		tokensNextAvailableNanos = 0
	end

	local maxTokensToAccumulate = tonumber(ARGV[3])

	local accumulatedTokens = redis.call("GET", KEYS[2])
	if not accumulatedTokens then
		accumulatedTokens = maxTokensToAccumulate
	end

	local currentTimeNanos = tonumber(ARGV[1])
	local nanosBetweenTokens = tonumber(ARGV[2])
	local requested = tonumber(ARGV[4])
	local maxWaitTime = tonumber(ARGV[5])
	local lifespan = tonumber(ARGV[6])
	local maxDebtNanos = tonumber(ARGV[7])
	local freshTokens = 0

	if currentTimeNanos > tokensNextAvailableNanos then
		freshTokens = math.floor((currentTimeNanos - tokensNextAvailableNanos) / nanosBetweenTokens)
		accumulatedTokens = math.min(maxTokensToAccumulate, accumulatedTokens + freshTokens)
		tokensNextAvailableNanos = currentTimeNanos
	end

	local waitTime = tokensNextAvailableNanos - currentTimeNanos
	local accumulatedTokensUsed = math.min(accumulatedTokens, requested)
	local tokensToWaitFor = requested - accumulatedTokensUsed
	local futureWaitNanos = tokensToWaitFor * nanosBetweenTokens

	tokensNextAvailableNanos = tokensNextAvailableNanos + futureWaitNanos
	accumulatedTokens = accumulatedTokens - accumulatedTokensUsed

	if (tokensNextAvailableNanos - currentTimeNanos > maxDebtNanos) or (waitTime > 0 and waitTime > maxWaitTime) then
    	waitTime = -1
	else
		if lifespan > 0 then
			redis.call("SET", KEYS[1], tokensNextAvailableNanos, "PX", lifespan)
			redis.call("SET", KEYS[2], math.floor(accumulatedTokens), "PX", lifespan)
		else
			redis.call("SET", KEYS[1], tokensNextAvailableNanos)
			redis.call("SET", KEYS[2], math.floor(accumulatedTokens))
		end
	end

	return waitTime
	`
	s := c.ScriptLoad(lua)
	if s.Err() != nil {
		logging.Fatalf("Unable to load LUA script into Redis; error=%v", s.Err())
		return
	}

	sha = s.Val()
	logging.Printf("Loaded LUA script into Redis; script SHA %v", sha)
	return
}
<|endoftext|>


<commit_before>// © 2014 Steve McCoy under the MIT license. See LICENSE for details.

package pngplus

import (
	"bytes"
	"encoding/binary"
	"io"
	"hash/crc32"
	"testing"
)

func TestEncodeBinary(t *testing.T) {
	var w bytes.Buffer
	s := []byte("hello")

	err := EncodeBinary(&w, s)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	bs := w.Bytes()

	var n uint32
	err = binary.Read(bytes.NewBuffer(bs), binary.BigEndian, &n)
	if err != nil {
		t.Fatalf("unexpected length read error: %v", err)
	}

	if n != uint32(len(s)) {
		t.Fatalf("wrong chunk length: %d vs. %d", n, len(s))
	}

	content := string(bs[4:len(bs)-4])
	expected := "biNyhello"
	if content != expected {
		t.Fatalf("wrong content: %q vs. %q", content, expected)
	}

	var crc uint32
	err = binary.Read(bytes.NewBuffer(bs[8+len(s):]), binary.BigEndian, &crc)
	if err != nil {
		t.Fatalf("unexpected crc read error: %v", err)
	}

	excrc := crc32.ChecksumIEEE([]byte(expected))
	if crc != excrc {
		t.Fatalf("wrong crc: %d vs. %d", crc, excrc)
	}
}

func TestDecodeBinary(t *testing.T) {
	var w bytes.Buffer

	// Very short
	_, err := DecodeBinary(&w)
	if err != io.EOF {
		t.Errorf("expected EOF, got %v", err)
	}

	s := []byte("hello")

	// Normal
	w.Reset()
	err = EncodeBinary(&w, s)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	result, err := DecodeBinary(&w)
	if err != nil {
		t.Fatalf("failed to decode: %v", err)
	}
	if string(result) != "hello" {
		t.Fatalf("unexpected result: %q", string(result))
	}
	if w.Len() != 0 {
		t.Fatalf("unexpected remains in w: %d", w.Len())
	}

	// Non-biNy
	w.Reset()
	err = EncodeBinary(&w, s)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	w.Bytes()[4] = 'q'

	result, err = DecodeBinary(&w)
	if err != ErrNotBinary {
		t.Fatalf("expected ErrNotBinary, got: %v", err)
	}
	if w.Len() != 0 {
		t.Fatalf("expected chunk to be skipped, have %d remaining", w.Len())
	}

	// Non-biNy and short
	w.Reset()
	err = EncodeBinary(&w, s)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	w.Bytes()[4] = 'q'

	result, err = DecodeBinary(&io.LimitedReader{R: &w, N: int64(w.Len())-1})
	if err != io.EOF {
		t.Fatalf("expected EOF, got: %v", err)
	}
}
<commit_msg>Add a test for short read during chunk content.<commit_after>// © 2014 Steve McCoy under the MIT license. See LICENSE for details.

package pngplus

import (
	"bytes"
	"encoding/binary"
	"io"
	"hash/crc32"
	"testing"
)

func TestEncodeBinary(t *testing.T) {
	var w bytes.Buffer
	s := []byte("hello")

	err := EncodeBinary(&w, s)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	bs := w.Bytes()

	var n uint32
	err = binary.Read(bytes.NewBuffer(bs), binary.BigEndian, &n)
	if err != nil {
		t.Fatalf("unexpected length read error: %v", err)
	}

	if n != uint32(len(s)) {
		t.Fatalf("wrong chunk length: %d vs. %d", n, len(s))
	}

	content := string(bs[4:len(bs)-4])
	expected := "biNyhello"
	if content != expected {
		t.Fatalf("wrong content: %q vs. %q", content, expected)
	}

	var crc uint32
	err = binary.Read(bytes.NewBuffer(bs[8+len(s):]), binary.BigEndian, &crc)
	if err != nil {
		t.Fatalf("unexpected crc read error: %v", err)
	}

	excrc := crc32.ChecksumIEEE([]byte(expected))
	if crc != excrc {
		t.Fatalf("wrong crc: %d vs. %d", crc, excrc)
	}
}

func TestDecodeBinary(t *testing.T) {
	var w bytes.Buffer

	// Very short
	_, err := DecodeBinary(&w)
	if err != io.EOF {
		t.Errorf("expected EOF, got %v", err)
	}

	s := []byte("hello")

	// Normal
	w.Reset()
	err = EncodeBinary(&w, s)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	result, err := DecodeBinary(&w)
	if err != nil {
		t.Fatalf("failed to decode: %v", err)
	}
	if string(result) != "hello" {
		t.Fatalf("unexpected result: %q", string(result))
	}
	if w.Len() != 0 {
		t.Fatalf("unexpected remains in w: %d", w.Len())
	}

	// Non-biNy
	w.Reset()
	err = EncodeBinary(&w, s)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	w.Bytes()[4] = 'q'

	result, err = DecodeBinary(&w)
	if err != ErrNotBinary {
		t.Fatalf("expected ErrNotBinary, got: %v", err)
	}
	if w.Len() != 0 {
		t.Fatalf("expected chunk to be skipped, have %d remaining", w.Len())
	}

	// Non-biNy and short
	w.Reset()
	err = EncodeBinary(&w, s)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	w.Bytes()[4] = 'q'

	result, err = DecodeBinary(&io.LimitedReader{R: &w, N: int64(w.Len())-1})
	if err != io.EOF {
		t.Fatalf("expected EOF, got: %v", err)
	}

	// biNy and short during content
	w.Reset()
	err = EncodeBinary(&w, s)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	result, err = DecodeBinary(&io.LimitedReader{R: &w, N: int64(w.Len())-5})
	if err != io.ErrUnexpectedEOF {
		t.Fatalf("expected UnexpectedEOF, got: %v", err)
	}

}
<|endoftext|>


<commit_before>// Copyright 2019 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.

// Command bbagent is Buildbucket's agent running in swarming.
//
// This executable creates a luciexe 'host' environment, and runs the
// Buildbucket build's exe within this environment. Please see
// https://go.chromium.org/luci/luciexe for details about the 'luciexe'
// protocol.
//
// This command is an implementation detail of Buildbucket.
package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"
	"sync/atomic"
	"time"

	"net/http"
	_ "net/http/pprof"

	"github.com/golang/protobuf/jsonpb"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/fieldmaskpb"
	"google.golang.org/protobuf/types/known/timestamppb"

	"go.chromium.org/luci/buildbucket"
	"go.chromium.org/luci/buildbucket/cmd/bbagent/bbinput"
	bbpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/logging/gologger"
	"go.chromium.org/luci/common/sync/dispatcher"
	"go.chromium.org/luci/common/system/environ"
	"go.chromium.org/luci/grpc/grpcutil"
	"go.chromium.org/luci/lucictx"
	"go.chromium.org/luci/luciexe"
	"go.chromium.org/luci/luciexe/host"
	"go.chromium.org/luci/luciexe/invoke"
)

var maxReserveDuration = 2 * time.Minute

func main() {
	go func() {
		// serves "/debug" endpoints for pprof.
		log.Println(http.ListenAndServe("localhost:6060", nil))
	}()
	os.Exit(mainImpl())
}

func mainImpl() int {
	ctx := logging.SetLevel(gologger.StdConfig.Use(context.Background()), logging.Info)

	check := func(err error) {
		if err != nil {
			logging.Errorf(ctx, err.Error())
			os.Exit(1)
		}
	}

	outputFile := luciexe.AddOutputFlagToSet(flag.CommandLine)
	flag.Parse()
	args := flag.Args()

	if len(args) != 1 {
		check(errors.Reason("expected 1 argument: got %d", len(args)).Err())
	}

	input, err := bbinput.Parse(args[0])
	check(errors.Annotate(err, "could not unmarshal BBAgentArgs").Err())

	sctx, err := lucictx.SwitchLocalAccount(ctx, "system")
	check(errors.Annotate(err, "could not switch to 'system' account in LUCI_CONTEXT").Err())

	// We start with retries disabled because the dispatcher.Channel will handle
	// them during the execution of the user process; In particular we want
	// dispatcher.Channel to be able to move on to a newer version of the Build if
	// it encounters transient errors, rather than retrying a potentially stale
	// Build state.
	//
	// We enable them again after the user process has finished.
	bbclientRetriesEnabled := false
	bbclient, secrets, err := newBuildsClient(sctx, input.Build.Infra.Buildbucket, &bbclientRetriesEnabled)
	check(errors.Annotate(err, "could not connect to Buildbucket").Err())
	logdogOutput, err := mkLogdogOutput(sctx, input.Build.Infra.Logdog)
	check(errors.Annotate(err, "could not create logdog output").Err())

	var (
		cctx   context.Context
		cancel func()
	)
	if dl := lucictx.GetDeadline(ctx); dl.GetSoftDeadline() != 0 {
		softDeadline := dl.SoftDeadlineTime()
		gracePeriod := time.Duration(dl.GetGracePeriod() * float64(time.Second))
		cctx, cancel = context.WithDeadline(ctx, softDeadline.Add(gracePeriod))
	} else {
		cctx, cancel = context.WithCancel(ctx)
	}
	defer cancel()

	dispatcherOpts, dispatcherErrCh := channelOpts(cctx)
	buildsCh, err := dispatcher.NewChannel(cctx, dispatcherOpts, mkSendFn(cctx, secrets, bbclient))
	check(errors.Annotate(err, "could not create builds dispatcher channel").Err())
	defer buildsCh.CloseAndDrain(cctx)

	// from this point forward we want to try to report errors to buildbucket,
	// too.
	check = func(err error) {
		if err != nil {
			logging.Errorf(cctx, err.Error())
			buildsCh.C <- &bbpb.Build{
				Status:          bbpb.Status_INFRA_FAILURE,
				SummaryMarkdown: fmt.Sprintf("fatal error in startup: %s", err),
			}
			buildsCh.CloseAndDrain(cctx)
			os.Exit(1)
		}
	}

	cctx = setResultDBContext(cctx, input.Build, secrets)
	prepareInputBuild(cctx, input.Build)

	opts := &host.Options{
		BaseBuild:      input.Build,
		ButlerLogLevel: logging.Warning,
		ViewerURL: fmt.Sprintf("https://%s/build/%d",
			input.Build.Infra.Buildbucket.Hostname, input.Build.Id),
		LogdogOutput: logdogOutput,
		ExeAuth:      host.DefaultExeAuth("bbagent", input.KnownPublicGerritHosts),
	}
	cwd, err := os.Getwd()
	check(errors.Annotate(err, "getting cwd").Err())
	opts.BaseDir = filepath.Join(cwd, "x")

	exeArgs := append(([]string)(nil), input.Build.Exe.Cmd...)
	payloadPath := input.PayloadPath
	if len(exeArgs) == 0 {
		// TODO(iannucci): delete me with ExecutablePath.
		var exe string
		payloadPath, exe = path.Split(input.ExecutablePath)
		exeArgs = []string{exe}
	}
	exePath, err := filepath.Abs(filepath.Join(payloadPath, exeArgs[0]))
	check(errors.Annotate(err, "absoluting exe path %q", input.ExecutablePath).Err())
	if runtime.GOOS == "windows" {
		exePath, err = resolveExe(exePath)
		check(errors.Annotate(err, "resolving %q", input.ExecutablePath).Err())
	}
	exeArgs[0] = exePath

	initialJSONPB, err := (&jsonpb.Marshaler{
		OrigName: true, Indent: "  ",
	}).MarshalToString(input)
	check(errors.Annotate(err, "marshalling input args").Err())
	logging.Infof(ctx, "Input args:\n%s", initialJSONPB)

	shutdownCh := make(chan struct{})
	var statusDetails *bbpb.StatusDetails
	builds, err := host.Run(cctx, opts, func(ctx context.Context, hostOpts host.Options) error {
		logging.Infof(ctx, "running luciexe: %q", exeArgs)
		logging.Infof(ctx, "  (cache dir): %q", input.CacheDir)
		invokeOpts := &invoke.Options{
			BaseDir:  hostOpts.BaseDir,
			CacheDir: input.CacheDir,
		}
		reserve := calcDeadlineReserve(ctx)
		deadlineEventCh, dctx, shutdown := lucictx.AdjustDeadline(ctx, reserve, 500*time.Millisecond)
		go func() {
			select {
			case <-shutdownCh:
				shutdown()
			case <-dctx.Done():
			}
		}()
		if newSoftDeadline := lucictx.GetDeadline(dctx).SoftDeadlineTime(); !newSoftDeadline.IsZero() {
			logging.Infof(dctx,
				"attempted to reserve %s from soft deadline for bbagent cleanup; new soft deadline: %s",
				reserve, newSoftDeadline)
		}
		subp, err := invoke.Start(dctx, exeArgs, input.Build, invokeOpts, deadlineEventCh)
		if err != nil {
			return err
		}
		var build *bbpb.Build
		build, err = subp.Wait()
		statusDetails = build.StatusDetails
		return err
	})
	if err != nil {
		check(errors.Annotate(err, "could not start luciexe host environment").Err())
	}

	var (
		finalBuild                *bbpb.Build = proto.Clone(input.Build).(*bbpb.Build)
		fatalUpdateBuildErrorSlot atomic.Value
	)

	go func() {
		// Monitors the `dispatcherErrCh` and checks for fatal error determined by
		// `fatalPred`. Stops the build shuttling and shuts down the luciexe if a
		// fatal error is received.
		stopped := false
		fatalPred := func(err error) bool {
			// TODO(crbug.com/1140612): Figure out a better solution to handle the
			// InvalidArgument error at the beginning of the build when build status
			// hasn't switched to STARTED yet.
			// Possible solution: keep calling GetBuild and start shuttling the build
			// once GetBuild returns STARTED status. For now, simply tolerate such
			// error.
			if grpcutil.Code(err) == codes.InvalidArgument && !strings.Contains(err.Error(), "cannot update steps of a SCHEDULED build") {
				return true
			}
			return false
		}
		for {
			select {
			case err := <-dispatcherErrCh:
				if !stopped && fatalPred(err) {
					close(shutdownCh)
					fatalUpdateBuildErrorSlot.Store(err)
					stopped = true
				}
			case <-cctx.Done():
				return
			}
		}
	}()

	// Now all we do is shuttle builds through to the buildbucket client channel
	// until there are no more builds to shuttle.
	for build := range builds {
		if fatalUpdateBuildErrorSlot.Load() == nil {
			buildsCh.C <- build
			finalBuild = build
		}
	}
	buildsCh.CloseAndDrain(cctx)

	// Now that the builds channel has been closed, update bb directly.
	updateMask := []string{
		"build.status",
		"build.status_details",
		"build.summary_markdown",
	}
	var retcode int

	fatalUpdateBuildErr, _ := fatalUpdateBuildErrorSlot.Load().(error)
	if finalizeBuild(ctx, finalBuild, fatalUpdateBuildErr, statusDetails, outputFile) {
		updateMask = append(updateMask, "build.steps", "build.output")
	} else {
		// finalizeBuild indicated that something is really wrong; Omit steps and
		// output from the final push to minimize potential issues.
		retcode = 1
	}
	bbclientRetriesEnabled = true
	_, bbErr := bbclient.UpdateBuild(
		metadata.NewOutgoingContext(cctx, metadata.Pairs(buildbucket.BuildTokenHeader, secrets.BuildToken)),
		&bbpb.UpdateBuildRequest{
			Build:      finalBuild,
			UpdateMask: &fieldmaskpb.FieldMask{Paths: updateMask},
		})
	if bbErr != nil {
		logging.Errorf(cctx, "Failed to report error %s to Buildbucket due to %s", err, bbErr)
		retcode = 2
	}

	return retcode
}

// finalizeBuild returns true if fatalErr is nil and there's no additional
// errors finalizing the build.
func finalizeBuild(ctx context.Context, finalBuild *bbpb.Build, fatalErr error, statusDetails *bbpb.StatusDetails, outputFile *luciexe.OutputFlag) bool {
	proto.Merge(finalBuild.StatusDetails, statusDetails)

	// set final times
	now := timestamppb.New(clock.Now(ctx))
	finalBuild.UpdateTime = now
	finalBuild.EndTime = now

	var finalErrs errors.MultiError
	if fatalErr != nil {
		finalErrs = append(finalErrs, errors.Annotate(fatalErr, "fatal error in builbucket.UpdateBuild").Err())
	}
	if err := outputFile.Write(finalBuild); err != nil {
		finalErrs = append(finalErrs, errors.Annotate(err, "writing final build").Err())
	}

	if len(finalErrs) > 0 {
		errors.Log(ctx, finalErrs)

		// we had some really bad error, just downgrade status and add a message to
		// summary markdown.
		finalBuild.Status = bbpb.Status_INFRA_FAILURE
		originalSM := finalBuild.SummaryMarkdown
		finalBuild.SummaryMarkdown = fmt.Sprintf("FATAL: %s", finalErrs.Error())
		if originalSM != "" {
			finalBuild.SummaryMarkdown += "\n\n" + originalSM
		}
	}

	return len(finalErrs) == 0
}

func prepareInputBuild(ctx context.Context, build *bbpb.Build) {
	// mark started
	build.Status = bbpb.Status_STARTED
	now := timestamppb.New(clock.Now(ctx))
	build.StartTime, build.UpdateTime = now, now
	// TODO(iannucci): this is sketchy, but we preemptively add the log entries
	// for the top level user stdout/stderr streams.
	//
	// Really, `invoke.Start` is the one that knows how to arrange the
	// Output.Logs, but host.Run makes a copy of this build immediately. Find
	// a way to set these up nicely (maybe have opts.BaseBuild be a function
	// returning an immutable bbpb.Build?).
	build.Output = &bbpb.Build_Output{
		Logs: []*bbpb.Log{
			{Name: "stdout", Url: "stdout"},
			{Name: "stderr", Url: "stderr"},
		},
	}
	populateSwarmingInfoFromEnv(build, environ.System())
	return
}

// Returns min(1% of the remaining time towards current soft deadline,
// `maxReserveDuration`). Returns 0 If soft deadline doesn't exist or
// has already been exceeded.
func calcDeadlineReserve(ctx context.Context) time.Duration {
	curSoftDeadline := lucictx.GetDeadline(ctx).SoftDeadlineTime()
	if now := clock.Now(ctx).UTC(); !curSoftDeadline.IsZero() && now.Before(curSoftDeadline) {
		if reserve := curSoftDeadline.Sub(now) / 100; reserve < maxReserveDuration {
			return reserve
		}
		return maxReserveDuration
	}
	return 0
}

func resolveExe(path string) (string, error) {
	if filepath.Ext(path) != "" {
		return path, nil
	}

	lme := errors.NewLazyMultiError(2)
	for i, ext := range []string{".exe", ".bat"} {
		candidate := path + ext
		if _, err := os.Stat(candidate); !lme.Assign(i, err) {
			return candidate, nil
		}
	}

	me := lme.Get().(errors.MultiError)
	return path, errors.Reason("cannot find .exe (%q) or .bat (%q)", me[0], me[1]).Err()
}
<commit_msg>[bbagent] Fix nil proto merge.<commit_after>// Copyright 2019 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.

// Command bbagent is Buildbucket's agent running in swarming.
//
// This executable creates a luciexe 'host' environment, and runs the
// Buildbucket build's exe within this environment. Please see
// https://go.chromium.org/luci/luciexe for details about the 'luciexe'
// protocol.
//
// This command is an implementation detail of Buildbucket.
package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"
	"sync/atomic"
	"time"

	"net/http"
	_ "net/http/pprof"

	"github.com/golang/protobuf/jsonpb"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/fieldmaskpb"
	"google.golang.org/protobuf/types/known/timestamppb"

	"go.chromium.org/luci/buildbucket"
	"go.chromium.org/luci/buildbucket/cmd/bbagent/bbinput"
	bbpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/logging/gologger"
	"go.chromium.org/luci/common/sync/dispatcher"
	"go.chromium.org/luci/common/system/environ"
	"go.chromium.org/luci/grpc/grpcutil"
	"go.chromium.org/luci/lucictx"
	"go.chromium.org/luci/luciexe"
	"go.chromium.org/luci/luciexe/host"
	"go.chromium.org/luci/luciexe/invoke"
)

var maxReserveDuration = 2 * time.Minute

func main() {
	go func() {
		// serves "/debug" endpoints for pprof.
		log.Println(http.ListenAndServe("localhost:6060", nil))
	}()
	os.Exit(mainImpl())
}

func mainImpl() int {
	ctx := logging.SetLevel(gologger.StdConfig.Use(context.Background()), logging.Info)

	check := func(err error) {
		if err != nil {
			logging.Errorf(ctx, err.Error())
			os.Exit(1)
		}
	}

	outputFile := luciexe.AddOutputFlagToSet(flag.CommandLine)
	flag.Parse()
	args := flag.Args()

	if len(args) != 1 {
		check(errors.Reason("expected 1 argument: got %d", len(args)).Err())
	}

	input, err := bbinput.Parse(args[0])
	check(errors.Annotate(err, "could not unmarshal BBAgentArgs").Err())

	sctx, err := lucictx.SwitchLocalAccount(ctx, "system")
	check(errors.Annotate(err, "could not switch to 'system' account in LUCI_CONTEXT").Err())

	// We start with retries disabled because the dispatcher.Channel will handle
	// them during the execution of the user process; In particular we want
	// dispatcher.Channel to be able to move on to a newer version of the Build if
	// it encounters transient errors, rather than retrying a potentially stale
	// Build state.
	//
	// We enable them again after the user process has finished.
	bbclientRetriesEnabled := false
	bbclient, secrets, err := newBuildsClient(sctx, input.Build.Infra.Buildbucket, &bbclientRetriesEnabled)
	check(errors.Annotate(err, "could not connect to Buildbucket").Err())
	logdogOutput, err := mkLogdogOutput(sctx, input.Build.Infra.Logdog)
	check(errors.Annotate(err, "could not create logdog output").Err())

	var (
		cctx   context.Context
		cancel func()
	)
	if dl := lucictx.GetDeadline(ctx); dl.GetSoftDeadline() != 0 {
		softDeadline := dl.SoftDeadlineTime()
		gracePeriod := time.Duration(dl.GetGracePeriod() * float64(time.Second))
		cctx, cancel = context.WithDeadline(ctx, softDeadline.Add(gracePeriod))
	} else {
		cctx, cancel = context.WithCancel(ctx)
	}
	defer cancel()

	dispatcherOpts, dispatcherErrCh := channelOpts(cctx)
	buildsCh, err := dispatcher.NewChannel(cctx, dispatcherOpts, mkSendFn(cctx, secrets, bbclient))
	check(errors.Annotate(err, "could not create builds dispatcher channel").Err())
	defer buildsCh.CloseAndDrain(cctx)

	// from this point forward we want to try to report errors to buildbucket,
	// too.
	check = func(err error) {
		if err != nil {
			logging.Errorf(cctx, err.Error())
			buildsCh.C <- &bbpb.Build{
				Status:          bbpb.Status_INFRA_FAILURE,
				SummaryMarkdown: fmt.Sprintf("fatal error in startup: %s", err),
			}
			buildsCh.CloseAndDrain(cctx)
			os.Exit(1)
		}
	}

	cctx = setResultDBContext(cctx, input.Build, secrets)
	prepareInputBuild(cctx, input.Build)

	opts := &host.Options{
		BaseBuild:      input.Build,
		ButlerLogLevel: logging.Warning,
		ViewerURL: fmt.Sprintf("https://%s/build/%d",
			input.Build.Infra.Buildbucket.Hostname, input.Build.Id),
		LogdogOutput: logdogOutput,
		ExeAuth:      host.DefaultExeAuth("bbagent", input.KnownPublicGerritHosts),
	}
	cwd, err := os.Getwd()
	check(errors.Annotate(err, "getting cwd").Err())
	opts.BaseDir = filepath.Join(cwd, "x")

	exeArgs := append(([]string)(nil), input.Build.Exe.Cmd...)
	payloadPath := input.PayloadPath
	if len(exeArgs) == 0 {
		// TODO(iannucci): delete me with ExecutablePath.
		var exe string
		payloadPath, exe = path.Split(input.ExecutablePath)
		exeArgs = []string{exe}
	}
	exePath, err := filepath.Abs(filepath.Join(payloadPath, exeArgs[0]))
	check(errors.Annotate(err, "absoluting exe path %q", input.ExecutablePath).Err())
	if runtime.GOOS == "windows" {
		exePath, err = resolveExe(exePath)
		check(errors.Annotate(err, "resolving %q", input.ExecutablePath).Err())
	}
	exeArgs[0] = exePath

	initialJSONPB, err := (&jsonpb.Marshaler{
		OrigName: true, Indent: "  ",
	}).MarshalToString(input)
	check(errors.Annotate(err, "marshalling input args").Err())
	logging.Infof(ctx, "Input args:\n%s", initialJSONPB)

	shutdownCh := make(chan struct{})
	var statusDetails *bbpb.StatusDetails
	builds, err := host.Run(cctx, opts, func(ctx context.Context, hostOpts host.Options) error {
		logging.Infof(ctx, "running luciexe: %q", exeArgs)
		logging.Infof(ctx, "  (cache dir): %q", input.CacheDir)
		invokeOpts := &invoke.Options{
			BaseDir:  hostOpts.BaseDir,
			CacheDir: input.CacheDir,
		}
		reserve := calcDeadlineReserve(ctx)
		deadlineEventCh, dctx, shutdown := lucictx.AdjustDeadline(ctx, reserve, 500*time.Millisecond)
		go func() {
			select {
			case <-shutdownCh:
				shutdown()
			case <-dctx.Done():
			}
		}()
		if newSoftDeadline := lucictx.GetDeadline(dctx).SoftDeadlineTime(); !newSoftDeadline.IsZero() {
			logging.Infof(dctx,
				"attempted to reserve %s from soft deadline for bbagent cleanup; new soft deadline: %s",
				reserve, newSoftDeadline)
		}
		subp, err := invoke.Start(dctx, exeArgs, input.Build, invokeOpts, deadlineEventCh)
		if err != nil {
			return err
		}
		var build *bbpb.Build
		build, err = subp.Wait()
		statusDetails = build.StatusDetails
		return err
	})
	if err != nil {
		check(errors.Annotate(err, "could not start luciexe host environment").Err())
	}

	var (
		finalBuild                *bbpb.Build = proto.Clone(input.Build).(*bbpb.Build)
		fatalUpdateBuildErrorSlot atomic.Value
	)

	go func() {
		// Monitors the `dispatcherErrCh` and checks for fatal error determined by
		// `fatalPred`. Stops the build shuttling and shuts down the luciexe if a
		// fatal error is received.
		stopped := false
		fatalPred := func(err error) bool {
			// TODO(crbug.com/1140612): Figure out a better solution to handle the
			// InvalidArgument error at the beginning of the build when build status
			// hasn't switched to STARTED yet.
			// Possible solution: keep calling GetBuild and start shuttling the build
			// once GetBuild returns STARTED status. For now, simply tolerate such
			// error.
			if grpcutil.Code(err) == codes.InvalidArgument && !strings.Contains(err.Error(), "cannot update steps of a SCHEDULED build") {
				return true
			}
			return false
		}
		for {
			select {
			case err := <-dispatcherErrCh:
				if !stopped && fatalPred(err) {
					close(shutdownCh)
					fatalUpdateBuildErrorSlot.Store(err)
					stopped = true
				}
			case <-cctx.Done():
				return
			}
		}
	}()

	// Now all we do is shuttle builds through to the buildbucket client channel
	// until there are no more builds to shuttle.
	for build := range builds {
		if fatalUpdateBuildErrorSlot.Load() == nil {
			buildsCh.C <- build
			finalBuild = build
		}
	}
	buildsCh.CloseAndDrain(cctx)

	// Now that the builds channel has been closed, update bb directly.
	updateMask := []string{
		"build.status",
		"build.status_details",
		"build.summary_markdown",
	}
	var retcode int

	fatalUpdateBuildErr, _ := fatalUpdateBuildErrorSlot.Load().(error)
	if finalizeBuild(ctx, finalBuild, fatalUpdateBuildErr, statusDetails, outputFile) {
		updateMask = append(updateMask, "build.steps", "build.output")
	} else {
		// finalizeBuild indicated that something is really wrong; Omit steps and
		// output from the final push to minimize potential issues.
		retcode = 1
	}
	bbclientRetriesEnabled = true
	_, bbErr := bbclient.UpdateBuild(
		metadata.NewOutgoingContext(cctx, metadata.Pairs(buildbucket.BuildTokenHeader, secrets.BuildToken)),
		&bbpb.UpdateBuildRequest{
			Build:      finalBuild,
			UpdateMask: &fieldmaskpb.FieldMask{Paths: updateMask},
		})
	if bbErr != nil {
		logging.Errorf(cctx, "Failed to report error %s to Buildbucket due to %s", err, bbErr)
		retcode = 2
	}

	return retcode
}

// finalizeBuild returns true if fatalErr is nil and there's no additional
// errors finalizing the build.
func finalizeBuild(ctx context.Context, finalBuild *bbpb.Build, fatalErr error, statusDetails *bbpb.StatusDetails, outputFile *luciexe.OutputFlag) bool {
	if statusDetails != nil {
		if finalBuild.StatusDetails == nil {
			finalBuild.StatusDetails = &bbpb.StatusDetails{}
		}
		proto.Merge(finalBuild.StatusDetails, statusDetails)
	}

	// set final times
	now := timestamppb.New(clock.Now(ctx))
	finalBuild.UpdateTime = now
	finalBuild.EndTime = now

	var finalErrs errors.MultiError
	if fatalErr != nil {
		finalErrs = append(finalErrs, errors.Annotate(fatalErr, "fatal error in builbucket.UpdateBuild").Err())
	}
	if err := outputFile.Write(finalBuild); err != nil {
		finalErrs = append(finalErrs, errors.Annotate(err, "writing final build").Err())
	}

	if len(finalErrs) > 0 {
		errors.Log(ctx, finalErrs)

		// we had some really bad error, just downgrade status and add a message to
		// summary markdown.
		finalBuild.Status = bbpb.Status_INFRA_FAILURE
		originalSM := finalBuild.SummaryMarkdown
		finalBuild.SummaryMarkdown = fmt.Sprintf("FATAL: %s", finalErrs.Error())
		if originalSM != "" {
			finalBuild.SummaryMarkdown += "\n\n" + originalSM
		}
	}

	return len(finalErrs) == 0
}

func prepareInputBuild(ctx context.Context, build *bbpb.Build) {
	// mark started
	build.Status = bbpb.Status_STARTED
	now := timestamppb.New(clock.Now(ctx))
	build.StartTime, build.UpdateTime = now, now
	// TODO(iannucci): this is sketchy, but we preemptively add the log entries
	// for the top level user stdout/stderr streams.
	//
	// Really, `invoke.Start` is the one that knows how to arrange the
	// Output.Logs, but host.Run makes a copy of this build immediately. Find
	// a way to set these up nicely (maybe have opts.BaseBuild be a function
	// returning an immutable bbpb.Build?).
	build.Output = &bbpb.Build_Output{
		Logs: []*bbpb.Log{
			{Name: "stdout", Url: "stdout"},
			{Name: "stderr", Url: "stderr"},
		},
	}
	populateSwarmingInfoFromEnv(build, environ.System())
	return
}

// Returns min(1% of the remaining time towards current soft deadline,
// `maxReserveDuration`). Returns 0 If soft deadline doesn't exist or
// has already been exceeded.
func calcDeadlineReserve(ctx context.Context) time.Duration {
	curSoftDeadline := lucictx.GetDeadline(ctx).SoftDeadlineTime()
	if now := clock.Now(ctx).UTC(); !curSoftDeadline.IsZero() && now.Before(curSoftDeadline) {
		if reserve := curSoftDeadline.Sub(now) / 100; reserve < maxReserveDuration {
			return reserve
		}
		return maxReserveDuration
	}
	return 0
}

func resolveExe(path string) (string, error) {
	if filepath.Ext(path) != "" {
		return path, nil
	}

	lme := errors.NewLazyMultiError(2)
	for i, ext := range []string{".exe", ".bat"} {
		candidate := path + ext
		if _, err := os.Stat(candidate); !lme.Assign(i, err) {
			return candidate, nil
		}
	}

	me := lme.Get().(errors.MultiError)
	return path, errors.Reason("cannot find .exe (%q) or .bat (%q)", me[0], me[1]).Err()
}
<|endoftext|>


<commit_before>package pki

import (
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"fmt"
	"time"

	"github.com/hashicorp/vault/helper/errutil"
	"github.com/hashicorp/vault/logical"
)

type revocationInfo struct {
	CertificateBytes  []byte    `json:"certificate_bytes"`
	RevocationTime    int64     `json:"revocation_time"`
	RevocationTimeUTC time.Time `json:"revocation_time_utc"`
}

// Revokes a cert, and tries to be smart about error recovery
func revokeCert(b *backend, req *logical.Request, serial string, fromLease bool) (*logical.Response, error) {
	// As this backend is self-contained and this function does not hook into
	// third parties to manage users or resources, if the mount is tainted,
	// revocation doesn't matter anyways -- the CRL that would be written will
	// be immediately blown away by the view being cleared. So we can simply
	// fast path a successful exit.
	if b.System().Tainted() {
		return nil, nil
	}

	alreadyRevoked := false
	var revInfo revocationInfo

	certEntry, err := fetchCertBySerial(req, "revoked/", serial)
	if err != nil {
		switch err.(type) {
		case errutil.UserError:
			return logical.ErrorResponse(err.Error()), nil
		case errutil.InternalError:
			return nil, err
		}
	}
	if certEntry != nil {
		// Set the revocation info to the existing values
		alreadyRevoked = true

		revEntry, err := req.Storage.Get("revoked/" + serial)
		if revEntry == nil || err != nil {
			return nil, fmt.Errorf("Error getting existing revocation info")
		}

		err = revEntry.DecodeJSON(&revInfo)
		if err != nil {
			return nil, fmt.Errorf("Error decoding existing revocation info")
		}
	}

	if !alreadyRevoked {
		certEntry, err = fetchCertBySerial(req, "certs/", serial)
		if err != nil {
			switch err.(type) {
			case errutil.UserError:
				return logical.ErrorResponse(err.Error()), nil
			case errutil.InternalError:
				return nil, err
			}
		}
		if certEntry == nil {
			return logical.ErrorResponse(fmt.Sprintf("certificate with serial %s not found", serial)), nil
		}

		cert, err := x509.ParseCertificate(certEntry.Value)
		if err != nil {
			return nil, fmt.Errorf("Error parsing certificate")
		}
		if cert == nil {
			return nil, fmt.Errorf("Got a nil certificate")
		}

		if cert.NotAfter.Before(time.Now()) {
			return nil, nil
		}

		// Compatibility: Don't revoke CAs if they had leases. New CAs going
		// forward aren't issued leases.
		if cert.IsCA && fromLease {
			return nil, nil
		}

		currTime := time.Now()
		revInfo.CertificateBytes = certEntry.Value
		revInfo.RevocationTime = currTime.Unix()
		revInfo.RevocationTimeUTC = currTime.UTC()

		certEntry, err = logical.StorageEntryJSON("revoked/"+serial, revInfo)
		if err != nil {
			return nil, fmt.Errorf("Error creating revocation entry")
		}

		err = req.Storage.Put(certEntry)
		if err != nil {
			return nil, fmt.Errorf("Error saving revoked certificate to new location")
		}

	}

	crlErr := buildCRL(b, req)
	switch crlErr.(type) {
	case errutil.UserError:
		return logical.ErrorResponse(fmt.Sprintf("Error during CRL building: %s", crlErr)), nil
	case errutil.InternalError:
		return nil, fmt.Errorf("Error encountered during CRL building: %s", crlErr)
	}

	resp := &logical.Response{
		Data: map[string]interface{}{
			"revocation_time": revInfo.RevocationTime,
		},
	}
	if !revInfo.RevocationTimeUTC.IsZero() {
		resp.Data["revocation_time_rfc3339"] = revInfo.RevocationTimeUTC.Format(time.RFC3339Nano)
	}
	return resp, nil
}

// Builds a CRL by going through the list of revoked certificates and building
// a new CRL with the stored revocation times and serial numbers.
func buildCRL(b *backend, req *logical.Request) error {
	revokedSerials, err := req.Storage.List("revoked/")
	if err != nil {
		return errutil.InternalError{Err: fmt.Sprintf("Error fetching list of revoked certs: %s", err)}
	}

	revokedCerts := []pkix.RevokedCertificate{}
	var revInfo revocationInfo
	for _, serial := range revokedSerials {
		revokedEntry, err := req.Storage.Get("revoked/" + serial)
		if err != nil {
			return errutil.InternalError{Err: fmt.Sprintf("Unable to fetch revoked cert with serial %s: %s", serial, err)}
		}
		if revokedEntry == nil {
			return errutil.InternalError{Err: fmt.Sprintf("Revoked certificate entry for serial %s is nil", serial)}
		}
		if revokedEntry.Value == nil || len(revokedEntry.Value) == 0 {
			// TODO: In this case, remove it and continue? How likely is this to
			// happen? Alternately, could skip it entirely, or could implement a
			// delete function so that there is a way to remove these
			return errutil.InternalError{Err: fmt.Sprintf("Found revoked serial but actual certificate is empty")}
		}

		err = revokedEntry.DecodeJSON(&revInfo)
		if err != nil {
			return errutil.InternalError{Err: fmt.Sprintf("Error decoding revocation entry for serial %s: %s", serial, err)}
		}

		revokedCert, err := x509.ParseCertificate(revInfo.CertificateBytes)
		if err != nil {
			return errutil.InternalError{Err: fmt.Sprintf("Unable to parse stored revoked certificate with serial %s: %s", serial, err)}
		}

		// NOTE: We have to change this to UTC time because the CRL standard
		// mandates it but Go will happily encode the CRL without this.
		newRevCert := pkix.RevokedCertificate{
			SerialNumber: revokedCert.SerialNumber,
		}
		if !revInfo.RevocationTimeUTC.IsZero() {
			newRevCert.RevocationTime = revInfo.RevocationTimeUTC
		} else {
			newRevCert.RevocationTime = time.Unix(revInfo.RevocationTime, 0).UTC()
		}
		revokedCerts = append(revokedCerts, newRevCert)
	}

	signingBundle, caErr := fetchCAInfo(req)
	switch caErr.(type) {
	case errutil.UserError:
		return errutil.UserError{Err: fmt.Sprintf("Could not fetch the CA certificate: %s", caErr)}
	case errutil.InternalError:
		return errutil.InternalError{Err: fmt.Sprintf("Error fetching CA certificate: %s", caErr)}
	}

	crlLifetime := b.crlLifetime
	crlInfo, err := b.CRL(req.Storage)
	if err != nil {
		return errutil.InternalError{Err: fmt.Sprintf("Error fetching CRL config information: %s", err)}
	}
	if crlInfo != nil {
		crlDur, err := time.ParseDuration(crlInfo.Expiry)
		if err != nil {
			return errutil.InternalError{Err: fmt.Sprintf("Error parsing CRL duration of %s", crlInfo.Expiry)}
		}
		crlLifetime = crlDur
	}

	crlBytes, err := signingBundle.Certificate.CreateCRL(rand.Reader, signingBundle.PrivateKey, revokedCerts, time.Now(), time.Now().Add(crlLifetime))
	if err != nil {
		return errutil.InternalError{Err: fmt.Sprintf("Error creating new CRL: %s", err)}
	}

	err = req.Storage.Put(&logical.StorageEntry{
		Key:   "crl",
		Value: crlBytes,
	})
	if err != nil {
		return errutil.InternalError{Err: fmt.Sprintf("Error storing CRL: %s", err)}
	}

	return nil
}
<commit_msg>pki: Avoiding a storage read<commit_after>package pki

import (
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"fmt"
	"time"

	"github.com/hashicorp/vault/helper/errutil"
	"github.com/hashicorp/vault/logical"
)

type revocationInfo struct {
	CertificateBytes  []byte    `json:"certificate_bytes"`
	RevocationTime    int64     `json:"revocation_time"`
	RevocationTimeUTC time.Time `json:"revocation_time_utc"`
}

// Revokes a cert, and tries to be smart about error recovery
func revokeCert(b *backend, req *logical.Request, serial string, fromLease bool) (*logical.Response, error) {
	// As this backend is self-contained and this function does not hook into
	// third parties to manage users or resources, if the mount is tainted,
	// revocation doesn't matter anyways -- the CRL that would be written will
	// be immediately blown away by the view being cleared. So we can simply
	// fast path a successful exit.
	if b.System().Tainted() {
		return nil, nil
	}

	alreadyRevoked := false
	var revInfo revocationInfo

	revEntry, err := fetchCertBySerial(req, "revoked/", serial)
	if err != nil {
		switch err.(type) {
		case errutil.UserError:
			return logical.ErrorResponse(err.Error()), nil
		case errutil.InternalError:
			return nil, err
		}
	}
	if revEntry != nil {
		// Set the revocation info to the existing values
		alreadyRevoked = true
		err = revEntry.DecodeJSON(&revInfo)
		if err != nil {
			return nil, fmt.Errorf("Error decoding existing revocation info")
		}
	}

	if !alreadyRevoked {
		certEntry, err := fetchCertBySerial(req, "certs/", serial)
		if err != nil {
			switch err.(type) {
			case errutil.UserError:
				return logical.ErrorResponse(err.Error()), nil
			case errutil.InternalError:
				return nil, err
			}
		}
		if certEntry == nil {
			return logical.ErrorResponse(fmt.Sprintf("certificate with serial %s not found", serial)), nil
		}

		cert, err := x509.ParseCertificate(certEntry.Value)
		if err != nil {
			return nil, fmt.Errorf("Error parsing certificate")
		}
		if cert == nil {
			return nil, fmt.Errorf("Got a nil certificate")
		}

		if cert.NotAfter.Before(time.Now()) {
			return nil, nil
		}

		// Compatibility: Don't revoke CAs if they had leases. New CAs going
		// forward aren't issued leases.
		if cert.IsCA && fromLease {
			return nil, nil
		}

		currTime := time.Now()
		revInfo.CertificateBytes = certEntry.Value
		revInfo.RevocationTime = currTime.Unix()
		revInfo.RevocationTimeUTC = currTime.UTC()

		revEntry, err := logical.StorageEntryJSON("revoked/"+serial, revInfo)
		if err != nil {
			return nil, fmt.Errorf("Error creating revocation entry")
		}

		err = req.Storage.Put(revEntry)
		if err != nil {
			return nil, fmt.Errorf("Error saving revoked certificate to new location")
		}

	}

	crlErr := buildCRL(b, req)
	switch crlErr.(type) {
	case errutil.UserError:
		return logical.ErrorResponse(fmt.Sprintf("Error during CRL building: %s", crlErr)), nil
	case errutil.InternalError:
		return nil, fmt.Errorf("Error encountered during CRL building: %s", crlErr)
	}

	resp := &logical.Response{
		Data: map[string]interface{}{
			"revocation_time": revInfo.RevocationTime,
		},
	}
	if !revInfo.RevocationTimeUTC.IsZero() {
		resp.Data["revocation_time_rfc3339"] = revInfo.RevocationTimeUTC.Format(time.RFC3339Nano)
	}
	return resp, nil
}

// Builds a CRL by going through the list of revoked certificates and building
// a new CRL with the stored revocation times and serial numbers.
func buildCRL(b *backend, req *logical.Request) error {
	revokedSerials, err := req.Storage.List("revoked/")
	if err != nil {
		return errutil.InternalError{Err: fmt.Sprintf("Error fetching list of revoked certs: %s", err)}
	}

	revokedCerts := []pkix.RevokedCertificate{}
	var revInfo revocationInfo
	for _, serial := range revokedSerials {
		revokedEntry, err := req.Storage.Get("revoked/" + serial)
		if err != nil {
			return errutil.InternalError{Err: fmt.Sprintf("Unable to fetch revoked cert with serial %s: %s", serial, err)}
		}
		if revokedEntry == nil {
			return errutil.InternalError{Err: fmt.Sprintf("Revoked certificate entry for serial %s is nil", serial)}
		}
		if revokedEntry.Value == nil || len(revokedEntry.Value) == 0 {
			// TODO: In this case, remove it and continue? How likely is this to
			// happen? Alternately, could skip it entirely, or could implement a
			// delete function so that there is a way to remove these
			return errutil.InternalError{Err: fmt.Sprintf("Found revoked serial but actual certificate is empty")}
		}

		err = revokedEntry.DecodeJSON(&revInfo)
		if err != nil {
			return errutil.InternalError{Err: fmt.Sprintf("Error decoding revocation entry for serial %s: %s", serial, err)}
		}

		revokedCert, err := x509.ParseCertificate(revInfo.CertificateBytes)
		if err != nil {
			return errutil.InternalError{Err: fmt.Sprintf("Unable to parse stored revoked certificate with serial %s: %s", serial, err)}
		}

		// NOTE: We have to change this to UTC time because the CRL standard
		// mandates it but Go will happily encode the CRL without this.
		newRevCert := pkix.RevokedCertificate{
			SerialNumber: revokedCert.SerialNumber,
		}
		if !revInfo.RevocationTimeUTC.IsZero() {
			newRevCert.RevocationTime = revInfo.RevocationTimeUTC
		} else {
			newRevCert.RevocationTime = time.Unix(revInfo.RevocationTime, 0).UTC()
		}
		revokedCerts = append(revokedCerts, newRevCert)
	}

	signingBundle, caErr := fetchCAInfo(req)
	switch caErr.(type) {
	case errutil.UserError:
		return errutil.UserError{Err: fmt.Sprintf("Could not fetch the CA certificate: %s", caErr)}
	case errutil.InternalError:
		return errutil.InternalError{Err: fmt.Sprintf("Error fetching CA certificate: %s", caErr)}
	}

	crlLifetime := b.crlLifetime
	crlInfo, err := b.CRL(req.Storage)
	if err != nil {
		return errutil.InternalError{Err: fmt.Sprintf("Error fetching CRL config information: %s", err)}
	}
	if crlInfo != nil {
		crlDur, err := time.ParseDuration(crlInfo.Expiry)
		if err != nil {
			return errutil.InternalError{Err: fmt.Sprintf("Error parsing CRL duration of %s", crlInfo.Expiry)}
		}
		crlLifetime = crlDur
	}

	crlBytes, err := signingBundle.Certificate.CreateCRL(rand.Reader, signingBundle.PrivateKey, revokedCerts, time.Now(), time.Now().Add(crlLifetime))
	if err != nil {
		return errutil.InternalError{Err: fmt.Sprintf("Error creating new CRL: %s", err)}
	}

	err = req.Storage.Put(&logical.StorageEntry{
		Key:   "crl",
		Value: crlBytes,
	})
	if err != nil {
		return errutil.InternalError{Err: fmt.Sprintf("Error storing CRL: %s", err)}
	}

	return nil
}
<|endoftext|>