<commit_before>package intercept

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"errors"
	"github.com/nbio/st"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"
)

var errRead = errors.New("read error")

type errorReader struct{}

func (r *errorReader) Read(p []byte) (int, error) {
	return 0, errRead
}

type user struct {
	XMLName xml.Name `xml:"Person" json:"-"`
	Name    string
}

func TestNewRequestModifier(t *testing.T) {
	h := http.Header{}
	h.Set("foo", "bar")
	req := &http.Request{Header: h}
	modifier := NewRequestModifier(req)
	st.Expect(t, modifier.Request, req)
	st.Expect(t, modifier.Header, h)
}

func TestReadString(t *testing.T) {
	bodyStr := `{"name":"Rick"}`
	strReader := strings.NewReader(bodyStr)
	body := ioutil.NopCloser(strReader)
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	str, err := modifier.ReadString()
	st.Expect(t, err, nil)
	st.Expect(t, str, bodyStr)
}

func TestReadStringError(t *testing.T) {
	body := ioutil.NopCloser(&errorReader{})
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	str, err := modifier.ReadString()
	st.Expect(t, err, errRead)
	st.Expect(t, str, "")
}

func TestReadBytes(t *testing.T) {
	bodyBytes := []byte(`{"name":"Rick"}`)
	strReader := bytes.NewBuffer(bodyBytes)
	body := ioutil.NopCloser(strReader)
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	str, err := modifier.ReadBytes()
	st.Expect(t, err, nil)
	st.Expect(t, str, bodyBytes)
}

func TestReadBytesError(t *testing.T) {
	body := ioutil.NopCloser(&errorReader{})
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	buf, err := modifier.ReadBytes()
	st.Expect(t, err, errRead)
	st.Expect(t, len(buf), 0)
}

func TestDecodeJSON(t *testing.T) {
	bodyBytes := []byte(`{"name":"Rick"}`)
	strReader := bytes.NewBuffer(bodyBytes)
	body := ioutil.NopCloser(strReader)
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	u := user{}
	err := modifier.DecodeJSON(&u)
	st.Expect(t, err, nil)
	st.Expect(t, u.Name, "Rick")
}

func TestDecodeJSONErrorFromReadBytes(t *testing.T) {
	body := ioutil.NopCloser(&errorReader{})
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	u := user{}
	err := modifier.DecodeJSON(&u)
	st.Expect(t, err, errRead)
	st.Expect(t, u.Name, "")
}

func TestDecodeJSONEOF(t *testing.T) {
	bodyBytes := []byte("")
	strReader := bytes.NewBuffer(bodyBytes)
	body := ioutil.NopCloser(strReader)
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	u := user{}
	err := modifier.DecodeJSON(&u)
	st.Expect(t, err, nil)
	st.Expect(t, u.Name, "")
}

func TestDecodeJSONErrorFromDecode(t *testing.T) {
	bodyBytes := []byte(`/`)
	strReader := bytes.NewBuffer(bodyBytes)
	body := ioutil.NopCloser(strReader)
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	u := user{}
	err := modifier.DecodeJSON(&u)
	_, ok := (err).(*json.SyntaxError)
	st.Expect(t, ok, true)
	st.Expect(t, err.Error(), "invalid character '/' looking for beginning of value")
	st.Expect(t, u.Name, "")
}

func TestDecodeXML(t *testing.T) {
	bodyBytes := []byte(`<Person><Name>Rick</Name></Person>`)
	strReader := bytes.NewBuffer(bodyBytes)
	body := ioutil.NopCloser(strReader)
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	u := user{}
	err := modifier.DecodeXML(&u, nil)
	st.Expect(t, err, nil)
	st.Expect(t, u.Name, "Rick")
}

func TestDecodeXMLErrorFromReadBytes(t *testing.T) {
	body := ioutil.NopCloser(&errorReader{})
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	u := user{}
	err := modifier.DecodeXML(&u, nil)
	st.Expect(t, err, errRead)
	st.Expect(t, u.Name, "")
}

func TestDecodeXMLErrorFromDecode(t *testing.T) {
	bodyBytes := []byte(`]]>`)
	strReader := bytes.NewBuffer(bodyBytes)
	body := ioutil.NopCloser(strReader)
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	u := user{}
	err := modifier.DecodeXML(&u, nil)
	_, ok := (err).(*xml.SyntaxError)
	st.Expect(t, ok, true)
	st.Expect(t, err.Error(), "XML syntax error on line 1: unescaped ]]> not in CDATA section")
	st.Expect(t, u.Name, "")
}

func TestDecodeXMLEOF(t *testing.T) {
	bodyBytes := []byte("")
	strReader := bytes.NewBuffer(bodyBytes)
	body := ioutil.NopCloser(strReader)
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	u := user{}
	err := modifier.DecodeXML(&u, nil)
	st.Expect(t, err, nil)
	st.Expect(t, u.Name, "")
}

func TestBytes(t *testing.T) {
	req := &http.Request{Header: http.Header{}}
	modifier := NewRequestModifier(req)
	modifier.Bytes([]byte("hello"))
	modifiedBody, err := ioutil.ReadAll(req.Body)
	st.Expect(t, err, nil)
	st.Expect(t, string(modifiedBody), "hello")
}

func TestStringGet(t *testing.T) {
	req := &http.Request{Method: "GET", Header: http.Header{}}
	modifier := NewRequestModifier(req)
	modifier.String("hello")
	st.Expect(t, req.Body, nil)
}

func TestStringHead(t *testing.T) {
	req := &http.Request{Method: "HEAD", Header: http.Header{}}
	modifier := NewRequestModifier(req)
	modifier.String("hello")
	st.Expect(t, req.Body, nil)
}

func TestString(t *testing.T) {
	req := &http.Request{Method: "POST", Header: http.Header{}}
	modifier := NewRequestModifier(req)
	modifier.String("hello")
	modifiedBody, err := ioutil.ReadAll(req.Body)
	st.Expect(t, err, nil)
	st.Expect(t, string(modifiedBody), "hello")
}

func TestJSONWithStructAsParameter(t *testing.T) {
	req := &http.Request{Header: http.Header{}}
	modifier := NewRequestModifier(req)
	u := &user{Name: "Rick"}
	err := modifier.JSON(u)
	st.Expect(t, err, nil)
	modifiedBody, err := ioutil.ReadAll(req.Body)
	st.Expect(t, err, nil)
	expectedBody := "{\"Name\":\"Rick\"}\n"
	st.Expect(t, string(modifiedBody), expectedBody)
	st.Expect(t, req.ContentLength, int64(len(expectedBody)))
	st.Expect(t, req.Header.Get("Content-Type"), "application/json")
}

func TestJSONWithStringAsParameter(t *testing.T) {
	req := &http.Request{Header: http.Header{}}
	modifier := NewRequestModifier(req)
	input := `{"Name":"Rick"}`
	err := modifier.JSON(input)
	st.Expect(t, err, nil)
	modifiedBody, err := ioutil.ReadAll(req.Body)
	st.Expect(t, err, nil)
	st.Expect(t, string(modifiedBody), input)
	st.Expect(t, req.ContentLength, int64(len(input)))
	st.Expect(t, req.Header.Get("Content-Type"), "application/json")
}

func TestJSONWithBytesAsParameter(t *testing.T) {
	req := &http.Request{Header: http.Header{}}
	modifier := NewRequestModifier(req)
	input := []byte(`{"Name":"Rick"}`)
	err := modifier.JSON(input)
	st.Expect(t, err, nil)
	modifiedBody, err := ioutil.ReadAll(req.Body)
	st.Expect(t, err, nil)
	st.Expect(t, modifiedBody, input)
	st.Expect(t, req.ContentLength, int64(len(input)))
	st.Expect(t, req.Header.Get("Content-Type"), "application/json")
}

func TestJSONEncodingError(t *testing.T) {
	req := &http.Request{Header: http.Header{}}
	modifier := NewRequestModifier(req)
	input := make(map[int]int)
	err := modifier.JSON(input)
	_, ok := err.(*json.UnsupportedTypeError)
	st.Expect(t, ok, true)
	st.Expect(t, err.Error(), "json: unsupported type: map[int]int")
}
<commit_msg>Add tests for RequestModifier.XML.<commit_after>package intercept

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"errors"
	"github.com/nbio/st"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"
)

var errRead = errors.New("read error")

type errorReader struct{}

func (r *errorReader) Read(p []byte) (int, error) {
	return 0, errRead
}

type user struct {
	XMLName xml.Name `xml:"Person" json:"-"`
	Name    string
}

func TestNewRequestModifier(t *testing.T) {
	h := http.Header{}
	h.Set("foo", "bar")
	req := &http.Request{Header: h}
	modifier := NewRequestModifier(req)
	st.Expect(t, modifier.Request, req)
	st.Expect(t, modifier.Header, h)
}

func TestReadString(t *testing.T) {
	bodyStr := `{"name":"Rick"}`
	strReader := strings.NewReader(bodyStr)
	body := ioutil.NopCloser(strReader)
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	str, err := modifier.ReadString()
	st.Expect(t, err, nil)
	st.Expect(t, str, bodyStr)
}

func TestReadStringError(t *testing.T) {
	body := ioutil.NopCloser(&errorReader{})
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	str, err := modifier.ReadString()
	st.Expect(t, err, errRead)
	st.Expect(t, str, "")
}

func TestReadBytes(t *testing.T) {
	bodyBytes := []byte(`{"name":"Rick"}`)
	strReader := bytes.NewBuffer(bodyBytes)
	body := ioutil.NopCloser(strReader)
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	str, err := modifier.ReadBytes()
	st.Expect(t, err, nil)
	st.Expect(t, str, bodyBytes)
}

func TestReadBytesError(t *testing.T) {
	body := ioutil.NopCloser(&errorReader{})
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	buf, err := modifier.ReadBytes()
	st.Expect(t, err, errRead)
	st.Expect(t, len(buf), 0)
}

func TestDecodeJSON(t *testing.T) {
	bodyBytes := []byte(`{"name":"Rick"}`)
	strReader := bytes.NewBuffer(bodyBytes)
	body := ioutil.NopCloser(strReader)
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	u := user{}
	err := modifier.DecodeJSON(&u)
	st.Expect(t, err, nil)
	st.Expect(t, u.Name, "Rick")
}

func TestDecodeJSONErrorFromReadBytes(t *testing.T) {
	body := ioutil.NopCloser(&errorReader{})
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	u := user{}
	err := modifier.DecodeJSON(&u)
	st.Expect(t, err, errRead)
	st.Expect(t, u.Name, "")
}

func TestDecodeJSONEOF(t *testing.T) {
	bodyBytes := []byte("")
	strReader := bytes.NewBuffer(bodyBytes)
	body := ioutil.NopCloser(strReader)
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	u := user{}
	err := modifier.DecodeJSON(&u)
	st.Expect(t, err, nil)
	st.Expect(t, u.Name, "")
}

func TestDecodeJSONErrorFromDecode(t *testing.T) {
	bodyBytes := []byte(`/`)
	strReader := bytes.NewBuffer(bodyBytes)
	body := ioutil.NopCloser(strReader)
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	u := user{}
	err := modifier.DecodeJSON(&u)
	_, ok := (err).(*json.SyntaxError)
	st.Expect(t, ok, true)
	st.Expect(t, err.Error(), "invalid character '/' looking for beginning of value")
	st.Expect(t, u.Name, "")
}

func TestDecodeXML(t *testing.T) {
	bodyBytes := []byte(`<Person><Name>Rick</Name></Person>`)
	strReader := bytes.NewBuffer(bodyBytes)
	body := ioutil.NopCloser(strReader)
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	u := user{}
	err := modifier.DecodeXML(&u, nil)
	st.Expect(t, err, nil)
	st.Expect(t, u.Name, "Rick")
}

func TestDecodeXMLErrorFromReadBytes(t *testing.T) {
	body := ioutil.NopCloser(&errorReader{})
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	u := user{}
	err := modifier.DecodeXML(&u, nil)
	st.Expect(t, err, errRead)
	st.Expect(t, u.Name, "")
}

func TestDecodeXMLErrorFromDecode(t *testing.T) {
	bodyBytes := []byte(`]]>`)
	strReader := bytes.NewBuffer(bodyBytes)
	body := ioutil.NopCloser(strReader)
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	u := user{}
	err := modifier.DecodeXML(&u, nil)
	_, ok := (err).(*xml.SyntaxError)
	st.Expect(t, ok, true)
	st.Expect(t, err.Error(), "XML syntax error on line 1: unescaped ]]> not in CDATA section")
	st.Expect(t, u.Name, "")
}

func TestDecodeXMLEOF(t *testing.T) {
	bodyBytes := []byte("")
	strReader := bytes.NewBuffer(bodyBytes)
	body := ioutil.NopCloser(strReader)
	req := &http.Request{Header: http.Header{}, Body: body}
	modifier := NewRequestModifier(req)
	u := user{}
	err := modifier.DecodeXML(&u, nil)
	st.Expect(t, err, nil)
	st.Expect(t, u.Name, "")
}

func TestBytes(t *testing.T) {
	req := &http.Request{Header: http.Header{}}
	modifier := NewRequestModifier(req)
	modifier.Bytes([]byte("hello"))
	modifiedBody, err := ioutil.ReadAll(req.Body)
	st.Expect(t, err, nil)
	st.Expect(t, string(modifiedBody), "hello")
}

func TestStringGet(t *testing.T) {
	req := &http.Request{Method: "GET", Header: http.Header{}}
	modifier := NewRequestModifier(req)
	modifier.String("hello")
	st.Expect(t, req.Body, nil)
}

func TestStringHead(t *testing.T) {
	req := &http.Request{Method: "HEAD", Header: http.Header{}}
	modifier := NewRequestModifier(req)
	modifier.String("hello")
	st.Expect(t, req.Body, nil)
}

func TestString(t *testing.T) {
	req := &http.Request{Method: "POST", Header: http.Header{}}
	modifier := NewRequestModifier(req)
	modifier.String("hello")
	modifiedBody, err := ioutil.ReadAll(req.Body)
	st.Expect(t, err, nil)
	st.Expect(t, string(modifiedBody), "hello")
}

func TestJSONWithStructAsParameter(t *testing.T) {
	req := &http.Request{Header: http.Header{}}
	modifier := NewRequestModifier(req)
	u := &user{Name: "Rick"}
	err := modifier.JSON(u)
	st.Expect(t, err, nil)
	modifiedBody, err := ioutil.ReadAll(req.Body)
	st.Expect(t, err, nil)
	expectedBody := "{\"Name\":\"Rick\"}\n"
	st.Expect(t, string(modifiedBody), expectedBody)
	st.Expect(t, req.ContentLength, int64(len(expectedBody)))
	st.Expect(t, req.Header.Get("Content-Type"), "application/json")
}

func TestJSONWithStringAsParameter(t *testing.T) {
	req := &http.Request{Header: http.Header{}}
	modifier := NewRequestModifier(req)
	input := `{"Name":"Rick"}`
	err := modifier.JSON(input)
	st.Expect(t, err, nil)
	modifiedBody, err := ioutil.ReadAll(req.Body)
	st.Expect(t, err, nil)
	st.Expect(t, string(modifiedBody), input)
	st.Expect(t, req.ContentLength, int64(len(input)))
	st.Expect(t, req.Header.Get("Content-Type"), "application/json")
}

func TestJSONWithBytesAsParameter(t *testing.T) {
	req := &http.Request{Header: http.Header{}}
	modifier := NewRequestModifier(req)
	input := []byte(`{"Name":"Rick"}`)
	err := modifier.JSON(input)
	st.Expect(t, err, nil)
	modifiedBody, err := ioutil.ReadAll(req.Body)
	st.Expect(t, err, nil)
	st.Expect(t, modifiedBody, input)
	st.Expect(t, req.ContentLength, int64(len(input)))
	st.Expect(t, req.Header.Get("Content-Type"), "application/json")
}

func TestJSONEncodingError(t *testing.T) {
	req := &http.Request{Header: http.Header{}}
	modifier := NewRequestModifier(req)
	input := make(map[int]int)
	err := modifier.JSON(input)
	_, ok := err.(*json.UnsupportedTypeError)
	st.Expect(t, ok, true)
	st.Expect(t, err.Error(), "json: unsupported type: map[int]int")
}

func TestXMLWithStructAsParameter(t *testing.T) {
	req := &http.Request{Header: http.Header{}}
	modifier := NewRequestModifier(req)
	u := &user{Name: "Rick"}
	err := modifier.XML(u)
	st.Expect(t, err, nil)
	modifiedBody, err := ioutil.ReadAll(req.Body)
	st.Expect(t, err, nil)
	expectedBody := `<Person><Name>Rick</Name></Person>`
	st.Expect(t, string(modifiedBody), expectedBody)
	st.Expect(t, req.ContentLength, int64(len(expectedBody)))
	st.Expect(t, req.Header.Get("Content-Type"), "application/xml")
}

func TestXMLWithStringAsParameter(t *testing.T) {
	req := &http.Request{Header: http.Header{}}
	modifier := NewRequestModifier(req)
	input := `<Person><Name>Rick</Name></Person>`
	err := modifier.XML(input)
	st.Expect(t, err, nil)
	modifiedBody, err := ioutil.ReadAll(req.Body)
	st.Expect(t, err, nil)
	st.Expect(t, string(modifiedBody), input)
	st.Expect(t, req.ContentLength, int64(len(input)))
	st.Expect(t, req.Header.Get("Content-Type"), "application/xml")
}

func TestXMLWithBytesAsParameter(t *testing.T) {
	req := &http.Request{Header: http.Header{}}
	modifier := NewRequestModifier(req)
	input := []byte(`<Person><Name>Rick</Name></Person>`)
	err := modifier.XML(input)
	st.Expect(t, err, nil)
	modifiedBody, err := ioutil.ReadAll(req.Body)
	st.Expect(t, err, nil)
	st.Expect(t, string(modifiedBody), string(input))
	st.Expect(t, req.ContentLength, int64(len(input)))
	st.Expect(t, req.Header.Get("Content-Type"), "application/xml")
}

func TestXMLEncodingError(t *testing.T) {
	req := &http.Request{Header: http.Header{}}
	modifier := NewRequestModifier(req)
	input := make(map[int]int)
	err := modifier.XML(input)
	_, ok := err.(*xml.UnsupportedTypeError)
	st.Expect(t, ok, true)
	st.Expect(t, err.Error(), "xml: unsupported type: map[int]int")
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/spf13/cobra"

	"github.com/ovh/cds/cli"
)

var adminDatabaseCmd = cli.Command{
	Name:  "database",
	Short: "Manage CDS Database",
}

func adminDatabase() *cobra.Command {
	return cli.NewCommand(adminDatabaseCmd, nil, []*cobra.Command{
		cli.NewCommand(adminDatabaseUnlockCmd, adminDatabaseUnlockFunc, nil),
		cli.NewCommand(adminDatabaseDeleteMigrationCmd, adminDatabaseDeleteFunc, nil),
		cli.NewListCommand(adminDatabaseMigrationsList, adminDatabaseMigrationsListFunc, nil),
		cli.NewGetCommand(adminDatabaseSignatureResume, adminDatabaseSignatureResumeFunc, nil),
		cli.NewCommand(adminDatabaseSignatureRoll, adminDatabaseSignatureRollFunc, nil),
		cli.NewGetCommand(adminDatabaseEncryptionResume, adminDatabaseEncryptionResumeFunc, nil),
		cli.NewCommand(adminDatabaseEncryptionRoll, adminDatabaseEncryptionRollFunc, nil),
	})
}

var adminDatabaseUnlockCmd = cli.Command{
	Name:  "unlock",
	Short: "Unlock a pending migration (Use with caution)",
	Args: []cli.Arg{
		{Name: "id"},
	},
}

func adminDatabaseUnlockFunc(v cli.Values) error {
	return client.AdminDatabaseMigrationUnlock(v.GetString("id"))
}

var adminDatabaseDeleteMigrationCmd = cli.Command{
	Name:  "delete",
	Short: "Delete a database migration from table gorp_migration (use with caution)",
	Args: []cli.Arg{
		{Name: "id"},
	},
}

func adminDatabaseDeleteFunc(v cli.Values) error {
	return client.AdminDatabaseMigrationDelete(v.GetString("id"))
}

var adminDatabaseMigrationsList = cli.Command{
	Name:  "list",
	Short: "List all CDS DB migrations",
}

func adminDatabaseMigrationsListFunc(_ cli.Values) (cli.ListResult, error) {
	migrations, err := client.AdminDatabaseMigrationsList()
	if err != nil {
		return nil, err
	}
	return cli.AsListResult(migrations), nil
}

var adminDatabaseSignatureResume = cli.Command{
	Name:  "list-signed-data",
	Short: "List all signed data in database",
}

func adminDatabaseSignatureResumeFunc(_ cli.Values) (interface{}, error) {
	return client.AdminDatabaseSignaturesResume()
}

var adminDatabaseSignatureRoll = cli.Command{
	Name:  "roll-signed-data",
	Short: "Roll a signed data in database",
	VariadicArgs: cli.Arg{
		Name: "entity",
	},
}

func adminDatabaseSignatureRollFunc(args cli.Values) error {

	entities := args.GetStringSlice("entity")
	if len(entities) == 0 {
		return client.AdminDatabaseSignaturesRollAllEntities()
	}

	for _, e := range entities {
		if err := client.AdminDatabaseSignaturesRollEntity(e); err != nil {
			return err
		}
	}

	return nil

}

var adminDatabaseEncryptionResume = cli.Command{
	Name:  "list-encrypted-data",
	Short: "List all encrypted data in database",
}

func adminDatabaseEncryptionResumeFunc(_ cli.Values) (interface{}, error) {
	return client.AdminDatabaseListEncryptedEntities()
}

var adminDatabaseEncryptionRoll = cli.Command{
	Name:  "roll-encrypteddata",
	Short: "Roll a encrypted data in database",
	VariadicArgs: cli.Arg{
		Name: "entity",
	},
}

func adminDatabaseEncryptionRollFunc(args cli.Values) error {
	entities := args.GetStringSlice("entity")
	if len(entities) == 0 {
		return client.AdminDatabaseRollAllEncryptedEntities()
	}

	for _, e := range entities {
		if err := client.AdminDatabaseRollEncryptedEntity(e); err != nil {
			return err
		}
	}

	return nil

}
<commit_msg>fix(cdsctl): display list-encrypted-data (#5215)<commit_after>package main

import (
	"fmt"

	"github.com/spf13/cobra"

	"github.com/ovh/cds/cli"
)

var adminDatabaseCmd = cli.Command{
	Name:  "database",
	Short: "Manage CDS Database",
}

func adminDatabase() *cobra.Command {
	return cli.NewCommand(adminDatabaseCmd, nil, []*cobra.Command{
		cli.NewCommand(adminDatabaseUnlockCmd, adminDatabaseUnlockFunc, nil),
		cli.NewCommand(adminDatabaseDeleteMigrationCmd, adminDatabaseDeleteFunc, nil),
		cli.NewListCommand(adminDatabaseMigrationsList, adminDatabaseMigrationsListFunc, nil),
		cli.NewGetCommand(adminDatabaseSignatureResume, adminDatabaseSignatureResumeFunc, nil),
		cli.NewCommand(adminDatabaseSignatureRoll, adminDatabaseSignatureRollFunc, nil),
		cli.NewCommand(adminDatabaseEncryptionResume, adminDatabaseEncryptionResumeFunc, nil),
		cli.NewCommand(adminDatabaseEncryptionRoll, adminDatabaseEncryptionRollFunc, nil),
	})
}

var adminDatabaseUnlockCmd = cli.Command{
	Name:  "unlock",
	Short: "Unlock a pending migration (Use with caution)",
	Args: []cli.Arg{
		{Name: "id"},
	},
}

func adminDatabaseUnlockFunc(v cli.Values) error {
	return client.AdminDatabaseMigrationUnlock(v.GetString("id"))
}

var adminDatabaseDeleteMigrationCmd = cli.Command{
	Name:  "delete",
	Short: "Delete a database migration from table gorp_migration (use with caution)",
	Args: []cli.Arg{
		{Name: "id"},
	},
}

func adminDatabaseDeleteFunc(v cli.Values) error {
	return client.AdminDatabaseMigrationDelete(v.GetString("id"))
}

var adminDatabaseMigrationsList = cli.Command{
	Name:  "list",
	Short: "List all CDS DB migrations",
}

func adminDatabaseMigrationsListFunc(_ cli.Values) (cli.ListResult, error) {
	migrations, err := client.AdminDatabaseMigrationsList()
	if err != nil {
		return nil, err
	}
	return cli.AsListResult(migrations), nil
}

var adminDatabaseSignatureResume = cli.Command{
	Name:  "list-signed-data",
	Short: "List all signed data in database",
}

func adminDatabaseSignatureResumeFunc(_ cli.Values) (interface{}, error) {
	return client.AdminDatabaseSignaturesResume()
}

var adminDatabaseSignatureRoll = cli.Command{
	Name:  "roll-signed-data",
	Short: "Roll a signed data in database",
	VariadicArgs: cli.Arg{
		Name: "entity",
	},
}

func adminDatabaseSignatureRollFunc(args cli.Values) error {
	entities := args.GetStringSlice("entity")
	if len(entities) == 0 {
		return client.AdminDatabaseSignaturesRollAllEntities()
	}

	for _, e := range entities {
		if err := client.AdminDatabaseSignaturesRollEntity(e); err != nil {
			return err
		}
	}

	return nil
}

var adminDatabaseEncryptionResume = cli.Command{
	Name:  "list-encrypted-data",
	Short: "List all encrypted data in database",
}

func adminDatabaseEncryptionResumeFunc(_ cli.Values) error {
	entities, err := client.AdminDatabaseListEncryptedEntities()
	for _, e := range entities {
		fmt.Println(e)
	}
	return err
}

var adminDatabaseEncryptionRoll = cli.Command{
	Name:  "roll-encrypteddata",
	Short: "Roll a encrypted data in database",
	VariadicArgs: cli.Arg{
		Name: "entity",
	},
}

func adminDatabaseEncryptionRollFunc(args cli.Values) error {
	entities := args.GetStringSlice("entity")
	if len(entities) == 0 {
		return client.AdminDatabaseRollAllEncryptedEntities()
	}

	for _, e := range entities {
		if err := client.AdminDatabaseRollEncryptedEntity(e); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package commands

import (
	"flag"
	"fmt"
	"net/rpc"
	"os"
	"time"
)

// Status represents the Status of a node in the cluser
type Status struct {
	CRole     string    // the nodes 'role' in the cluster (primary, secondary, monitor)
	DataDir   string    // directory of the postgres database
	DBRole    string    // the 'role' of the running pgsql instance inside the node (master, slave)
	Ip        string    // advertise_ip
	PGPort    int       //
	State     string    // the current state of the node
	UpdatedAt time.Time // the last time the node state was updated
}

// ClusterListCommand satisfies the Command interface for listing nodes in yoke
type ClusterListCommand struct{}

// Help prints detailed help text for the cluster list command
func (c *ClusterListCommand) Help() {
	fmt.Printf(`
Description:
  Returns status information for all nodes in the cluster

Usage:
  cli list
  cli cluster:list

  ex. cli list
  `)
}

// Run displays select information about all of the nodes in a cluster
func (c *ClusterListCommand) Run(opts []string) {

	// flags
	flags := flag.NewFlagSet("flags", flag.ContinueOnError)
	flags.Usage = func() { c.Help() }

	var fHost string
	flags.StringVar(&fHost, "h", "localhost", "")
	flags.StringVar(&fHost, "host", "localhost", "")

	var fPort string
	flags.StringVar(&fPort, "p", "4401", "")
	flags.StringVar(&fPort, "port", "4401", "")

	if err := flags.Parse(opts); err != nil {
		fmt.Println("[cli.ClusterList.run] Failed to parse flags!", err)
	}

	// create an RPC client that will connect to the matching node
	client, err := rpc.Dial("tcp", fmt.Sprintf("%s:%s", fHost, fPort))
	if err != nil {
		fmt.Println("[cli.ClusterList.run] Failed to dial!", err)
		os.Exit(1)
	}

	//
	defer client.Close()

	var members = &[]Status{}

	//
	if err := client.Call("Status.RPCCluster", "", members); err != nil {
		fmt.Println("[cli.ClusterList.run] Failed to call!", err)
		os.Exit(1)
	}

	//
	fmt.Println(`
Cluster Role |   Cluster IP    |       State        |  Postgres Role  |  Postgres Port  |      Last Updated
-----------------------------------------------------------------------------------------------------------------`)
	for _, member := range *members {
		fmt.Printf("%-12s | %-15s | %-18s | %-15s | %-15d | %-25s\n", member.CRole, member.Ip, member.State, member.DBRole, member.PGPort, member.UpdatedAt.Format("01.02.06 (15:04:05) MST"))
	}

	fmt.Println("")
}
<commit_msg>breaking apart state and status in cluster output<commit_after>package commands

import (
	"flag"
	"fmt"
	"net/rpc"
	"os"
	"regexp"
	"time"
)

// Status represents the Status of a node in the cluser
type Status struct {
	CRole     string    // the nodes 'role' in the cluster (primary, secondary, monitor)
	DataDir   string    // directory of the postgres database
	DBRole    string    // the 'role' of the running pgsql instance inside the node (master, slave)
	Ip        string    // advertise_ip
	PGPort    int       //
	State     string    // the current state of the node
	UpdatedAt time.Time // the last time the node state was updated
}

// ClusterListCommand satisfies the Command interface for listing nodes in yoke
type ClusterListCommand struct{}

// Help prints detailed help text for the cluster list command
func (c *ClusterListCommand) Help() {
	fmt.Printf(`
Description:
  Returns status information for all nodes in the cluster

Usage:
  cli list
  cli cluster:list

  ex. cli list
  `)
}

// Run displays select information about all of the nodes in a cluster
func (c *ClusterListCommand) Run(opts []string) {

	// flags
	flags := flag.NewFlagSet("flags", flag.ContinueOnError)
	flags.Usage = func() { c.Help() }

	var fHost string
	flags.StringVar(&fHost, "h", "localhost", "")
	flags.StringVar(&fHost, "host", "localhost", "")

	var fPort string
	flags.StringVar(&fPort, "p", "4401", "")
	flags.StringVar(&fPort, "port", "4401", "")

	if err := flags.Parse(opts); err != nil {
		fmt.Println("[cli.ClusterList.run] Failed to parse flags!", err)
	}

	// create an RPC client that will connect to the matching node
	client, err := rpc.Dial("tcp", fmt.Sprintf("%s:%s", fHost, fPort))
	if err != nil {
		fmt.Println("[cli.ClusterList.run] Failed to dial!", err)
		os.Exit(1)
	}

	//
	defer client.Close()

	var members = &[]Status{}

	//
	if err := client.Call("Status.RPCCluster", "", members); err != nil {
		fmt.Println("[cli.ClusterList.run] Failed to call!", err)
		os.Exit(1)
	}

	reFindErrorBody := regexp.MustCompile(`^\((.*)\)(.*)$`)

	//
	fmt.Println(`
Cluster Role |   Cluster IP    |     State     |    Status    |  Postgres Role  |  Postgres Port  |      Last Updated
---------------------------------------------------------------------------------------------------------------------------`)
	for _, member := range *members {

		state := "--"
		status := "--"

		subMatch := reFindErrorBody.FindStringSubmatch(member.State)
		if subMatch != nil {
			state = subMatch[1]
			status = subMatch[2]
		}

		fmt.Printf("%-12s | %-15s | %-13s | %-12s | %-15s | %-15d | %-25s\n", member.CRole, member.Ip, state, status, member.DBRole, member.PGPort, member.UpdatedAt.Format("01.02.06 (15:04:05) MST"))
	}

	fmt.Println("")
}
<|endoftext|>


<commit_before>package gokiq

import (
	"github.com/garyburd/redigo/redis"
	"testing"
	"time"
)

func resetRedis(pool *redis.Pool) {
	conn := pool.Get()
	defer conn.Close()
	conn.Do("FLUSHDB")
}

func TestGokiq(t *testing.T) {
	pool := redis.NewPool(func() (redis.Conn, error) {
		c, err := redis.Dial("tcp", ":6379")
		if err != nil {
			return nil, err
		}
		return c, err
	}, 3)

	conn := pool.Get()
	defer conn.Close()
	defer pool.Close()

	job := NewJob("HardWorder", "default", []string{})
	job.Enqueue(pool)

	count, _ := redis.Int(conn.Do("SISMEMBER", "queues", job.Queue))

	if count != 1 {
		t.Error("Expected queues list to have the correct queue but didn't found it.")
	}

	count, _ = redis.Int(conn.Do("LLEN", "queue:default"))

	if count != 1 {
		t.Errorf("Expected the queue to have exactly one job but found %d", count)
	}

	now := time.Now()
	job.EnqueueAt(now, pool)

	score, _ := redis.Int64(conn.Do("ZSCORE", "schedule", job.toJSON()))

	if score != now.Unix() {
		t.Errorf("Expected the timestamp to be %d but go %d", score, now.Unix())
	}

	resetRedis(pool)
}
<commit_msg>Test JSON output<commit_after>package gokiq

import (
	"fmt"
	"github.com/garyburd/redigo/redis"
	"testing"
	"time"
)

func resetRedis(pool *redis.Pool) {
	conn := pool.Get()
	defer conn.Close()
	conn.Do("FLUSHDB")
}

func TestGokiq(t *testing.T) {
	pool := redis.NewPool(func() (redis.Conn, error) {
		c, err := redis.Dial("tcp", ":6379")
		if err != nil {
			return nil, err
		}
		return c, err
	}, 3)

	conn := pool.Get()
	defer conn.Close()
	defer pool.Close()

	job := NewJob("HardWorder", "default", []string{})
	job.Enqueue(pool)

	expected := fmt.Sprintf(`{"jid":"%s","retry":false,"queue":"default","class":"HardWorder","args":[],"enqueued_at":%d}`,
		job.JID,
		job.EnqueuedAt)
	actual := job.toJSON()

	if expected != actual {
		t.Errorf("Excepted JSON to be %s, got %s", expected, job.toJSON())
	}

	count, _ := redis.Int(conn.Do("SISMEMBER", "queues", job.Queue))

	if count != 1 {
		t.Error("Expected queues list to have the correct queue but didn't found it.")
	}

	count, _ = redis.Int(conn.Do("LLEN", "queue:default"))

	if count != 1 {
		t.Errorf("Expected the queue to have exactly one job but found %d", count)
	}

	now := time.Now()
	job.EnqueueAt(now, pool)

	score, _ := redis.Int64(conn.Do("ZSCORE", "schedule", job.toJSON()))

	if score != now.Unix() {
		t.Errorf("Expected the timestamp to be %d but go %d", score, now.Unix())
	}

	resetRedis(pool)
}
<|endoftext|>


<commit_before>package rfc

/*
 * ZLint Copyright 2021 Regents of the University of Michigan
 *
 * 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.
 */

import (
	"unicode"

	"github.com/zmap/zcrypto/x509"
	"github.com/zmap/zlint/v3/lint"
	"github.com/zmap/zlint/v3/util"
)

type IANURIIA5String struct{}

/************************************************
When the issuerAltName extension contains a URI, the name MUST be
stored in the uniformResourceIdentifier (an IA5String).
************************************************/

func init() {
	lint.RegisterLint(&lint.Lint{
		Name:          "e_ext_ian_uri_not_ia5",
		Description:   "When subjectAltName contains a URI, the name MUST be an IA5 string",
		Citation:      "RFC5280: 4.2.1.7",
		Source:        lint.RFC5280,
		EffectiveDate: util.RFC5280Date,
		Lint:          &IANURIIA5String{},
	})
}

func (l *IANURIIA5String) Initialize() error {
	return nil
}

func (l *IANURIIA5String) CheckApplies(c *x509.Certificate) bool {
	return util.IsExtInCert(c, util.IssuerAlternateNameOID)
}

func (l *IANURIIA5String) Execute(c *x509.Certificate) *lint.LintResult {
	for _, uri := range c.IANURIs {
		for _, c := range uri {
			if c > unicode.MaxASCII {
				return &lint.LintResult{Status: lint.Error}
			}
		}
	}
	return &lint.LintResult{Status: lint.Pass}
}
<commit_msg>lints: fix description of e_ext_ian_uri_not_ia5 (#554)<commit_after>package rfc

/*
 * ZLint Copyright 2021 Regents of the University of Michigan
 *
 * 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.
 */

import (
	"unicode"

	"github.com/zmap/zcrypto/x509"
	"github.com/zmap/zlint/v3/lint"
	"github.com/zmap/zlint/v3/util"
)

type IANURIIA5String struct{}

/************************************************
When the issuerAltName extension contains a URI, the name MUST be
stored in the uniformResourceIdentifier (an IA5String).
************************************************/

func init() {
	lint.RegisterLint(&lint.Lint{
		Name:          "e_ext_ian_uri_not_ia5",
		Description:   "When issuer alternative name contains a URI, the name MUST be an IA5 string",
		Citation:      "RFC5280: 4.2.1.7",
		Source:        lint.RFC5280,
		EffectiveDate: util.RFC5280Date,
		Lint:          &IANURIIA5String{},
	})
}

func (l *IANURIIA5String) Initialize() error {
	return nil
}

func (l *IANURIIA5String) CheckApplies(c *x509.Certificate) bool {
	return util.IsExtInCert(c, util.IssuerAlternateNameOID)
}

func (l *IANURIIA5String) Execute(c *x509.Certificate) *lint.LintResult {
	for _, uri := range c.IANURIs {
		for _, c := range uri {
			if c > unicode.MaxASCII {
				return &lint.LintResult{Status: lint.Error}
			}
		}
	}
	return &lint.LintResult{Status: lint.Pass}
}
<|endoftext|>


<commit_before>package validation

import (
	"fmt"
	"reflect"
	"strings"
)

import (
	bvalidation "github.com/astaxie/beego/validation"
)

// Recursively validates a struct against `valid:"..." tags`
// The arg `parentType` is actually only used as an optional single argument
// Anything more than one is ignored
func Validate(obj interface{}, parentType ...string) (errs []error) {
	// Actual validation happens here
	validator := bvalidation.Validation{}
	valid, err := validator.Valid(obj)
	if err != nil {
		errs = append(errs, err)
		return
	}

	// The actual type name of this object
	Type := reflect.TypeOf(obj).String()

	// If validation above failed
	if !valid {
		// If this function was called with at least one argument
		// Prepend that argument to the type name
		// Helps us find context in recursive calls
		if len(parentType) > 0 {
			parts := strings.Split(Type, ".")
			Type = fmt.Sprintf("%s.%s", parentType[0], parts[len(parts)-1])
		}

		// Collect all the validation errors and return them
		for _, e := range validator.Errors {
			errFmt := "%s validation failed: `%s` %s (actual value: %#v)"
			errs = append(errs,
				fmt.Errorf(errFmt, Type, e.Field, e.Message, e.Value))
		}
	}

	// Figure out if this object is a pointer
	// If so, derefernce it
	v := reflect.ValueOf(obj)
	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}

	// Find sub-structs and validate them, too
	for i := range make([]struct{}, v.NumField()) {
		field := v.Field(i)
		if field.Kind() == reflect.Struct {
			errs = append(errs, Validate(field.Interface(), Type)...)
		} else if field.Kind() == reflect.Slice {
			if field.Len() > 0 {
				firstItem := field.Index(0)
				if firstItem.Kind() == reflect.Struct {
					for j := range make([]struct{}, field.Len()) {
						item := field.Index(j).Interface()
						errs = append(errs, Validate(item, Type)...)
					}
				}
			}
		}
	}
	return
}
<commit_msg>Better type display.<commit_after>package validation

import (
	"fmt"
	"reflect"
	"strings"
)

import (
	bvalidation "github.com/astaxie/beego/validation"
)

// Recursively validates a struct against `valid:"..." tags`
// The arg `parentType` is actually only used as an optional single argument
// Anything more than one is ignored
func Validate(obj interface{}, parentType ...string) (errs []error) {
	// Actual validation happens here
	validator := bvalidation.Validation{}
	valid, err := validator.Valid(obj)
	if err != nil {
		errs = append(errs, err)
		return
	}

	// Figure out if this object is a pointer
	// If so, derefernce it
	v := reflect.ValueOf(obj)
	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}

	// The actual type name of this object
	Type := v.Type().String()

	// If validation above failed
	if !valid {
		// If this function was called with at least one argument
		// Prepend that argument to the type name
		// Helps us find context in recursive calls
		if len(parentType) > 0 {
			parts := strings.Split(Type, ".")
			Type = fmt.Sprintf("%s.%s", parentType[0], parts[len(parts)-1])
		}

		// Collect all the validation errors and return them
		for _, e := range validator.Errors {
			errFmt := "%s validation failed: `%s` %s (actual value: %#v)"
			errs = append(errs,
				fmt.Errorf(errFmt, Type, e.Field, e.Message, e.Value))
		}
	}

	// Find sub-structs and validate them, too
	for i := range make([]struct{}, v.NumField()) {
		field := v.Field(i)
		if field.Kind() == reflect.Struct {
			errs = append(errs, Validate(field.Interface(), Type)...)
		} else if field.Kind() == reflect.Slice {
			if field.Len() > 0 {
				firstItem := field.Index(0)
				if firstItem.Kind() == reflect.Struct {
					for j := range make([]struct{}, field.Len()) {
						item := field.Index(j).Interface()
						errs = append(errs, Validate(item, Type)...)
					}
				}
			}
		}
	}
	return
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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.

// The integration driver provides a suite of tests to run against a registered runner.
package main

import (
	"context"
	"flag"
	"regexp"
	"sync"
	"sync/atomic"

	"github.com/apache/beam/sdks/go/pkg/beam"
	"github.com/apache/beam/sdks/go/pkg/beam/io/filesystem/memfs"
	"github.com/apache/beam/sdks/go/pkg/beam/log"
	"github.com/apache/beam/sdks/go/pkg/beam/x/beamx"
	"github.com/apache/beam/sdks/go/test/integration/wordcount"
)

var (
	cpus   = flag.Int("cpus", 10, "Number of tests to run in parallel")
	filter = flag.String("filter", ".*", "Test filer to run a subset of tests")
)

const old_pond = "memfs://old_pond"

func init() {
	memfs.Write(old_pond, []byte("old pond \na frog leaps in\nwater's sound\n"))
}

type namedPipeline struct {
	name string
	p    *beam.Pipeline
}

func main() {
	flag.Parse()
	beam.Init()

	if *cpus < 1 {
		*cpus = 1
	}

	pipelines := []namedPipeline{
		{"wordcount:memfs", wordcount.WordCount(old_pond, "+Qj8iAnV5BI2A4sbzUbb6Q==", 8)},
		{"wordcount:kinglear", wordcount.WordCount("gs://apache-beam-samples/shakespeare/kinglear.txt", "7ZCh5ih9m8IW1w+iS8sRKg==", 4749)},
	}

	re := regexp.MustCompile(*filter)

	ch := make(chan namedPipeline, len(pipelines))
	for _, np := range pipelines {
		if re.MatchString(np.name) {
			ch <- np
		}
	}
	close(ch)

	ctx := context.Background()

	var failures int32
	var wg sync.WaitGroup
	for i := 0; i < *cpus; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for np := range ch {
				log.Infof(ctx, "Running %v test ..", np.name)

				if err := beamx.Run(ctx, np.p); err != nil {
					atomic.AddInt32(&failures, 1)
					log.Errorf(ctx, "Test %v failed: %v", np.name, err)
				} else {
					log.Infof(ctx, "Test %v completed", np.name)
				}
			}
		}()
	}
	wg.Wait()

	if failures > 0 {
		log.Exitf(ctx, "Result: %v tests failed", failures)
	}
	log.Infof(ctx, "Result: all tests passed!")
}
<commit_msg>Tweak Go integration test driver<commit_after>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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.

// The integration driver provides a suite of tests to run against a registered runner.
package main

import (
	"context"
	"flag"
	"regexp"
	"sync"
	"sync/atomic"

	"github.com/apache/beam/sdks/go/pkg/beam"
	"github.com/apache/beam/sdks/go/pkg/beam/io/filesystem/memfs"
	"github.com/apache/beam/sdks/go/pkg/beam/log"
	"github.com/apache/beam/sdks/go/pkg/beam/x/beamx"
	"github.com/apache/beam/sdks/go/test/integration/wordcount"
)

var (
	parallel = flag.Int("parallel", 10, "Number of tests to run in parallel")
	filter   = flag.String("filter", ".*", "Test filer to run a subset of tests")
)

const old_pond = "memfs://old_pond"

func init() {
	memfs.Write(old_pond, []byte("old pond \na frog leaps in\nwater's sound\n"))
}

type namedPipeline struct {
	name string
	p    *beam.Pipeline
}

func main() {
	flag.Parse()
	beam.Init()

	if *parallel < 1 {
		*parallel = 1
	}

	pipelines := []namedPipeline{
		{"wordcount:memfs", wordcount.WordCount(old_pond, "+Qj8iAnV5BI2A4sbzUbb6Q==", 8)},
		{"wordcount:kinglear", wordcount.WordCount("gs://apache-beam-samples/shakespeare/kinglear.txt", "7ZCh5ih9m8IW1w+iS8sRKg==", 4749)},
	}

	re := regexp.MustCompile(*filter)

	ch := make(chan namedPipeline, len(pipelines))
	for _, np := range pipelines {
		if re.MatchString(np.name) {
			ch <- np
		}
	}
	close(ch)

	ctx := context.Background()

	var failures int32
	var wg sync.WaitGroup
	for i := 0; i < *parallel; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for np := range ch {
				log.Infof(ctx, "Running %v test ..", np.name)

				if err := beamx.Run(ctx, np.p); err != nil {
					atomic.AddInt32(&failures, 1)
					log.Errorf(ctx, "Test %v failed: %v", np.name, err)
				} else {
					log.Infof(ctx, "Test %v completed", np.name)
				}
			}
		}()
	}
	wg.Wait()

	if failures > 0 {
		log.Exitf(ctx, "Result: %v tests failed", failures)
	}
	log.Infof(ctx, "Result: all tests passed!")
}
<|endoftext|>


<commit_before>package password

import (
    "regexp"
    "math/rand"
    "time"
    "encoding/base64"
    "crypto/sha256"
    "strings"
)

func ValidatePassword(s string) bool {
    pattern := "^(?=.*\\d)(?=.*[a-z])(?=.*[A-Z]).{8,20}$"
    matched, _ := regexp.MatchString(pattern, s)
    return matched
}

const letterBytes = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
const (
    letterIdxBits = 6                    // 6 bits to represent a letter index
    letterIdxMask = 1<<letterIdxBits - 1 // All 1-bits, as many as letterIdxBits
    letterIdxMax  = 63 / letterIdxBits   // # of letter indices fitting in 63 bits
)
var src = rand.NewSource(time.Now().UnixNano())
func MakeSalt(n int) string {
    salt := make([]byte, n)
    // A src.Int63() generates 63 random bits, enough for letterIdxMax characters!
    for i, cache, remain := n-1, src.Int63(), letterIdxMax; i >= 0; {
        if remain == 0 {
            cache, remain = src.Int63(), letterIdxMax
        }
        if idx := int(cache & letterIdxMask); idx < len(letterBytes) {
            salt[i] = letterBytes[idx]
            i--
        }
        cache >>= letterIdxBits
        remain--
    }

    return string(salt)
}

func AutoAddSalt(s string) (string, string, bool) {
    if !ValidatePassword(s) { return "", "", false }
    minLen := 8
    passB := []byte(s)
    salted := make([]byte, len(passB)+minLen)
    salt := MakeSalt(minLen)
    for i := 0; i<minLen; i++ {
        salted[2*i] = passB[i]
        salted[2*i+1] = salt[i]
    }
    for i := minLen; i<len(passB)-minLen; i++ {
        salted[2*minLen+i] = passB[minLen+i]
    }
    return string(salted), salt, true
}

func HashAutoSalt(s string) (string, string, bool) {
    salted, salt, valid := AutoAddSalt(s)
    if !valid { return "", "", false }
    h := sha256.New()
    hash := string(base64.StdEncoding.EncodeToString(h.Sum([]byte(salted))))
    return hash, salt, true
}

func HashWithSalt(s string, salt string) (string, string, bool) {
    if !ValidatePassword(s) { return "", "", false }
    h := sha256.New()
    hash := string(base64.StdEncoding.EncodeToString(h.Sum([]byte(salt))))
    return hash, salt, true
}

func Compare(pass string, salt string, hash string) (bool, bool) {
    hashed, _, valid := HashWithSalt(pass, salt)
    return strings.Compare(hashed, hash) == 0, valid
}
<commit_msg>just left it here first, the regex for password is not working, as not supported<commit_after>package password

import (
    "regexp"
    "math/rand"
    "time"
    "encoding/base64"
    "crypto/sha256"
    "strings"
    //"fmt"
)

func ValidatePassword(s string) bool {
    pattern := "^(?=.*\\d)(?=.*[a-z])(?=.*[A-Z]).{8,20}$"
    compiledp, _ := regexp.Compile(pattern)
    matched := compiledp.MatchString(s)
    return matched
}

const letterBytes = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
const (
    letterIdxBits = 6                    // 6 bits to represent a letter index
    letterIdxMask = 1<<letterIdxBits - 1 // All 1-bits, as many as letterIdxBits
    letterIdxMax  = 63 / letterIdxBits   // # of letter indices fitting in 63 bits
)
var src = rand.NewSource(time.Now().UnixNano())
func MakeSalt(n int) string {
    salt := make([]byte, n)
    // A src.Int63() generates 63 random bits, enough for letterIdxMax characters!
    for i, cache, remain := n-1, src.Int63(), letterIdxMax; i >= 0; {
        if remain == 0 {
            cache, remain = src.Int63(), letterIdxMax
        }
        if idx := int(cache & letterIdxMask); idx < len(letterBytes) {
            salt[i] = letterBytes[idx]
            i--
        }
        cache >>= letterIdxBits
        remain--
    }

    return string(salt)
}

func AutoAddSalt(s string) (string, string, bool) {
    if !ValidatePassword(s) { return "", "", false }
    minLen := 8
    passB := []byte(s)
    salted := make([]byte, len(passB)+minLen)
    salt := MakeSalt(minLen)
    for i := 0; i<minLen; i++ {
        salted[2*i] = passB[i]
        salted[2*i+1] = salt[i]
    }
    for i := minLen; i<len(passB)-minLen; i++ {
        salted[2*minLen+i] = passB[minLen+i]
    }
    return string(salted), salt, true
}

func HashAutoSalt(s string) (string, string, bool) {
    salted, salt, valid := AutoAddSalt(s)
    if !valid { return "", "", false }
    h := sha256.New()
    hash := string(base64.StdEncoding.EncodeToString(h.Sum([]byte(salted))))
    return hash, salt, true
}

func HashWithSalt(s string, salt string) (string, string, bool) {
    if !ValidatePassword(s) { return "", "", false }
    h := sha256.New()
    hash := string(base64.StdEncoding.EncodeToString(h.Sum([]byte(salt))))
    return hash, salt, true
}

func Compare(pass string, salt string, hash string) (bool, bool) {
    hashed, _, valid := HashWithSalt(pass, salt)
    return strings.Compare(hashed, hash) == 0, valid
}
<|endoftext|>


<commit_before>/* Copyright (c) 2014-2015, Daniel Martí <mvdan@mvdan.cc> */
/* See LICENSE for licensing information */

package main

import (
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/mvdan/bytesize"
)

const (
	// Length of the random hexadecimal ids assigned to pastes. At least 4.
	idSize = 8
	// Number of times to try getting an unused random paste id
	randTries = 10
	// Name of the HTTP form field when uploading a paste
	fieldName = "paste"
	// Content-Type when serving pastes
	contentType = "text/plain; charset=utf-8"
	// Report usage stats how often
	statsReport = 1 * time.Minute
	// How long to wait before retrying to delete a file
	deleteRetry = 2 * time.Minute

	// HTTP response strings
	invalidID     = "invalid paste id"
	unknownAction = "unsupported action"
)

var (
	siteURL = flag.String("u", "http://localhost:8080", "URL of the site")
	listen = flag.String("l", ":8080", "Host and port to listen to")
	lifeTime = flag.Duration("t", 24*time.Hour, "Lifetime of the pastes")
	maxNumber = flag.Int("m", 0, "Maximum number of pastes to store at once")

	maxSize    = 1 * bytesize.MB
	maxStorage = 1 * bytesize.GB
)

func init() {
	flag.Var(&maxSize, "s", "Maximum size of pastes")
	flag.Var(&maxStorage, "M", "Maximum storage size to use at once")
}

// Binary representation of an identifier for a paste
type ID [idSize / 2]byte

// Parse a hexadecimal string into an ID. Return the resulting ID and an
// error, if any.
func IDFromString(hexID string) (id ID, err error) {
	if len(hexID) != idSize {
		return id, fmt.Errorf("invalid id at %s", hexID)
	}
	b, err := hex.DecodeString(hexID)
	if err != nil || len(b) != idSize/2 {
		return id, fmt.Errorf("invalid id at %s", hexID)
	}
	copy(id[:], b)
	return id, nil
}

func (id ID) String() string {
	return hex.EncodeToString(id[:])
}

func describeLimits() string {
	var limits []string
	if maxSize > 0 {
		limits = append(limits, fmt.Sprintf("Maximum size per paste is %s.", maxSize))
	}
	if *lifeTime > 0 {
		limits = append(limits, fmt.Sprintf("Pastes will be deleted after %s.", *lifeTime))
	}
	if len(limits) > 0 {
		return strings.Join(limits, " ") + "\n\n"
	}
	return ""
}

func getContentFromForm(r *http.Request) ([]byte, error) {
	if value := r.FormValue(fieldName); len(value) > 0 {
		return []byte(value), nil
	}
	if f, _, err := r.FormFile(fieldName); err == nil {
		defer f.Close()
		content, err := ioutil.ReadAll(f)
		if err == nil && len(content) > 0 {
			return content, nil
		}
	}
	return nil, errors.New("no paste provided")
}

func setHeaders(header http.Header, id ID, paste Paste) {
	modTime := paste.ModTime()
	header.Set("Etag", fmt.Sprintf("%d-%s", modTime.Unix(), id))
	if *lifeTime > 0 {
		deathTime := modTime.Add(*lifeTime)
		lifeLeft := deathTime.Sub(time.Now())
		header.Set("Expires", deathTime.UTC().Format(http.TimeFormat))
		header.Set("Cache-Control", fmt.Sprintf(
			"max-age=%.f, must-revalidate", lifeLeft.Seconds()))
	}
	header.Set("Content-Type", contentType)
}

func handleGet(store Store, w http.ResponseWriter, r *http.Request) {
	if _, e := templates[r.URL.Path]; e {
		err := tmpl.ExecuteTemplate(w, r.URL.Path,
			struct{ SiteURL, LimitDesc, FieldName string }{
				*siteURL, describeLimits(), fieldName})
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
		}
		return
	}
	id, err := IDFromString(r.URL.Path[1:])
	if err != nil {
		http.Error(w, invalidID, http.StatusBadRequest)
		return
	}
	paste, err := store.Get(id)
	if err == ErrPasteNotFound {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	} else if err != nil {
		log.Printf("Unknown error on GET: %s", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer paste.Close()
	setHeaders(w.Header(), id, paste)
	http.ServeContent(w, r, "", paste.ModTime(), paste)
}

func handlePost(store Store, stats *Stats, w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, int64(maxSize))
	content, err := getContentFromForm(r)
	size := int64(len(content))
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if err := stats.makeSpaceFor(size); err != nil {
		http.Error(w, err.Error(), http.StatusServiceUnavailable)
	}
	id, err := store.Put(content)
	if err != nil {
		log.Printf("Unknown error on POST: %s", err)
		stats.freeSpace(size)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	setupPasteDeletion(store, stats, id, size, *lifeTime)
	fmt.Fprintf(w, "%s/%s\n", *siteURL, id)
}

func newHandler(store Store, stats *Stats) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case "GET":
			handleGet(store, w, r)
		case "POST":
			handlePost(store, stats, w, r)
		default:
			http.Error(w, unknownAction, http.StatusBadRequest)
		}
	})
}

func setupStore(stats *Stats, storageType string, args []string) (Store, error) {
	params, e := map[string]map[string]string{
		"fs": {
			"dir": "pastes",
		},
		"fs-mmap": {
			"dir": "pastes",
		},
		"mem": {},
	}[storageType]
	if !e {
		return nil, fmt.Errorf("unknown storage type '%s'", storageType)
	}
	if len(args) > len(params) {
		return nil, fmt.Errorf("too many arguments given for %s", storageType)
	}
	for k := range params {
		if len(args) == 0 {
			break
		}
		params[k] = args[0]
		args = args[1:]
	}
	switch storageType {
	case "fs":
		log.Printf("Starting up file store in the directory '%s'", params["dir"])
		return NewFileStore(stats, params["dir"])
	case "fs-mmap":
		log.Printf("Starting up mmapped file store in the directory '%s'", params["dir"])
		return NewMmapStore(stats, params["dir"])
	case "mem":
		log.Printf("Starting up in-memory store")
		return NewMemStore()
	}
	return nil, nil
}

func main() {
	flag.Parse()
	if maxStorage > 1*bytesize.EB {
		log.Fatalf("Specified a maximum storage size that would overflow int64!")
	}
	stats := Stats{
		maxNumber:  *maxNumber,
		maxStorage: int64(maxStorage),
	}
	log.Printf("siteURL    = %s", *siteURL)
	log.Printf("listen     = %s", *listen)
	log.Printf("lifeTime   = %s", *lifeTime)
	log.Printf("maxSize    = %s", maxSize)
	log.Printf("maxNumber  = %d", *maxNumber)
	log.Printf("maxStorage = %s", maxStorage)

	args := flag.Args()
	if len(args) == 0 {
		args = []string{"fs"}
	}
	store, err := setupStore(&stats, args[0], args[1:])
	if err != nil {
		log.Fatalf("Could not setup paste store: %s", err)
	}
	loadTemplates()
	ticker := time.NewTicker(statsReport)
	go func() {
		for _ = range ticker.C {
			log.Println(stats.Report())
		}
	}()
	http.HandleFunc("/", newHandler(store, &stats))
	log.Println("Up and running!")
	log.Println(stats.Report())
	log.Fatal(http.ListenAndServe(*listen, nil))
}
<commit_msg>Also warn on overflowing maxSize<commit_after>/* Copyright (c) 2014-2015, Daniel Martí <mvdan@mvdan.cc> */
/* See LICENSE for licensing information */

package main

import (
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/mvdan/bytesize"
)

const (
	// Length of the random hexadecimal ids assigned to pastes. At least 4.
	idSize = 8
	// Number of times to try getting an unused random paste id
	randTries = 10
	// Name of the HTTP form field when uploading a paste
	fieldName = "paste"
	// Content-Type when serving pastes
	contentType = "text/plain; charset=utf-8"
	// Report usage stats how often
	statsReport = 1 * time.Minute
	// How long to wait before retrying to delete a file
	deleteRetry = 2 * time.Minute

	// HTTP response strings
	invalidID     = "invalid paste id"
	unknownAction = "unsupported action"
)

var (
	siteURL = flag.String("u", "http://localhost:8080", "URL of the site")
	listen = flag.String("l", ":8080", "Host and port to listen to")
	lifeTime = flag.Duration("t", 24*time.Hour, "Lifetime of the pastes")
	maxNumber = flag.Int("m", 0, "Maximum number of pastes to store at once")

	maxSize    = 1 * bytesize.MB
	maxStorage = 1 * bytesize.GB
)

func init() {
	flag.Var(&maxSize, "s", "Maximum size of pastes")
	flag.Var(&maxStorage, "M", "Maximum storage size to use at once")
}

// Binary representation of an identifier for a paste
type ID [idSize / 2]byte

// Parse a hexadecimal string into an ID. Return the resulting ID and an
// error, if any.
func IDFromString(hexID string) (id ID, err error) {
	if len(hexID) != idSize {
		return id, fmt.Errorf("invalid id at %s", hexID)
	}
	b, err := hex.DecodeString(hexID)
	if err != nil || len(b) != idSize/2 {
		return id, fmt.Errorf("invalid id at %s", hexID)
	}
	copy(id[:], b)
	return id, nil
}

func (id ID) String() string {
	return hex.EncodeToString(id[:])
}

func describeLimits() string {
	var limits []string
	if maxSize > 0 {
		limits = append(limits, fmt.Sprintf("Maximum size per paste is %s.", maxSize))
	}
	if *lifeTime > 0 {
		limits = append(limits, fmt.Sprintf("Pastes will be deleted after %s.", *lifeTime))
	}
	if len(limits) > 0 {
		return strings.Join(limits, " ") + "\n\n"
	}
	return ""
}

func getContentFromForm(r *http.Request) ([]byte, error) {
	if value := r.FormValue(fieldName); len(value) > 0 {
		return []byte(value), nil
	}
	if f, _, err := r.FormFile(fieldName); err == nil {
		defer f.Close()
		content, err := ioutil.ReadAll(f)
		if err == nil && len(content) > 0 {
			return content, nil
		}
	}
	return nil, errors.New("no paste provided")
}

func setHeaders(header http.Header, id ID, paste Paste) {
	modTime := paste.ModTime()
	header.Set("Etag", fmt.Sprintf("%d-%s", modTime.Unix(), id))
	if *lifeTime > 0 {
		deathTime := modTime.Add(*lifeTime)
		lifeLeft := deathTime.Sub(time.Now())
		header.Set("Expires", deathTime.UTC().Format(http.TimeFormat))
		header.Set("Cache-Control", fmt.Sprintf(
			"max-age=%.f, must-revalidate", lifeLeft.Seconds()))
	}
	header.Set("Content-Type", contentType)
}

func handleGet(store Store, w http.ResponseWriter, r *http.Request) {
	if _, e := templates[r.URL.Path]; e {
		err := tmpl.ExecuteTemplate(w, r.URL.Path,
			struct{ SiteURL, LimitDesc, FieldName string }{
				*siteURL, describeLimits(), fieldName})
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
		}
		return
	}
	id, err := IDFromString(r.URL.Path[1:])
	if err != nil {
		http.Error(w, invalidID, http.StatusBadRequest)
		return
	}
	paste, err := store.Get(id)
	if err == ErrPasteNotFound {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	} else if err != nil {
		log.Printf("Unknown error on GET: %s", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer paste.Close()
	setHeaders(w.Header(), id, paste)
	http.ServeContent(w, r, "", paste.ModTime(), paste)
}

func handlePost(store Store, stats *Stats, w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, int64(maxSize))
	content, err := getContentFromForm(r)
	size := int64(len(content))
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if err := stats.makeSpaceFor(size); err != nil {
		http.Error(w, err.Error(), http.StatusServiceUnavailable)
	}
	id, err := store.Put(content)
	if err != nil {
		log.Printf("Unknown error on POST: %s", err)
		stats.freeSpace(size)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	setupPasteDeletion(store, stats, id, size, *lifeTime)
	fmt.Fprintf(w, "%s/%s\n", *siteURL, id)
}

func newHandler(store Store, stats *Stats) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case "GET":
			handleGet(store, w, r)
		case "POST":
			handlePost(store, stats, w, r)
		default:
			http.Error(w, unknownAction, http.StatusBadRequest)
		}
	})
}

func setupStore(stats *Stats, storageType string, args []string) (Store, error) {
	params, e := map[string]map[string]string{
		"fs": {
			"dir": "pastes",
		},
		"fs-mmap": {
			"dir": "pastes",
		},
		"mem": {},
	}[storageType]
	if !e {
		return nil, fmt.Errorf("unknown storage type '%s'", storageType)
	}
	if len(args) > len(params) {
		return nil, fmt.Errorf("too many arguments given for %s", storageType)
	}
	for k := range params {
		if len(args) == 0 {
			break
		}
		params[k] = args[0]
		args = args[1:]
	}
	switch storageType {
	case "fs":
		log.Printf("Starting up file store in the directory '%s'", params["dir"])
		return NewFileStore(stats, params["dir"])
	case "fs-mmap":
		log.Printf("Starting up mmapped file store in the directory '%s'", params["dir"])
		return NewMmapStore(stats, params["dir"])
	case "mem":
		log.Printf("Starting up in-memory store")
		return NewMemStore()
	}
	return nil, nil
}

func main() {
	flag.Parse()
	if maxStorage > 1*bytesize.EB {
		log.Fatalf("Specified a maximum storage size that would overflow int64!")
	}
	if maxSize > 1*bytesize.EB {
		log.Fatalf("Specified a maximum paste size that would overflow int64!")
	}
	stats := Stats{
		maxNumber:  *maxNumber,
		maxStorage: int64(maxStorage),
	}
	log.Printf("siteURL    = %s", *siteURL)
	log.Printf("listen     = %s", *listen)
	log.Printf("lifeTime   = %s", *lifeTime)
	log.Printf("maxSize    = %s", maxSize)
	log.Printf("maxNumber  = %d", *maxNumber)
	log.Printf("maxStorage = %s", maxStorage)

	args := flag.Args()
	if len(args) == 0 {
		args = []string{"fs"}
	}
	store, err := setupStore(&stats, args[0], args[1:])
	if err != nil {
		log.Fatalf("Could not setup paste store: %s", err)
	}
	loadTemplates()
	ticker := time.NewTicker(statsReport)
	go func() {
		for _ = range ticker.C {
			log.Println(stats.Report())
		}
	}()
	http.HandleFunc("/", newHandler(store, &stats))
	log.Println("Up and running!")
	log.Println(stats.Report())
	log.Fatal(http.ListenAndServe(*listen, nil))
}
<|endoftext|>


<commit_before>/* Copyright (c) 2014, Daniel Martí <mvdan@mvdan.cc> */
/* See LICENSE for licensing information */

package main

import (
	"errors"
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	indexTmpl = "index.html"
	formTmpl  = "form.html"
	chars     = "abcdefghijklmnopqrstuvwxyz0123456789"
	randTries = 10

	// GET error messages
	invalidId     = "Invalid paste id."
	pasteNotFound = "Paste doesn't exist."
	unknownError  = "Something went terribly wrong."
	// POST error messages
	missingForm = "Paste could not be found inside the posted form."
)

var (
	siteUrl, listen, dataDir string
	lifeTime                 time.Duration
	maxSizeStr               string
	idSize                   int
	maxSize                  ByteSize

	validId       *regexp.Regexp
	regexByteSize = regexp.MustCompile(`^([\d\.]+)\s*([KM]?B|[BKM])$`)
	indexTemplate *template.Template
	formTemplate  *template.Template
	data          = struct {
		sync.RWMutex
		m map[Id]PasteInfo
	}{m: make(map[Id]PasteInfo)}
	customRand *rand.Rand
)

func init() {
	flag.StringVar(&siteUrl, "u", "http://localhost:8080", "URL of the site")
	flag.StringVar(&listen, "l", "localhost:8080", "Host and port to listen to")
	flag.StringVar(&dataDir, "d", "data", "Directory to store all the pastes in")
	flag.DurationVar(&lifeTime, "t", 12*time.Hour, "Lifetime of the pastes (units: s,m,h)")
	flag.StringVar(&maxSizeStr, "s", "1M", "Maximum size of POSTs in bytes (units: B,K,M)")
	flag.IntVar(&idSize, "i", 8, "Size of the paste ids (between 6 and 256)")
	validId = regexp.MustCompile("^[a-zA-Z0-9]{" + strconv.Itoa(idSize) + "}$")
	customRand = rand.New(rand.NewSource(time.Now().UnixNano()))
}

type PasteInfo struct {
	ModTime   time.Time
	DeathTime time.Time
	Size      ByteSize
}

type Id string

func IdFromPath(idPath string) (Id, error) {
	parts := strings.Split(idPath, string(filepath.Separator))
	if len(parts) != 3 {
		return "", errors.New("Found invalid number of directories at " + idPath)
	}
	rawId := parts[0] + parts[1] + parts[2]
	if !validId.MatchString(rawId) {
		return "", errors.New("Found invalid id " + rawId)
	}
	return Id(rawId), nil
}

func RandomId() (Id, error) {
	s := make([]byte, idSize)
	var id Id
	data.RLock()
	for try := 0; try < randTries; try++ {
		var offset int = 0
	RandLoop:
		for {
			r := customRand.Int63()
			for i := 0; i < 8; i++ {
				randbyte := int(r&0xff) % len(chars)
				s[offset] = chars[randbyte]
				offset++
				if offset == idSize {
					break RandLoop
				}
				r >>= 8
			}
		}
		id = Id(s)
		if _, e := data.m[id]; !e {
			data.RUnlock()
			return id, nil
		}
	}
	return id, fmt.Errorf("Gave up trying to find an unused random id after %d tries", randTries)
}

func (id Id) String() string {
	return string(id)
}

func (id Id) Path() string {
	return path.Join(string(id[0:2]), string(id[2:4]), string(id[4:]))
}

func (id Id) EndLife() {
	data.Lock()
	err := os.Remove(id.Path())
	if err == nil {
		delete(data.m, id)
		data.Unlock()
		log.Printf("Removed paste: %s", id)
	} else {
		data.Unlock()
		log.Printf("Could not end the life of %s: %s", id, err)
		id.EndLifeAfter(2 * time.Minute)
	}
}

func (id Id) EndLifeAfter(duration time.Duration) {
	timer := time.NewTimer(duration)
	go func() {
		<-timer.C
		id.EndLife()
	}()
}

type ByteSize int64

const (
	B ByteSize = 1 << (10 * iota)
	KB
	MB
)

func parseByteSize(str string) (ByteSize, error) {
	if !regexByteSize.MatchString(str) {
		return 0, errors.New("Could not parse size in bytes")
	}
	parts := regexByteSize.FindStringSubmatch(str)
	size, _ := strconv.ParseFloat(string(parts[1]), 64)

	switch string(parts[2]) {
	case "KB", "K":
		size *= float64(KB)
	case "MB", "M":
		size *= float64(MB)
	}
	return ByteSize(size), nil
}

func (b ByteSize) String() string {
	switch {
	case b >= MB:
		return fmt.Sprintf("%.2f MB", float64(b)/float64(MB))
	case b >= KB:
		return fmt.Sprintf("%.2f KB", float64(b)/float64(KB))
	}
	return fmt.Sprintf("%d B", b)
}

func handler(w http.ResponseWriter, r *http.Request) {
	var err error
	switch r.Method {
	case "GET":
		switch r.URL.Path {
		case "/":
			indexTemplate.Execute(w, siteUrl)
			return
		case "/form":
			formTemplate.Execute(w, siteUrl)
			return
		}
		rawId := r.URL.Path[1:]
		if !validId.MatchString(rawId) {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", invalidId)
			return
		}
		id := Id(strings.ToLower(rawId))
		data.RLock()
		pasteInfo, e := data.m[id]
		data.RUnlock()
		if !e {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "%s\n", pasteNotFound)
			return
		}
		etag := fmt.Sprintf("%d-%s", pasteInfo.ModTime.Unix(), id)
		if inm := r.Header.Get("If-None-Match"); inm != "" {
			if etag == inm || inm == "*" {
				w.WriteHeader(http.StatusNotModified)
				return
			}
		}
		pastePath := id.Path()
		pasteFile, err := os.Open(pastePath)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		defer pasteFile.Close()
		w.Header().Set("Etag", etag)
		w.Header().Set("Content-Type", "text-plain; charset=utf-8")
		http.ServeContent(w, r, "", pasteInfo.ModTime, pasteFile)

	case "POST":
		r.Body = http.MaxBytesReader(w, r.Body, int64(maxSize))
		var id Id
		var content string
		id, err = RandomId()
		if err != nil {
			log.Println(err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		if err = r.ParseMultipartForm(int64(maxSize)); err != nil {
			log.Printf("Could not parse POST multipart form: %s", err)
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", err)
			return
		}
		if vs, found := r.Form["paste"]; found {
			content = vs[0]
		} else {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", missingForm)
			return
		}
		pastePath := id.Path()
		dir, _ := path.Split(pastePath)
		if err = os.MkdirAll(dir, 0700); err != nil {
			log.Printf("Could not create directories leading to %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		deathTime := time.Now().Add(lifeTime)
		id.EndLifeAfter(lifeTime)
		pasteFile, err := os.OpenFile(pastePath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
		if err != nil {
			log.Printf("Could not create new paste pasteFile %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		defer pasteFile.Close()
		b, err := io.WriteString(pasteFile, content)
		if err != nil {
			log.Printf("Could not write data into %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		writtenSize := ByteSize(b)
		data.Lock()
		data.m[id] = PasteInfo{
			ModTime:   time.Now(),
			DeathTime: deathTime,
			Size:      writtenSize,
		}
		data.Unlock()
		log.Printf("Created a new paste: %s (%s)", id, writtenSize)
		fmt.Fprintf(w, "%s/%s\n", siteUrl, id)
	}
}

func walkFunc(filePath string, fileInfo os.FileInfo, err error) error {
	if err != nil {
		return err
	}
	if fileInfo.IsDir() {
		return nil
	}
	id, err := IdFromPath(filePath)
	if err != nil {
		return err
	}
	modTime := fileInfo.ModTime()
	deathTime := modTime.Add(lifeTime)
	now := time.Now()
	if deathTime.Before(now) {
		go id.EndLife()
		return nil
	}
	var lifeLeft time.Duration
	if deathTime.After(now.Add(lifeTime)) {
		lifeLeft = lifeTime
	} else {
		lifeLeft = deathTime.Sub(now)
	}
	size := ByteSize(fileInfo.Size())
	data.Lock()
	data.m[id] = PasteInfo{
		ModTime:   modTime,
		DeathTime: deathTime,
		Size:      size,
	}
	data.Unlock()
	log.Printf("Recovered paste %s (%s) from %s has %s left", id, size, modTime, lifeLeft)
	id.EndLifeAfter(lifeLeft)
	return nil
}

func main() {
	var err error
	flag.Parse()
	if idSize < 6 || idSize > 256 {
		log.Fatalf("Provided id size %d is not between 6 and 256", idSize)
	}
	if maxSize, err = parseByteSize(maxSizeStr); err != nil {
		log.Fatalf("Invalid max size '%s': %s", maxSizeStr, err)
	}
	if indexTemplate, err = template.ParseFiles(indexTmpl); err != nil {
		log.Fatalf("Could not load template %s: %s", indexTmpl, err)
	}
	if formTemplate, err = template.ParseFiles(formTmpl); err != nil {
		log.Fatalf("Could not load template %s: %s", formTmpl, err)
	}
	if err = os.MkdirAll(dataDir, 0700); err != nil {
		log.Fatalf("Could not create data directory %s: %s", dataDir, err)
	}
	if err = os.Chdir(dataDir); err != nil {
		log.Fatalf("Could not enter data directory %s: %s", dataDir, err)
	}
	log.Printf("idSize   = %d", idSize)
	log.Printf("maxSize  = %s", maxSize)
	log.Printf("siteUrl  = %s", siteUrl)
	log.Printf("listen   = %s", listen)
	log.Printf("dataDir  = %s", dataDir)
	log.Printf("lifeTime = %s", lifeTime)
	if err = filepath.Walk(".", walkFunc); err != nil {
		log.Fatalf("Could not recover data directory %s: %s", dataDir, err)
	}
	http.HandleFunc("/", handler)
	log.Fatal(http.ListenAndServe(listen, nil))
}
<commit_msg>Don't accept zero-length pastes<commit_after>/* Copyright (c) 2014, Daniel Martí <mvdan@mvdan.cc> */
/* See LICENSE for licensing information */

package main

import (
	"errors"
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	indexTmpl = "index.html"
	formTmpl  = "form.html"
	chars     = "abcdefghijklmnopqrstuvwxyz0123456789"
	randTries = 10

	// GET error messages
	invalidId     = "Invalid paste id."
	pasteNotFound = "Paste doesn't exist."
	unknownError  = "Something went terribly wrong."
	// POST error messages
	missingForm = "Paste could not be found inside the posted form."
)

var (
	siteUrl, listen, dataDir string
	lifeTime                 time.Duration
	maxSizeStr               string
	idSize                   int
	maxSize                  ByteSize

	validId       *regexp.Regexp
	regexByteSize = regexp.MustCompile(`^([\d\.]+)\s*([KM]?B|[BKM])$`)
	indexTemplate *template.Template
	formTemplate  *template.Template
	data          = struct {
		sync.RWMutex
		m map[Id]PasteInfo
	}{m: make(map[Id]PasteInfo)}
	customRand *rand.Rand
)

func init() {
	flag.StringVar(&siteUrl, "u", "http://localhost:8080", "URL of the site")
	flag.StringVar(&listen, "l", "localhost:8080", "Host and port to listen to")
	flag.StringVar(&dataDir, "d", "data", "Directory to store all the pastes in")
	flag.DurationVar(&lifeTime, "t", 12*time.Hour, "Lifetime of the pastes (units: s,m,h)")
	flag.StringVar(&maxSizeStr, "s", "1M", "Maximum size of POSTs in bytes (units: B,K,M)")
	flag.IntVar(&idSize, "i", 8, "Size of the paste ids (between 6 and 256)")
	validId = regexp.MustCompile("^[a-zA-Z0-9]{" + strconv.Itoa(idSize) + "}$")
	customRand = rand.New(rand.NewSource(time.Now().UnixNano()))
}

type PasteInfo struct {
	ModTime   time.Time
	DeathTime time.Time
	Size      ByteSize
}

type Id string

func IdFromPath(idPath string) (Id, error) {
	parts := strings.Split(idPath, string(filepath.Separator))
	if len(parts) != 3 {
		return "", errors.New("Found invalid number of directories at " + idPath)
	}
	rawId := parts[0] + parts[1] + parts[2]
	if !validId.MatchString(rawId) {
		return "", errors.New("Found invalid id " + rawId)
	}
	return Id(rawId), nil
}

func RandomId() (Id, error) {
	s := make([]byte, idSize)
	var id Id
	data.RLock()
	for try := 0; try < randTries; try++ {
		var offset int = 0
	RandLoop:
		for {
			r := customRand.Int63()
			for i := 0; i < 8; i++ {
				randbyte := int(r&0xff) % len(chars)
				s[offset] = chars[randbyte]
				offset++
				if offset == idSize {
					break RandLoop
				}
				r >>= 8
			}
		}
		id = Id(s)
		if _, e := data.m[id]; !e {
			data.RUnlock()
			return id, nil
		}
	}
	return id, fmt.Errorf("Gave up trying to find an unused random id after %d tries", randTries)
}

func (id Id) String() string {
	return string(id)
}

func (id Id) Path() string {
	return path.Join(string(id[0:2]), string(id[2:4]), string(id[4:]))
}

func (id Id) EndLife() {
	data.Lock()
	err := os.Remove(id.Path())
	if err == nil {
		delete(data.m, id)
		data.Unlock()
		log.Printf("Removed paste: %s", id)
	} else {
		data.Unlock()
		log.Printf("Could not end the life of %s: %s", id, err)
		id.EndLifeAfter(2 * time.Minute)
	}
}

func (id Id) EndLifeAfter(duration time.Duration) {
	timer := time.NewTimer(duration)
	go func() {
		<-timer.C
		id.EndLife()
	}()
}

type ByteSize int64

const (
	B ByteSize = 1 << (10 * iota)
	KB
	MB
)

func parseByteSize(str string) (ByteSize, error) {
	if !regexByteSize.MatchString(str) {
		return 0, errors.New("Could not parse size in bytes")
	}
	parts := regexByteSize.FindStringSubmatch(str)
	size, _ := strconv.ParseFloat(string(parts[1]), 64)

	switch string(parts[2]) {
	case "KB", "K":
		size *= float64(KB)
	case "MB", "M":
		size *= float64(MB)
	}
	return ByteSize(size), nil
}

func (b ByteSize) String() string {
	switch {
	case b >= MB:
		return fmt.Sprintf("%.2f MB", float64(b)/float64(MB))
	case b >= KB:
		return fmt.Sprintf("%.2f KB", float64(b)/float64(KB))
	}
	return fmt.Sprintf("%d B", b)
}

func handler(w http.ResponseWriter, r *http.Request) {
	var err error
	switch r.Method {
	case "GET":
		switch r.URL.Path {
		case "/":
			indexTemplate.Execute(w, siteUrl)
			return
		case "/form":
			formTemplate.Execute(w, siteUrl)
			return
		}
		rawId := r.URL.Path[1:]
		if !validId.MatchString(rawId) {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", invalidId)
			return
		}
		id := Id(strings.ToLower(rawId))
		data.RLock()
		pasteInfo, e := data.m[id]
		data.RUnlock()
		if !e {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "%s\n", pasteNotFound)
			return
		}
		etag := fmt.Sprintf("%d-%s", pasteInfo.ModTime.Unix(), id)
		if inm := r.Header.Get("If-None-Match"); inm != "" {
			if etag == inm || inm == "*" {
				w.WriteHeader(http.StatusNotModified)
				return
			}
		}
		pastePath := id.Path()
		pasteFile, err := os.Open(pastePath)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		defer pasteFile.Close()
		w.Header().Set("Etag", etag)
		w.Header().Set("Content-Type", "text-plain; charset=utf-8")
		http.ServeContent(w, r, "", pasteInfo.ModTime, pasteFile)

	case "POST":
		r.Body = http.MaxBytesReader(w, r.Body, int64(maxSize))
		var id Id
		var content string
		id, err = RandomId()
		if err != nil {
			log.Println(err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		if err = r.ParseMultipartForm(int64(maxSize)); err != nil {
			log.Printf("Could not parse POST multipart form: %s", err)
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", err)
			return
		}
		if vs, found := r.Form["paste"]; found && len(vs[0]) > 0 {
			content = vs[0]
		} else {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", missingForm)
			return
		}
		pastePath := id.Path()
		dir, _ := path.Split(pastePath)
		if err = os.MkdirAll(dir, 0700); err != nil {
			log.Printf("Could not create directories leading to %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		deathTime := time.Now().Add(lifeTime)
		id.EndLifeAfter(lifeTime)
		pasteFile, err := os.OpenFile(pastePath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
		if err != nil {
			log.Printf("Could not create new paste pasteFile %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		defer pasteFile.Close()
		b, err := io.WriteString(pasteFile, content)
		if err != nil {
			log.Printf("Could not write data into %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		writtenSize := ByteSize(b)
		data.Lock()
		data.m[id] = PasteInfo{
			ModTime:   time.Now(),
			DeathTime: deathTime,
			Size:      writtenSize,
		}
		data.Unlock()
		log.Printf("Created a new paste: %s (%s)", id, writtenSize)
		fmt.Fprintf(w, "%s/%s\n", siteUrl, id)
	}
}

func walkFunc(filePath string, fileInfo os.FileInfo, err error) error {
	if err != nil {
		return err
	}
	if fileInfo.IsDir() {
		return nil
	}
	id, err := IdFromPath(filePath)
	if err != nil {
		return err
	}
	modTime := fileInfo.ModTime()
	deathTime := modTime.Add(lifeTime)
	now := time.Now()
	if deathTime.Before(now) {
		go id.EndLife()
		return nil
	}
	var lifeLeft time.Duration
	if deathTime.After(now.Add(lifeTime)) {
		lifeLeft = lifeTime
	} else {
		lifeLeft = deathTime.Sub(now)
	}
	size := ByteSize(fileInfo.Size())
	data.Lock()
	data.m[id] = PasteInfo{
		ModTime:   modTime,
		DeathTime: deathTime,
		Size:      size,
	}
	data.Unlock()
	log.Printf("Recovered paste %s (%s) from %s has %s left", id, size, modTime, lifeLeft)
	id.EndLifeAfter(lifeLeft)
	return nil
}

func main() {
	var err error
	flag.Parse()
	if idSize < 6 || idSize > 256 {
		log.Fatalf("Provided id size %d is not between 6 and 256", idSize)
	}
	if maxSize, err = parseByteSize(maxSizeStr); err != nil {
		log.Fatalf("Invalid max size '%s': %s", maxSizeStr, err)
	}
	if indexTemplate, err = template.ParseFiles(indexTmpl); err != nil {
		log.Fatalf("Could not load template %s: %s", indexTmpl, err)
	}
	if formTemplate, err = template.ParseFiles(formTmpl); err != nil {
		log.Fatalf("Could not load template %s: %s", formTmpl, err)
	}
	if err = os.MkdirAll(dataDir, 0700); err != nil {
		log.Fatalf("Could not create data directory %s: %s", dataDir, err)
	}
	if err = os.Chdir(dataDir); err != nil {
		log.Fatalf("Could not enter data directory %s: %s", dataDir, err)
	}
	log.Printf("idSize   = %d", idSize)
	log.Printf("maxSize  = %s", maxSize)
	log.Printf("siteUrl  = %s", siteUrl)
	log.Printf("listen   = %s", listen)
	log.Printf("dataDir  = %s", dataDir)
	log.Printf("lifeTime = %s", lifeTime)
	if err = filepath.Walk(".", walkFunc); err != nil {
		log.Fatalf("Could not recover data directory %s: %s", dataDir, err)
	}
	http.HandleFunc("/", handler)
	log.Fatal(http.ListenAndServe(listen, nil))
}
<|endoftext|>


<commit_before>// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package lib

import (
	"context"
	"flag"

	"github.com/bazelbuild/remote-apis-sdks/go/pkg/client"
	"github.com/maruel/subcommands"

	"go.chromium.org/luci/client/cas"
	"go.chromium.org/luci/common/cli"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/runtime/profiling"
)

// AuthFlags is an interface to register auth flags and create RBE Client.
type AuthFlags interface {
	// Register registers auth flags to the given flag set. e.g. -service-account-json.
	Register(f *flag.FlagSet)

	// Parse parses auth flags.
	Parse() error

	// NewClient creates an authroised RBE Client.
	NewClient(ctx context.Context, instance string) (*client.Client, error)
}

var _ cli.ContextModificator = (*commonFlags)(nil)

type commonFlags struct {
	subcommands.CommandRunBase
	casFlags  cas.Flags
	logConfig logging.Config // for -log-level, used by ModifyContext
	profiler  profiling.Profiler
	authFlags AuthFlags
}

func (c *commonFlags) Init(authFlags AuthFlags) {
	c.casFlags.Init(&c.Flags)
	c.authFlags = authFlags
	c.authFlags.Register(&c.Flags)
	c.logConfig.Level = logging.Warning
	c.logConfig.AddFlags(&c.Flags)
	c.profiler.AddFlags(&c.Flags)
}

func (c *commonFlags) Parse() error {
	// extract glog flag used in remote-apis-sdks
	logtostderr := flag.Lookup("logtostderr")
	if logtostderr == nil {
		return errors.Reason("logtostderr flag for glog not found").Err()
	}
	v := flag.Lookup("v")
	if v == nil {
		return errors.Reason("v flag for glog not found").Err()
	}

	if c.logConfig.Level == logging.Debug {
		logtostderr.Value.Set("true")
		v.Value.Set("9")
	}

	if err := c.profiler.Start(); err != nil {
		return err
	}
	if err := c.authFlags.Parse(); err != nil {
		return err
	}

	return c.casFlags.Parse()
}

// ModifyContext implements cli.ContextModificator.
func (c *commonFlags) ModifyContext(ctx context.Context) context.Context {
	return c.logConfig.Set(ctx)
}
<commit_msg>[cas] move flag lookups to inside if logging.Debug condition.<commit_after>// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package lib

import (
	"context"
	"flag"

	"github.com/bazelbuild/remote-apis-sdks/go/pkg/client"
	"github.com/maruel/subcommands"

	"go.chromium.org/luci/client/cas"
	"go.chromium.org/luci/common/cli"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/runtime/profiling"
)

// AuthFlags is an interface to register auth flags and create RBE Client.
type AuthFlags interface {
	// Register registers auth flags to the given flag set. e.g. -service-account-json.
	Register(f *flag.FlagSet)

	// Parse parses auth flags.
	Parse() error

	// NewClient creates an authroised RBE Client.
	NewClient(ctx context.Context, instance string) (*client.Client, error)
}

var _ cli.ContextModificator = (*commonFlags)(nil)

type commonFlags struct {
	subcommands.CommandRunBase
	casFlags  cas.Flags
	logConfig logging.Config // for -log-level, used by ModifyContext
	profiler  profiling.Profiler
	authFlags AuthFlags
}

func (c *commonFlags) Init(authFlags AuthFlags) {
	c.casFlags.Init(&c.Flags)
	c.authFlags = authFlags
	c.authFlags.Register(&c.Flags)
	c.logConfig.Level = logging.Warning
	c.logConfig.AddFlags(&c.Flags)
	c.profiler.AddFlags(&c.Flags)
}

func (c *commonFlags) Parse() error {
	if c.logConfig.Level == logging.Debug {
		// extract glog flag used in remote-apis-sdks
		logtostderr := flag.Lookup("logtostderr")
		if logtostderr == nil {
			return errors.Reason("logtostderr flag for glog not found").Err()
		}
		v := flag.Lookup("v")
		if v == nil {
			return errors.Reason("v flag for glog not found").Err()
		}
		logtostderr.Value.Set("true")
		v.Value.Set("9")
	}

	if err := c.profiler.Start(); err != nil {
		return err
	}
	if err := c.authFlags.Parse(); err != nil {
		return err
	}

	return c.casFlags.Parse()
}

// ModifyContext implements cli.ContextModificator.
func (c *commonFlags) ModifyContext(ctx context.Context) context.Context {
	return c.logConfig.Set(ctx)
}
<|endoftext|>