<commit_before>package instagram

import (
	"testing"
)

func Test_GetAccoutByUsername_exist(t *testing.T) {
	for _, testCase := range []struct {
		username, fullName, id string
	}{
		{"instagram", "Instagram", "25025320"},
		{"solidlsnake", "Konstantin", "248188406"},
	} {
		account, err := GetAccountByUsername(testCase.username)
		if err != nil ||
			account.Username != testCase.username ||
			account.FullName != testCase.fullName ||
			account.ID != testCase.id {
			t.Error("Unexpected account info.")
		}
	}
}

func Test_GetAccoutByUsername_notExist(t *testing.T) {
	for _, testCase := range []struct {
		username string
	}{
		{"fhusdhfjashbfjfghyashf"},
		{"fhusadhfyasifjasduiash"},
	} {
		_, err := GetAccountByUsername(testCase.username)
		if err == nil {
			t.Error("Unexpected account info.")
		}
	}
}

func Test_GetMediaByUrl(t *testing.T) {
	for _, testCase := range []struct {
		url, code, username, mediaType string
	}{
		{
			"https://www.instagram.com/p/ceiqEstT6r/",
			"ceiqEstT6r",
			"solidlsnake",
			"image",
		},
		{
			"https://www.instagram.com/p/12376OtT5o/",
			"12376OtT5o",
			"solidlsnake",
			"video",
		},
	} {
		media, err := GetMediaByURL(testCase.url)
		if err != nil ||
			media.Code != testCase.code ||
			media.Owner.Username != testCase.username ||
			media.Type != testCase.mediaType {
			t.Error("Unexpected media info.")
		}
	}
}

func Test_GetMediaByCode(t *testing.T) {
	for _, testCase := range []struct {
		code, username, mediaType string
	}{
		{
			"ceiqEstT6r",
			"solidlsnake",
			"image",
		},
		{
			"12376OtT5o",
			"solidlsnake",
			"video",
		},
	} {
		media, err := GetMediaByCode(testCase.code)
		if err != nil ||
			media.Code != testCase.code ||
			media.Owner.Username != testCase.username ||
			media.Type != testCase.mediaType {
			t.Error("Unexpected media info.")
		}
	}
}

func Test_GetMediaByCode_notExist(t *testing.T) {
	for _, testCase := range []struct {
		code string
	}{
		{"aaaaaaaaaa"},
		{"abcdefghij"},
	} {
		_, err := GetMediaByCode(testCase.code)
		if err == nil {
			t.Error("Unexpected media info.")
		}
	}
}

func Test_GetUserMedia_quantity(t *testing.T) {
	account, _ := GetAccountByUsername("solidlsnake")
	count := int(account.MediaCount)

	for _, testCase := range []struct {
		username string
		quantity uint16
		expected int
	}{
		{"instagram", 10, 10},
		{"solidlsnake", 999, count},
	} {
		medias, err := GetAccountMedia(testCase.username, testCase.quantity)
		if err != nil || len(medias) != testCase.expected {
			t.Error("Wrong numbers of media.")
		}
	}
}

func Test_GetAllUserMedia_quantity(t *testing.T) {
	for _, testCase := range []struct {
		username string
	}{
		{"eminem"},
		{"solidlsnake"},
	} {
		account, _ := GetAccountByUsername(testCase.username)
		expected := int(account.MediaCount)
		medias, err := GetAllAccountMedia(testCase.username)
		if err != nil || len(medias) != expected {
			t.Error("Wrong numbers of media.")
		}
	}
}

func Test_GetLocationMedia_quantity(t *testing.T) {
	for _, testCase := range []struct {
		locationID string
		quantity   uint16
	}{
		{"17326249", 10},
		{"17326249", 25},
	} {
		medias, err := GetLocationMedia(testCase.locationID, testCase.quantity)
		if err != nil || len(medias) != int(testCase.quantity) {
			t.Error("Wrong numbers of media.")
		}
	}
}

func Test_GetLocationTopMedia_quantity(t *testing.T) {
	for _, testCase := range []struct {
		locationID string
	}{
		{"17326249"},
	} {
		medias, err := GetLocationTopMedia(testCase.locationID)
		if err != nil || len(medias) != 9 {
			t.Error("Wrong numbers of media.")
		}
	}
}

func Test_GetLocationById(t *testing.T) {
	for _, testCase := range []struct {
		locationID, slug string
	}{
		{"17326249", "moscow-russia"},
		{"212988663", "new-york-new-york"},
	} {
		location, err := GetLocationByID(testCase.locationID)
		if err != nil ||
			location.Slug != testCase.slug {
			t.Error("Wrong location info")
		}
	}
}

func Test_GetTagMedia_quantity(t *testing.T) {
	for _, testCase := range []struct {
		tag      string
		quantity uint16
	}{
		{"lol", 10},
		{"haha", 25},
	} {
		medias, err := GetTagMedia(testCase.tag, testCase.quantity)
		if err != nil || len(medias) != int(testCase.quantity) {
			t.Error("Wrong numbers of media.")
		}
	}
}

func Test_GetTagTopMedia_quantity(t *testing.T) {
	for _, testCase := range []struct {
		tag string
	}{
		{"lol"},
	} {
		medias, err := GetTagTopMedia(testCase.tag)
		if err != nil || len(medias) != 9 {
			t.Error("Wrong numbers of media.")
		}
	}
}

func Test_SearchForUsers(t *testing.T) {
	for _, testCase := range []struct {
		username, id string
	}{
		{"solidlsnake", "248188406" },
	} {
		accounts, err := SearchForUsers(testCase.username)
		if err != nil ||
			accounts[0].ID != testCase.id {
			t.Error("Wrong account information.")
		}
	}
}
<commit_msg>Better tests output<commit_after>package instagram

import (
	"testing"
)

func Test_GetAccoutByUsername_exist(t *testing.T) {
	for _, testCase := range []struct {
		username, fullName, id string
	}{
		{"instagram", "Instagram", "25025320"},
		{"solidlsnake", "Konstantin", "248188406"},
	} {
		account, err := GetAccountByUsername(testCase.username)
		if err != nil ||
			account.Username != testCase.username ||
			account.FullName != testCase.fullName ||
			account.ID != testCase.id {
			t.Error("Unexpected account info.")
		}
	}
}

func Test_GetAccoutByUsername_notExist(t *testing.T) {
	for _, testCase := range []struct {
		username string
	}{
		{"fhusdhfjashbfjfghyashf"},
		{"fhusadhfyasifjasduiash"},
	} {
		_, err := GetAccountByUsername(testCase.username)
		if err == nil {
			t.Error("Unexpected account info.")
		}
	}
}

func Test_GetMediaByUrl(t *testing.T) {
	for _, testCase := range []struct {
		url, code, username, mediaType string
	}{
		{
			"https://www.instagram.com/p/ceiqEstT6r/",
			"ceiqEstT6r",
			"solidlsnake",
			"image",
		},
		{
			"https://www.instagram.com/p/12376OtT5o/",
			"12376OtT5o",
			"solidlsnake",
			"video",
		},
	} {
		media, err := GetMediaByURL(testCase.url)
		if err != nil ||
			media.Code != testCase.code ||
			media.Owner.Username != testCase.username ||
			media.Type != testCase.mediaType {
			t.Error("Unexpected media info.")
		}
	}
}

func Test_GetMediaByCode(t *testing.T) {
	for _, testCase := range []struct {
		code, username, mediaType string
	}{
		{
			"ceiqEstT6r",
			"solidlsnake",
			"image",
		},
		{
			"12376OtT5o",
			"solidlsnake",
			"video",
		},
	} {
		media, err := GetMediaByCode(testCase.code)
		if err != nil ||
			media.Code != testCase.code ||
			media.Owner.Username != testCase.username ||
			media.Type != testCase.mediaType {
			t.Error("Unexpected media info.")
		}
	}
}

func Test_GetMediaByCode_notExist(t *testing.T) {
	for _, testCase := range []struct {
		code string
	}{
		{"aaaaaaaaaa"},
		{"abcdefghij"},
	} {
		_, err := GetMediaByCode(testCase.code)
		if err == nil {
			t.Error("Unexpected media info.")
		}
	}
}

func Test_GetUserMedia_quantity(t *testing.T) {
	account, _ := GetAccountByUsername("solidlsnake")
	count := int(account.MediaCount)

	for _, testCase := range []struct {
		username string
		quantity uint16
		expected int
	}{
		{"instagram", 10, 10},
		{"solidlsnake", 999, count},
	} {
		medias, err := GetAccountMedia(testCase.username, testCase.quantity)
		if err != nil || len(medias) != testCase.expected {
			t.Error("Wrong numbers of media.")
		}
	}
}

func Test_GetAllUserMedia_quantity(t *testing.T) {
	for _, testCase := range []struct {
		username string
	}{
		{"eminem"},
		{"solidlsnake"},
	} {
		account, _ := GetAccountByUsername(testCase.username)
		expected := int(account.MediaCount)
		medias, err := GetAllAccountMedia(testCase.username)
		if err != nil || len(medias) != expected {
			t.Error("Wrong numbers of media.")
		}
	}
}

func Test_GetLocationMedia_quantity(t *testing.T) {
	for _, testCase := range []struct {
		locationID string
		quantity   uint16
	}{
		{"17326249", 10},
		{"17326249", 25},
	} {
		medias, err := GetLocationMedia(testCase.locationID, testCase.quantity)
		if err != nil || len(medias) != int(testCase.quantity) {
			t.Errorf("Wrong numbers of media. Expect %d, get %s.", testCase.quantity, len(medias))
		}
	}
}

func Test_GetLocationTopMedia_quantity(t *testing.T) {
	for _, testCase := range []struct {
		locationID string
	}{
		{"17326249"},
	} {
		medias, err := GetLocationTopMedia(testCase.locationID)
		if err != nil || len(medias) != 9 {
			t.Error("Wrong numbers of media.")
		}
	}
}

func Test_GetLocationById(t *testing.T) {
	for _, testCase := range []struct {
		locationID, slug string
	}{
		{"17326249", "moscow-russia"},
		{"212988663", "new-york-new-york"},
	} {
		location, err := GetLocationByID(testCase.locationID)
		if err != nil ||
			location.Slug != testCase.slug {
			t.Errorf("Wrong location info. Expect slug %s, but get %s.", testCase.slug, location.Slug )
		}
	}
}

func Test_GetTagMedia_quantity(t *testing.T) {
	for _, testCase := range []struct {
		tag      string
		quantity uint16
	}{
		{"lol", 10},
		{"haha", 25},
	} {
		medias, err := GetTagMedia(testCase.tag, testCase.quantity)
		if err != nil || len(medias) != int(testCase.quantity) {
			t.Error("Wrong numbers of media.")
		}
	}
}

func Test_GetTagTopMedia_quantity(t *testing.T) {
	for _, testCase := range []struct {
		tag string
	}{
		{"lol"},
	} {
		medias, err := GetTagTopMedia(testCase.tag)
		if err != nil || len(medias) != 9 {
			t.Error("Wrong numbers of media.")
		}
	}
}

func Test_SearchForUsers(t *testing.T) {
	for _, testCase := range []struct {
		username, id string
	}{
		{"solidlsnake", "248188406" },
	} {
		accounts, err := SearchForUsers(testCase.username)
		if err != nil ||
			accounts[0].ID != testCase.id {
			t.Error("Wrong account information.")
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha256"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"runtime"
	"time"

	"github.com/aws/aws-lambda-go/lambda"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sns"
	"go.mozilla.org/autograph/formats"
	"go.mozilla.org/autograph/signer/apk"
	"go.mozilla.org/autograph/signer/apk2"
	"go.mozilla.org/autograph/signer/contentsignature"
	"go.mozilla.org/autograph/signer/contentsignaturepki"
	"go.mozilla.org/autograph/signer/genericrsa"
	"go.mozilla.org/autograph/signer/gpg2"
	"go.mozilla.org/autograph/signer/mar"
	"go.mozilla.org/autograph/signer/pgp"
	"go.mozilla.org/autograph/signer/rsapss"
	"go.mozilla.org/autograph/signer/xpi"
	"go.mozilla.org/hawk"
)

type configuration struct {
	url           string
	monitoringKey string
	env           string
	rootHash      string
	truststore    *x509.CertPool
	// hash keystore for verifying XPI dep signers
	depRootHash   string
	depTruststore *x509.CertPool
}

var conf configuration

const inputdata string = "AUTOGRAPH MONITORING"

var softNotifCache map[string]time.Time

func init() {
	// create a cache to avoid sending the same notifications over and over
	softNotifCache = make(map[string]time.Time)
}

func main() {
	conf.url = os.Getenv("AUTOGRAPH_URL")
	if conf.url == "" {
		log.Fatal("AUTOGRAPH_URL must be set to the base url of the autograph service")
	}
	conf.monitoringKey = os.Getenv("AUTOGRAPH_KEY")
	if conf.monitoringKey == "" {
		log.Fatal("AUTOGRAPH_KEY must be set to the api monitoring key")
	}
	conf.env = os.Getenv("AUTOGRAPH_ENV")
	switch conf.env {
	case "stage":
		conf.rootHash = firefoxPkiStageRootHash
		conf.truststore = x509.NewCertPool()
		conf.truststore.AppendCertsFromPEM([]byte(firefoxPkiStageRoot))
		conf.depRootHash = ""
		conf.depTruststore = nil
	case "prod":
		conf.rootHash = firefoxPkiProdRootHash
		conf.truststore = x509.NewCertPool()
		conf.truststore.AppendCertsFromPEM([]byte(firefoxPkiProdRoot))
		conf.depRootHash = firefoxPkiStageRootHash
		conf.depTruststore = x509.NewCertPool()
		conf.depTruststore.AppendCertsFromPEM([]byte(firefoxPkiStageRoot))
	default:
		conf.rootHash = autographDevRootHash
		conf.truststore = nil
		conf.depRootHash = ""
		conf.depTruststore = nil
	}
	if os.Getenv("AUTOGRAPH_ROOT_HASH") != "" {
		conf.rootHash = os.Getenv("AUTOGRAPH_ROOT_HASH")
	}
	if os.Getenv("LAMBDA_TASK_ROOT") != "" {
		// we are inside a lambda environment so run as lambda
		lambda.Start(Handler)
	} else {
		err := Handler()
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %v\n", err)
			os.Exit(1)
		}
		os.Exit(0)
	}
}

// Handler is a wrapper around monitor() that performs garbage collection
// before returning
func Handler() (err error) {
	defer func() {
		// force gc run
		// https://bugzilla.mozilla.org/show_bug.cgi?id=1621133
		t1 := time.Now()
		runtime.GC()
		log.Println("Garbage collected in", time.Now().Sub(t1))
	}()
	return monitor()
}

// monitor contacts the autograph service and verifies all monitoring signatures
func monitor() (err error) {
	log.Println("Retrieving monitoring data from", conf.url)
	req, err := http.NewRequest("GET", conf.url+"__monitor__", nil)
	if err != nil {
		return
	}

	// For client requests, setting this field prevents re-use of
	// TCP connections between requests to the same hosts, as if
	// Transport.DisableKeepAlives were set.
	req.Close = true

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", makeAuthHeader(req, "monitor", conf.monitoringKey))
	cli := &http.Client{}
	resp, err := cli.Do(req)
	if err != nil || resp == nil {
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusCreated {
		return fmt.Errorf("Request failed with %s: %s", resp.Status, resp.Body)
	}

	dec := json.NewDecoder(resp.Body)
	failed := false
	var failures []error
	for {
		var response formats.SignatureResponse
		if err := dec.Decode(&response); err == io.EOF {
			break
		} else if err != nil {
			log.Fatal(err)
		}
		switch response.Type {
		case contentsignature.Type:
			log.Printf("Verifying content signature from signer %q", response.SignerID)
			err = verifyContentSignature(conf.rootHash, response)
		case contentsignaturepki.Type:
			log.Printf("Verifying content signature pki from signer %q", response.SignerID)
			err = verifyContentSignature(conf.rootHash, response)
		case xpi.Type:
			log.Printf("Verifying XPI signature from signer %q", response.SignerID)
			err = verifyXPISignature(response.Signature)
		case apk.Type:
			log.Printf("Verifying APK signature from signer %q", response.SignerID)
			err = verifyAPKSignature(response.Signature)
		case apk2.Type:
			// we don't verify apk2 signatures because they can only be obtained on valid
			// APK files, which is too big to fit in the monitoring logic
			log.Printf("Skipping verification of APK2 signature from signer %q (we can't verify those)", response.SignerID)
			continue
		case mar.Type:
			log.Printf("Verifying MAR signature from signer %q", response.SignerID)
			err = verifyMARSignature(response.Signature, response.PublicKey)
		case genericrsa.Type:
			log.Printf("Verifying RSA signature from signer %q", response.SignerID)
			err = genericrsa.VerifyGenericRsaSignatureResponse([]byte(inputdata), response)
		case rsapss.Type:
			log.Printf("Verifying RSA-PSS signature from signer %q", response.SignerID)
			err = verifyRsapssSignature(response.Signature, response.PublicKey)
		case pgp.Type, gpg2.Type:
			// we don't verify pgp signatures because that requires building a keyring
			// using the public key which is hard to do using the current openpgp package
			log.Printf("Skipping verification of PGP signature from signer %q (we can't verify those)", response.SignerID)
			continue
		default:
			err = fmt.Errorf("unknown signature type %q", response.Type)
		}
		if err != nil {
			failed = true
			log.Printf("Response from signer %q does not pass: %v", response.SignerID, err)
			log.Printf("Response was: %+v", response)
			failures = append(failures, err)
		} else {
			log.Printf("Response from signer %q passes verification", response.SignerID)
		}
	}
	if failed {
		failure := "Errors found during monitoring:"
		for i, fail := range failures {
			failure += fmt.Sprintf("\n%d. %s", i+1, fail.Error())
		}
		return fmt.Errorf(failure)
	}
	log.Println("All signature responses passed, monitoring OK")
	return
}

func makeAuthHeader(req *http.Request, user, token string) string {
	auth := hawk.NewRequestAuth(req,
		&hawk.Credentials{
			ID:   user,
			Key:  token,
			Hash: sha256.New},
		0)
	auth.Ext = fmt.Sprintf("%d", time.Now().Nanosecond())
	payloadhash := auth.PayloadHash("application/json")
	payloadhash.Write([]byte(""))
	auth.SetHash(payloadhash)
	return auth.RequestHeader()
}

// send a message to a predefined sns topic
func sendSoftNotification(id string, format string, a ...interface{}) error {
	if ts, ok := softNotifCache[id]; ok {
		// don't send dup notifications for 24 hours
		if ts.Add(24 * time.Hour).After(time.Now()) {
			log.Printf("silencing soft notification ID %s", id)
			return nil
		}
	}
	if os.Getenv("LAMBDA_TASK_ROOT") == "" || os.Getenv("AUTOGRAPH_SOFT_NOTIFICATION_SNS") == "" {
		// We're not running in lambda or the conf isnt ready so don't try to publish to SQS
		log.Printf("soft notification ID %s: %s", id, fmt.Sprintf(format, a...))
		return nil
	}

	svc := sns.New(session.New())
	params := &sns.PublishInput{
		Message:  aws.String(fmt.Sprintf(format, a...)),
		TopicArn: aws.String(os.Getenv("AUTOGRAPH_SOFT_NOTIFICATION_SNS")),
	}
	_, err := svc.Publish(params)
	if err != nil {
		return err
	}
	log.Printf("Soft notification send to %q with body: %s", os.Getenv("AUTOGRAPH_SOFT_NOTIFICATION_SNS"), *params.Message)
	// add the notification to the cache
	softNotifCache[id] = time.Now()
	return nil
}
<commit_msg>tools: monitor: add content signature trust store and hash<commit_after>package main

import (
	"crypto/sha256"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"runtime"
	"time"

	"github.com/aws/aws-lambda-go/lambda"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sns"
	"go.mozilla.org/autograph/formats"
	"go.mozilla.org/autograph/signer/apk"
	"go.mozilla.org/autograph/signer/apk2"
	"go.mozilla.org/autograph/signer/contentsignature"
	"go.mozilla.org/autograph/signer/contentsignaturepki"
	"go.mozilla.org/autograph/signer/genericrsa"
	"go.mozilla.org/autograph/signer/gpg2"
	"go.mozilla.org/autograph/signer/mar"
	"go.mozilla.org/autograph/signer/pgp"
	"go.mozilla.org/autograph/signer/rsapss"
	"go.mozilla.org/autograph/signer/xpi"
	"go.mozilla.org/hawk"
)

type configuration struct {
	url           string
	monitoringKey string
	env           string
	rootHash      string
	truststore    *x509.CertPool

	// hash and keystore for verifying XPI dep signers
	depRootHash   string
	depTruststore *x509.CertPool

	// hash and keystore for verifying ContentSignature responses
	// when XPI and ContentSignature branches of the code signing
	// PKI use different roots
	contentSignatureRootHash   string
	contentSignatureTruststore *x509.CertPool
}

var conf configuration

const inputdata string = "AUTOGRAPH MONITORING"

var softNotifCache map[string]time.Time

func init() {
	// create a cache to avoid sending the same notifications over and over
	softNotifCache = make(map[string]time.Time)
}

func main() {
	conf.url = os.Getenv("AUTOGRAPH_URL")
	if conf.url == "" {
		log.Fatal("AUTOGRAPH_URL must be set to the base url of the autograph service")
	}
	conf.monitoringKey = os.Getenv("AUTOGRAPH_KEY")
	if conf.monitoringKey == "" {
		log.Fatal("AUTOGRAPH_KEY must be set to the api monitoring key")
	}

	// configure monitor to check responses against Fx stage or
	// prod or autograph dev code signing PKI roots and CA root
	// certs defined in constants.go
	conf.env = os.Getenv("AUTOGRAPH_ENV")
	switch conf.env {
	case "stage":
		conf.rootHash = firefoxPkiStageRootHash
		conf.truststore = x509.NewCertPool()
		conf.truststore.AppendCertsFromPEM([]byte(firefoxPkiStageRoot))
		conf.contentSignatureRootHash = firefoxPkiContentSignatureStageRootHash
		conf.contentSignatureTruststore = x509.NewCertPool()
		conf.contentSignatureTruststore.AppendCertsFromPEM([]byte(firefoxPkiContentSignatureStageRoot))
		conf.depRootHash = ""
		conf.depTruststore = nil
	case "prod":
		conf.rootHash = firefoxPkiProdRootHash
		conf.truststore = x509.NewCertPool()
		conf.truststore.AppendCertsFromPEM([]byte(firefoxPkiProdRoot))
		conf.contentSignatureRootHash = firefoxPkiProdRootHash
		conf.contentSignatureTruststore = conf.truststore
		conf.depRootHash = firefoxPkiStageRootHash
		conf.depTruststore = x509.NewCertPool()
		conf.depTruststore.AppendCertsFromPEM([]byte(firefoxPkiStageRoot))
	default:
		conf.rootHash = autographDevRootHash
		conf.truststore = nil
		conf.contentSignatureRootHash = autographDevRootHash
		conf.contentSignatureTruststore = nil
		conf.depRootHash = ""
		conf.depTruststore = nil
	}
	if os.Getenv("AUTOGRAPH_ROOT_HASH") != "" {
		conf.rootHash = os.Getenv("AUTOGRAPH_ROOT_HASH")
		conf.contentSignatureRootHash = conf.rootHash
	}
	if os.Getenv("LAMBDA_TASK_ROOT") != "" {
		// we are inside a lambda environment so run as lambda
		lambda.Start(Handler)
	} else {
		err := Handler()
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %v\n", err)
			os.Exit(1)
		}
		os.Exit(0)
	}
}

// Handler is a wrapper around monitor() that performs garbage collection
// before returning
func Handler() (err error) {
	defer func() {
		// force gc run
		// https://bugzilla.mozilla.org/show_bug.cgi?id=1621133
		t1 := time.Now()
		runtime.GC()
		log.Println("Garbage collected in", time.Now().Sub(t1))
	}()
	return monitor()
}

// monitor contacts the autograph service and verifies all monitoring signatures
func monitor() (err error) {
	log.Println("Retrieving monitoring data from", conf.url)
	req, err := http.NewRequest("GET", conf.url+"__monitor__", nil)
	if err != nil {
		return
	}

	// For client requests, setting this field prevents re-use of
	// TCP connections between requests to the same hosts, as if
	// Transport.DisableKeepAlives were set.
	req.Close = true

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", makeAuthHeader(req, "monitor", conf.monitoringKey))
	cli := &http.Client{}
	resp, err := cli.Do(req)
	if err != nil || resp == nil {
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusCreated {
		return fmt.Errorf("Request failed with %s: %s", resp.Status, resp.Body)
	}

	dec := json.NewDecoder(resp.Body)
	failed := false
	var failures []error
	for {
		var response formats.SignatureResponse
		if err := dec.Decode(&response); err == io.EOF {
			break
		} else if err != nil {
			log.Fatal(err)
		}
		switch response.Type {
		case contentsignature.Type:
			log.Printf("Verifying content signature from signer %q", response.SignerID)
			err = verifyContentSignature(conf.contentSignatureRootHash, response)
		case contentsignaturepki.Type:
			log.Printf("Verifying content signature pki from signer %q", response.SignerID)
			err = verifyContentSignature(conf.contentSignatureRootHash, response)
		case xpi.Type:
			log.Printf("Verifying XPI signature from signer %q", response.SignerID)
			err = verifyXPISignature(response.Signature)
		case apk.Type:
			log.Printf("Verifying APK signature from signer %q", response.SignerID)
			err = verifyAPKSignature(response.Signature)
		case apk2.Type:
			// we don't verify apk2 signatures because they can only be obtained on valid
			// APK files, which is too big to fit in the monitoring logic
			log.Printf("Skipping verification of APK2 signature from signer %q (we can't verify those)", response.SignerID)
			continue
		case mar.Type:
			log.Printf("Verifying MAR signature from signer %q", response.SignerID)
			err = verifyMARSignature(response.Signature, response.PublicKey)
		case genericrsa.Type:
			log.Printf("Verifying RSA signature from signer %q", response.SignerID)
			err = genericrsa.VerifyGenericRsaSignatureResponse([]byte(inputdata), response)
		case rsapss.Type:
			log.Printf("Verifying RSA-PSS signature from signer %q", response.SignerID)
			err = verifyRsapssSignature(response.Signature, response.PublicKey)
		case pgp.Type, gpg2.Type:
			// we don't verify pgp signatures because that requires building a keyring
			// using the public key which is hard to do using the current openpgp package
			log.Printf("Skipping verification of PGP signature from signer %q (we can't verify those)", response.SignerID)
			continue
		default:
			err = fmt.Errorf("unknown signature type %q", response.Type)
		}
		if err != nil {
			failed = true
			log.Printf("Response from signer %q does not pass: %v", response.SignerID, err)
			log.Printf("Response was: %+v", response)
			failures = append(failures, err)
		} else {
			log.Printf("Response from signer %q passes verification", response.SignerID)
		}
	}
	if failed {
		failure := "Errors found during monitoring:"
		for i, fail := range failures {
			failure += fmt.Sprintf("\n%d. %s", i+1, fail.Error())
		}
		return fmt.Errorf(failure)
	}
	log.Println("All signature responses passed, monitoring OK")
	return
}

func makeAuthHeader(req *http.Request, user, token string) string {
	auth := hawk.NewRequestAuth(req,
		&hawk.Credentials{
			ID:   user,
			Key:  token,
			Hash: sha256.New},
		0)
	auth.Ext = fmt.Sprintf("%d", time.Now().Nanosecond())
	payloadhash := auth.PayloadHash("application/json")
	payloadhash.Write([]byte(""))
	auth.SetHash(payloadhash)
	return auth.RequestHeader()
}

// send a message to a predefined sns topic
func sendSoftNotification(id string, format string, a ...interface{}) error {
	if ts, ok := softNotifCache[id]; ok {
		// don't send dup notifications for 24 hours
		if ts.Add(24 * time.Hour).After(time.Now()) {
			log.Printf("silencing soft notification ID %s", id)
			return nil
		}
	}
	if os.Getenv("LAMBDA_TASK_ROOT") == "" || os.Getenv("AUTOGRAPH_SOFT_NOTIFICATION_SNS") == "" {
		// We're not running in lambda or the conf isnt ready so don't try to publish to SQS
		log.Printf("soft notification ID %s: %s", id, fmt.Sprintf(format, a...))
		return nil
	}

	svc := sns.New(session.New())
	params := &sns.PublishInput{
		Message:  aws.String(fmt.Sprintf(format, a...)),
		TopicArn: aws.String(os.Getenv("AUTOGRAPH_SOFT_NOTIFICATION_SNS")),
	}
	_, err := svc.Publish(params)
	if err != nil {
		return err
	}
	log.Printf("Soft notification send to %q with body: %s", os.Getenv("AUTOGRAPH_SOFT_NOTIFICATION_SNS"), *params.Message)
	// add the notification to the cache
	softNotifCache[id] = time.Now()
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2017 The etcd-operator 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 controller

import (
	api "github.com/coreos/etcd-operator/pkg/apis/etcd/v1beta2"

	"github.com/sirupsen/logrus"
)

const (
	// Copy from deployment_controller.go:
	// maxRetries is the number of times a etcd backup will be retried before it is dropped out of the queue.
	// With the current rate-limiter in use (5ms*2^(maxRetries-1)) the following numbers represent the times
	// an etcd backup is going to be requeued:
	//
	// 5ms, 10ms, 20ms, 40ms, 80ms, 160ms, 320ms, 640ms, 1.3s, 2.6s, 5.1s, 10.2s, 20.4s, 41s, 82s
	maxRetries = 15
)

func (b *Backup) runWorker() {
	for b.processNextItem() {
	}
}

func (b *Backup) processNextItem() bool {
	// Wait until there is a new item in the working queue
	key, quit := b.queue.Get()
	if quit {
		return false
	}
	// Tell the queue that we are done with processing this key. This unblocks the key for other workers
	// This allows safe parallel processing because two pods with the same key are never processed in
	// parallel.
	defer b.queue.Done(key)
	err := b.processItem(key.(string))
	// Handle the error if something went wrong during the execution of the business logic
	b.handleErr(err, key)
	return true
}

func (b *Backup) processItem(key string) error {
	obj, exists, err := b.indexer.GetByKey(key)
	if err != nil {
		return err
	}
	if !exists {
		return nil
	}

	eb := obj.(*api.EtcdBackup)
	// don't process the CR if it has a status since
	// having a status means that the CR has been processed before.
	if eb.Status.Succeeded || len(eb.Status.Reason) != 0 {
		return nil
	}
	err = b.handleBackup(&eb.Spec)
	// Report backup status
	b.reportBackupStatus(err, eb)
	return err
}

func (b *Backup) reportBackupStatus(err error, eb *api.EtcdBackup) {
	if err != nil {
		eb.Status.Succeeded = false
		eb.Status.Reason = err.Error()
	} else {
		eb.Status.Succeeded = true
	}
	_, err = b.backupCRCli.EtcdV1beta2().EtcdBackups(b.namespace).Update(eb)
	if err != nil {
		b.logger.Warningf("failed to update status of backup CR %v : (%v)", eb.Name, err)
	}
}

func (b *Backup) handleErr(err error, key interface{}) {
	if err == nil {
		// Forget about the #AddRateLimited history of the key on every successful synchronization.
		// This ensures that future processing of updates for this key is not delayed because of
		// an outdated error history.
		b.queue.Forget(key)
		return
	}

	// This controller retries maxRetries times if something goes wrong. After that, it stops trying.
	if b.queue.NumRequeues(key) < maxRetries {
		b.logger.Errorf("error syncing etcd backup (%v): %v", key, err)

		// Re-enqueue the key rate limited. Based on the rate limiter on the
		// queue and the re-enqueue history, the key will be processed later again.
		b.queue.AddRateLimited(key)
		return
	}

	b.queue.Forget(key)
	// Report that, even after several retries, we could not successfully process this key
	b.logger.Infof("Dropping etcd backup (%v) out of the queue: %v", key, err)
}

func (b *Backup) handleBackup(spec *api.BackupSpec) error {
	switch spec.StorageType {
	case api.BackupStorageTypeS3:
		return handleS3(b.kubecli, spec.S3, b.namespace, spec.ClusterName)
	default:
		logrus.Fatalf("unknown StorageType: %v", spec.StorageType)
	}
	return nil
}
<commit_msg>controller/backup-operator: use berr instead of err for reportBackupStatus<commit_after>// Copyright 2017 The etcd-operator 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 controller

import (
	api "github.com/coreos/etcd-operator/pkg/apis/etcd/v1beta2"

	"github.com/sirupsen/logrus"
)

const (
	// Copy from deployment_controller.go:
	// maxRetries is the number of times a etcd backup will be retried before it is dropped out of the queue.
	// With the current rate-limiter in use (5ms*2^(maxRetries-1)) the following numbers represent the times
	// an etcd backup is going to be requeued:
	//
	// 5ms, 10ms, 20ms, 40ms, 80ms, 160ms, 320ms, 640ms, 1.3s, 2.6s, 5.1s, 10.2s, 20.4s, 41s, 82s
	maxRetries = 15
)

func (b *Backup) runWorker() {
	for b.processNextItem() {
	}
}

func (b *Backup) processNextItem() bool {
	// Wait until there is a new item in the working queue
	key, quit := b.queue.Get()
	if quit {
		return false
	}
	// Tell the queue that we are done with processing this key. This unblocks the key for other workers
	// This allows safe parallel processing because two pods with the same key are never processed in
	// parallel.
	defer b.queue.Done(key)
	err := b.processItem(key.(string))
	// Handle the error if something went wrong during the execution of the business logic
	b.handleErr(err, key)
	return true
}

func (b *Backup) processItem(key string) error {
	obj, exists, err := b.indexer.GetByKey(key)
	if err != nil {
		return err
	}
	if !exists {
		return nil
	}

	eb := obj.(*api.EtcdBackup)
	// don't process the CR if it has a status since
	// having a status means that the CR has been processed before.
	if eb.Status.Succeeded || len(eb.Status.Reason) != 0 {
		return nil
	}
	err = b.handleBackup(&eb.Spec)
	// Report backup status
	b.reportBackupStatus(err, eb)
	return err
}

func (b *Backup) reportBackupStatus(berr error, eb *api.EtcdBackup) {
	if berr != nil {
		eb.Status.Succeeded = false
		eb.Status.Reason = berr.Error()
	} else {
		eb.Status.Succeeded = true
	}
	_, err := b.backupCRCli.EtcdV1beta2().EtcdBackups(b.namespace).Update(eb)
	if err != nil {
		b.logger.Warningf("failed to update status of backup CR %v : (%v)", eb.Name, err)
	}
}

func (b *Backup) handleErr(err error, key interface{}) {
	if err == nil {
		// Forget about the #AddRateLimited history of the key on every successful synchronization.
		// This ensures that future processing of updates for this key is not delayed because of
		// an outdated error history.
		b.queue.Forget(key)
		return
	}

	// This controller retries maxRetries times if something goes wrong. After that, it stops trying.
	if b.queue.NumRequeues(key) < maxRetries {
		b.logger.Errorf("error syncing etcd backup (%v): %v", key, err)

		// Re-enqueue the key rate limited. Based on the rate limiter on the
		// queue and the re-enqueue history, the key will be processed later again.
		b.queue.AddRateLimited(key)
		return
	}

	b.queue.Forget(key)
	// Report that, even after several retries, we could not successfully process this key
	b.logger.Infof("Dropping etcd backup (%v) out of the queue: %v", key, err)
}

func (b *Backup) handleBackup(spec *api.BackupSpec) error {
	switch spec.StorageType {
	case api.BackupStorageTypeS3:
		return handleS3(b.kubecli, spec.S3, b.namespace, spec.ClusterName)
	default:
		logrus.Fatalf("unknown StorageType: %v", spec.StorageType)
	}
	return nil
}
<|endoftext|>


<commit_before>package micro

import (
	"os"
	"os/signal"
	rtime "runtime"
	"strings"
	"sync"

	"github.com/micro/go-micro/v2/auth"
	"github.com/micro/go-micro/v2/client"
	"github.com/micro/go-micro/v2/config/cmd"
	"github.com/micro/go-micro/v2/debug/service/handler"
	"github.com/micro/go-micro/v2/debug/stats"
	"github.com/micro/go-micro/v2/debug/trace"
	"github.com/micro/go-micro/v2/logger"
	"github.com/micro/go-micro/v2/plugin"
	"github.com/micro/go-micro/v2/runtime"
	"github.com/micro/go-micro/v2/server"
	"github.com/micro/go-micro/v2/store"
	signalutil "github.com/micro/go-micro/v2/util/signal"
	"github.com/micro/go-micro/v2/util/wrapper"
)

type service struct {
	opts Options

	once sync.Once
}

func newService(opts ...Option) Service {
	service := new(service)
	options := newOptions(opts...)

	// service name
	serviceName := options.Server.Options().Name

	// authFn returns the auth, we pass as a function since auth
	// has not yet been set at this point.
	authFn := func() auth.Auth { return options.Server.Options().Auth }

	// wrap client to inject From-Service header on any calls
	options.Client = wrapper.FromService(serviceName, options.Client)
	options.Client = wrapper.TraceCall(serviceName, trace.DefaultTracer, options.Client)
	options.Client = wrapper.AuthClient(serviceName, options.Server.Options().Id, authFn, options.Client)

	// wrap the server to provide handler stats
	options.Server.Init(
		server.WrapHandler(wrapper.HandlerStats(stats.DefaultStats)),
		server.WrapHandler(wrapper.TraceHandler(trace.DefaultTracer)),
		server.WrapHandler(wrapper.AuthHandler(authFn)),
	)

	// set opts
	service.opts = options

	return service
}

func (s *service) Name() string {
	return s.opts.Server.Options().Name
}

// Init initialises options. Additionally it calls cmd.Init
// which parses command line flags. cmd.Init is only called
// on first Init.
func (s *service) Init(opts ...Option) {
	// process options
	for _, o := range opts {
		o(&s.opts)
	}

	s.once.Do(func() {
		// setup the plugins
		for _, p := range strings.Split(os.Getenv("MICRO_PLUGIN"), ",") {
			if len(p) == 0 {
				continue
			}

			// load the plugin
			c, err := plugin.Load(p)
			if err != nil {
				logger.Fatal(err)
			}

			// initialise the plugin
			if err := plugin.Init(c); err != nil {
				logger.Fatal(err)
			}
		}

		// set cmd name
		if len(s.opts.Cmd.App().Name) == 0 {
			s.opts.Cmd.App().Name = s.Server().Options().Name
		}

		// Initialise the command flags, overriding new service
		if err := s.opts.Cmd.Init(
			cmd.Auth(&s.opts.Auth),
			cmd.Broker(&s.opts.Broker),
			cmd.Registry(&s.opts.Registry),
			cmd.Runtime(&s.opts.Runtime),
			cmd.Transport(&s.opts.Transport),
			cmd.Client(&s.opts.Client),
			cmd.Config(&s.opts.Config),
			cmd.Server(&s.opts.Server),
			cmd.Store(&s.opts.Store),
			cmd.Profile(&s.opts.Profile),
		); err != nil {
			logger.Fatal(err)
		}

		// Explicitly set the table name to the service name
		name := s.opts.Cmd.App().Name
		s.opts.Store.Init(store.Table(name))

		// Set the client for the micro clients
		s.opts.Auth.Init(auth.WithClient(s.Client()))
		s.opts.Runtime.Init(runtime.WithClient(s.Client()))
		s.opts.Store.Init(store.WithClient(s.Client()))
	})
}

func (s *service) Options() Options {
	return s.opts
}

func (s *service) Client() client.Client {
	return s.opts.Client
}

func (s *service) Server() server.Server {
	return s.opts.Server
}

func (s *service) String() string {
	return "micro"
}

func (s *service) Start() error {
	for _, fn := range s.opts.BeforeStart {
		if err := fn(); err != nil {
			return err
		}
	}

	if err := s.opts.Server.Start(); err != nil {
		return err
	}

	for _, fn := range s.opts.AfterStart {
		if err := fn(); err != nil {
			return err
		}
	}

	return nil
}

func (s *service) Stop() error {
	var gerr error

	for _, fn := range s.opts.BeforeStop {
		if err := fn(); err != nil {
			gerr = err
		}
	}

	if err := s.opts.Server.Stop(); err != nil {
		return err
	}

	for _, fn := range s.opts.AfterStop {
		if err := fn(); err != nil {
			gerr = err
		}
	}

	return gerr
}

func (s *service) Run() error {
	// register the debug handler
	s.opts.Server.Handle(
		s.opts.Server.NewHandler(
			handler.NewHandler(),
			server.InternalHandler(true),
		),
	)

	// start the profiler
	if s.opts.Profile != nil {
		// to view mutex contention
		rtime.SetMutexProfileFraction(5)
		// to view blocking profile
		rtime.SetBlockProfileRate(1)

		if err := s.opts.Profile.Start(); err != nil {
			return err
		}
		defer s.opts.Profile.Stop()
	}

	if logger.V(logger.InfoLevel, logger.DefaultLogger) {
		logger.Infof("Starting [service] %s", s.Name())
	}

	if err := s.Start(); err != nil {
		return err
	}

	ch := make(chan os.Signal, 1)
	if s.opts.Signal {
		signal.Notify(ch, signalutil.Shutdown()...)
	}

	select {
	// wait on kill signal
	case <-ch:
	// wait on context cancel
	case <-s.opts.Context.Done():
	}

	return s.Stop()
}
<commit_msg>Disable auth service client<commit_after>package micro

import (
	"os"
	"os/signal"
	rtime "runtime"
	"strings"
	"sync"

	"github.com/micro/go-micro/v2/auth"
	"github.com/micro/go-micro/v2/client"
	"github.com/micro/go-micro/v2/config/cmd"
	"github.com/micro/go-micro/v2/debug/service/handler"
	"github.com/micro/go-micro/v2/debug/stats"
	"github.com/micro/go-micro/v2/debug/trace"
	"github.com/micro/go-micro/v2/logger"
	"github.com/micro/go-micro/v2/plugin"
	"github.com/micro/go-micro/v2/runtime"
	"github.com/micro/go-micro/v2/server"
	"github.com/micro/go-micro/v2/store"
	signalutil "github.com/micro/go-micro/v2/util/signal"
	"github.com/micro/go-micro/v2/util/wrapper"
)

type service struct {
	opts Options

	once sync.Once
}

func newService(opts ...Option) Service {
	service := new(service)
	options := newOptions(opts...)

	// service name
	serviceName := options.Server.Options().Name

	// authFn returns the auth, we pass as a function since auth
	// has not yet been set at this point.
	authFn := func() auth.Auth { return options.Server.Options().Auth }

	// wrap client to inject From-Service header on any calls
	options.Client = wrapper.FromService(serviceName, options.Client)
	options.Client = wrapper.TraceCall(serviceName, trace.DefaultTracer, options.Client)
	options.Client = wrapper.AuthClient(serviceName, options.Server.Options().Id, authFn, options.Client)

	// wrap the server to provide handler stats
	options.Server.Init(
		server.WrapHandler(wrapper.HandlerStats(stats.DefaultStats)),
		server.WrapHandler(wrapper.TraceHandler(trace.DefaultTracer)),
		server.WrapHandler(wrapper.AuthHandler(authFn)),
	)

	// set opts
	service.opts = options

	return service
}

func (s *service) Name() string {
	return s.opts.Server.Options().Name
}

// Init initialises options. Additionally it calls cmd.Init
// which parses command line flags. cmd.Init is only called
// on first Init.
func (s *service) Init(opts ...Option) {
	// process options
	for _, o := range opts {
		o(&s.opts)
	}

	s.once.Do(func() {
		// setup the plugins
		for _, p := range strings.Split(os.Getenv("MICRO_PLUGIN"), ",") {
			if len(p) == 0 {
				continue
			}

			// load the plugin
			c, err := plugin.Load(p)
			if err != nil {
				logger.Fatal(err)
			}

			// initialise the plugin
			if err := plugin.Init(c); err != nil {
				logger.Fatal(err)
			}
		}

		// set cmd name
		if len(s.opts.Cmd.App().Name) == 0 {
			s.opts.Cmd.App().Name = s.Server().Options().Name
		}

		// Initialise the command flags, overriding new service
		if err := s.opts.Cmd.Init(
			cmd.Auth(&s.opts.Auth),
			cmd.Broker(&s.opts.Broker),
			cmd.Registry(&s.opts.Registry),
			cmd.Runtime(&s.opts.Runtime),
			cmd.Transport(&s.opts.Transport),
			cmd.Client(&s.opts.Client),
			cmd.Config(&s.opts.Config),
			cmd.Server(&s.opts.Server),
			cmd.Store(&s.opts.Store),
			cmd.Profile(&s.opts.Profile),
		); err != nil {
			logger.Fatal(err)
		}

		// Explicitly set the table name to the service name
		name := s.opts.Cmd.App().Name
		s.opts.Store.Init(store.Table(name))

		// Set the client for the micro clients
		// s.opts.Auth.Init(auth.WithClient(s.Client()))
		s.opts.Runtime.Init(runtime.WithClient(s.Client()))
		s.opts.Store.Init(store.WithClient(s.Client()))
	})
}

func (s *service) Options() Options {
	return s.opts
}

func (s *service) Client() client.Client {
	return s.opts.Client
}

func (s *service) Server() server.Server {
	return s.opts.Server
}

func (s *service) String() string {
	return "micro"
}

func (s *service) Start() error {
	for _, fn := range s.opts.BeforeStart {
		if err := fn(); err != nil {
			return err
		}
	}

	if err := s.opts.Server.Start(); err != nil {
		return err
	}

	for _, fn := range s.opts.AfterStart {
		if err := fn(); err != nil {
			return err
		}
	}

	return nil
}

func (s *service) Stop() error {
	var gerr error

	for _, fn := range s.opts.BeforeStop {
		if err := fn(); err != nil {
			gerr = err
		}
	}

	if err := s.opts.Server.Stop(); err != nil {
		return err
	}

	for _, fn := range s.opts.AfterStop {
		if err := fn(); err != nil {
			gerr = err
		}
	}

	return gerr
}

func (s *service) Run() error {
	// register the debug handler
	s.opts.Server.Handle(
		s.opts.Server.NewHandler(
			handler.NewHandler(),
			server.InternalHandler(true),
		),
	)

	// start the profiler
	if s.opts.Profile != nil {
		// to view mutex contention
		rtime.SetMutexProfileFraction(5)
		// to view blocking profile
		rtime.SetBlockProfileRate(1)

		if err := s.opts.Profile.Start(); err != nil {
			return err
		}
		defer s.opts.Profile.Stop()
	}

	if logger.V(logger.InfoLevel, logger.DefaultLogger) {
		logger.Infof("Starting [service] %s", s.Name())
	}

	if err := s.Start(); err != nil {
		return err
	}

	ch := make(chan os.Signal, 1)
	if s.opts.Signal {
		signal.Notify(ch, signalutil.Shutdown()...)
	}

	select {
	// wait on kill signal
	case <-ch:
	// wait on context cancel
	case <-s.opts.Context.Done():
	}

	return s.Stop()
}
<|endoftext|>


<commit_before>package sand

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	log "github.com/sirupsen/logrus"
)

const (
	iso8601         = "2006-01-02T15:04:05.999999999Z07:00"
	serviceCacheKey = "service-access-token"
)

var notAllowedResponse = map[string]interface{}{
	"allowed": false,
}

//Service can be used to verify a token with SAND
type Service struct {
	Client

	//The default resource name that this Service will check the token against.
	Resource string

	//Default context
	Context map[string]interface{}

	//The URL of the token verification endpoint, e.g., "https://oauth.example.com/warden/token/allowed"
	TokenVerifyURL string

	//The default expiry time for cache for invalid tokens and also valid tokens without expiry times
	//Default value is 3599 seconds
	//Only services need this because client tokens will always give expiry time
	DefaultExpDuration time.Duration

	//The scopes required for the service to access the token verification endpoint
	Scopes []string
}

// VerificationOption affects how tokens are verified
type VerificationOption struct {
	TargetScopes    []string
	Resource        string
	Action          string
	Context         map[string]interface{}
	NumRetry        *int
	SkipCachedToken bool
	Timeout         time.Duration
}

//NewService returns a Service struct.
func NewService(id, secret, tokenURL, resource, verifyURL string, scopes []string) (service *Service, err error) {
	client, err := NewClient(id, secret, tokenURL)
	if err != nil || resource == "" || verifyURL == "" {
		err = errors.New("NewService: missing required argument(s)")
		return
	}
	client.cacheType = "service"
	service = &Service{
		Client:             *client,
		Resource:           resource,
		Context:            map[string]interface{}{},
		TokenVerifyURL:     verifyURL,
		Scopes:             scopes,
		DefaultExpDuration: defaultExpiryTime,
	}
	return
}

//CheckRequest checks the bearer token of an incoming HTTP request and return response with 'allowed' true/false field.
//If the error is of type sand.ConnectionError, the service should respond with
//HTTP status code 502. Otherwise the client would perform unnecessary retries.
//Example with Gin:
//  func(c *gin.Context) {
//    response, err := sandService.CheckRequest(c.Request, []string{"scope1", "scope2"}, "action")
//    if err != nil || response["allowed"] != true {
//      c.JSON(sandService.ErrorCode(err), err)    //This would set 502 on ConnectionError
//    }
//  }
func (s *Service) CheckRequest(r *http.Request, targetScopes []string, action string) (map[string]interface{}, error) {
	return s.CheckRequestWithCustomRetry(r, targetScopes, action, s.DefaultRetryCount)
}

//CheckRequestWithCustomRetry allows specifying a positive number as number of retries to
//use instead of using DefaultRetryCount on a per-request basis.
//Using a negative number for numRetry is equivalent to the "Request" function
func (s *Service) CheckRequestWithCustomRetry(r *http.Request, targetScopes []string, action string, numRetry int) (map[string]interface{}, error) {
	return s.VerifyRequest(r, &VerificationOption{TargetScopes: targetScopes, Action: action, NumRetry: &numRetry})
}

//VerifyRequest takes the token in a request and verifies with SAND
//Remember to set a reasonable NumRetry value (>= 0) for the VerificationOption
func (s *Service) VerifyRequest(r *http.Request, opt *VerificationOption) (map[string]interface{}, error) {
	token := ExtractToken(r.Header.Get("Authorization"))
	rv, err := s.VerifyTokenWithCache(token, opt)
	if err != nil {
		log.Error(err)
	}
	return rv, err
}

//ErrorCode gets the HTTP error code based on the error type. By default it is
//401 unauthorized; if the error is connection error, then it returns 502
func (s *Service) ErrorCode(err error) int {
	if err != nil {
		//Return 502 on error
		return http.StatusBadGateway
	}
	return http.StatusUnauthorized
}

//VerifyTokenWithCache tries to get the result for this token from the cache first.
//If not found in cache, if will make a token verification request with Sand.
func (s *Service) VerifyTokenWithCache(token string, opt *VerificationOption) (map[string]interface{}, error) {
	s.buildOption(opt)
	if token == "" || opt.Resource == "" {
		return notAllowedResponse, nil
	}

	var ckey string
	if theCache != nil {
		//Calculate cache key for use later
		ckey = s.cacheKey(token, opt.TargetScopes, opt.Resource)
		//Read from cache
		result := theCache.Read(ckey, 0)
		response, ok := result.(map[string]interface{})
		if ok {
			return response, nil
		}
	}
	resp, err := s.verifyToken(token, opt)
	if err != nil {
		switch err.(type) {
		case ExpiredTokenError:
			log.Infof("(%s) Service was unauthorized (401) to access Sand's token verification endpoint. Retry once by getting a new service token.", s.ClientID)
			option := *opt
			option.SkipCachedToken = true
			if resp, err = s.verifyToken(token, &option); err != nil {
				return notAllowedResponse, err
			}
		default:
			return notAllowedResponse, err
		}
	}
	if resp == nil {
		log.Warnf("(%s) Service getting nil token verification response.", s.ClientID)
		return notAllowedResponse, err
	}

	if theCache != nil {
		//Write to cache
		if resp["allowed"] == true {
			exp := s.DefaultExpDuration
			if resp["exp"] != nil {
				expTime, ok := resp["exp"].(string)
				if ok {
					exp = s.expirationTime(expTime)
				}
			} else {
				log.Warnf("(%s) Token verification result missing 'exp' field. Use default expiration time.", s.ClientID)
			}
			theCache.Write(ckey, resp, exp)
		} else {
			theCache.Write(ckey, notAllowedResponse, s.DefaultExpDuration)
		}
	}
	return resp, nil
}

//verifyToken verifies with SAND to see if the token is allowed to access this service.
//When returned error is nil and the map return value is nil or {"allowed":false},
//VerifyTokenWithCache will return notAllowedResponse, so the service will respond
//to the client with 401.
func (s *Service) verifyToken(token string, opt *VerificationOption) (map[string]interface{}, error) {
	if token == "" || opt.Resource == "" {
		return nil, nil
	}
	s.buildNumRetry(opt)
	accessToken, err := s.Token(serviceCacheKey, s.Scopes, &tokenOption{NumRetry: *opt.NumRetry, Timeout: opt.Timeout})
	if err != nil {
		return nil, err
	}

	transport := http.DefaultTransport.(*http.Transport).Clone()
	transport.TLSClientConfig.MinVersion = s.SSLMinVersion
	client := &http.Client{
		Transport: transport,
		Timeout:   getTimeout(opt.Timeout, s.Timeout),
	}

	data := map[string]interface{}{
		"scopes":   opt.TargetScopes,
		"token":    token,
		"resource": opt.Resource,
		"action":   opt.Action,
		"context":  opt.Context,
	}
	dBytes, err := json.Marshal(data)
	if err != nil {
		return nil, fmt.Errorf("(%s) Service failed to marshal json data: %v", s.ClientID, err)
	}

	var body []byte
	for retryCount := 0; retryCount <= *opt.NumRetry; retryCount++ {
		req, err := http.NewRequest("POST", s.TokenVerifyURL, bytes.NewBuffer(dBytes))
		if err != nil {
			return nil, fmt.Errorf("(%s) Service failed to create token verification request: %v", s.ClientID, err)
		}
		req.Header.Add("Authorization", "Bearer "+accessToken)
		resp, err := client.Do(req)
		if err != nil {
			if retryCount < *opt.NumRetry {
				sleep := ExponentialRetryDuration(retryCount)
				time.Sleep(sleep)
				continue
			}
			return nil, InternalServerError{SandError{
				fmt.Sprintf("(%s) Service failed to verify the token: %v", s.ClientID, err),
			}}
		}
		defer resp.Body.Close()
		body, _ = ioutil.ReadAll(resp.Body)
		if resp.StatusCode != http.StatusOK {
			str := fmt.Sprintf("(%s) Error response from the authentication service: %d - %s", s.ClientID, resp.StatusCode, body)
			log.Warn(str)
			switch resp.StatusCode {
			case http.StatusInternalServerError:
				//When the response is 500, the client's token may be expired. So let the client retry
				//and return 401 by returning nil, so that the result is not cached.
				return nil, nil
			case http.StatusUnauthorized:
				return nil, ExpiredTokenError{SandError{
					fmt.Sprintf("(%s) Service cannot access the token verification endpoint at this time with the current service token", s.ClientID),
				}}
			case http.StatusForbidden:
				return nil, PermissionDeniedError{SandError{
					fmt.Sprintf("(%s) Service is denied access to the token verification endpoint", s.ClientID),
				}}
			default:
				if retryCount < *opt.NumRetry {
					sleep := ExponentialRetryDuration(retryCount)
					time.Sleep(sleep)
					continue
				}
				if resp.StatusCode == http.StatusTooManyRequests {
					return nil, TooManyRequestsError{SandError{
						fmt.Sprintf("(%s) Service cannot verify a token due to rate limiting", s.ClientID),
					}}
				}
				return nil, InternalServerError{SandError{fmt.Sprintf("(%s) %s", s.ClientID, str)}}
			}
		}
	}
	var result map[string]interface{}
	err = json.Unmarshal(body, &result)
	return result, err
}

//Set the defaults for values that are not given.
func (s *Service) buildOption(opt *VerificationOption) {
	if opt.Resource == "" {
		opt.Resource = s.Resource
	}
	if len(opt.Context) == 0 {
		opt.Context = s.Context
	}
	if len(opt.TargetScopes) == 0 {
		opt.TargetScopes = []string{}
	}
}

func (s *Service) buildNumRetry(opt *VerificationOption) {
	retry := s.DefaultRetryCount
	if opt.NumRetry != nil {
		retry = *opt.NumRetry
	}
	retry = s.tokenRequestRetryCount(retry)
	opt.NumRetry = &retry
}

//expirationTime computes the expiry time given the expiry time as a string
//Example time returned by SAND: {"exp":"2021-08-26T05:33:27.107035237Z"}
func (s *Service) expirationTime(expTime string) time.Duration {
	if expTime == "" {
		log.Warnf("(%s) 'exp' field is empty from the token verification result. Use default expiration time.", s.ClientID)
		return s.DefaultExpDuration
	}
	t, err := time.Parse(iso8601, expTime)
	if err != nil {
		log.Warnf("(%s) Error parsing the token expiration time: %s. Use default expiration time: %v", s.ClientID, expTime, err)
		return s.DefaultExpDuration
	}
	return t.Sub(time.Now())
}
<commit_msg>Retry at least once when service getting a token. Previously if the service fails to get a token, it'd respond with 502 and client won't retry. But the service should retry getting a token at least once.<commit_after>package sand

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	log "github.com/sirupsen/logrus"
)

const (
	iso8601         = "2006-01-02T15:04:05.999999999Z07:00"
	serviceCacheKey = "service-access-token"
)

var notAllowedResponse = map[string]interface{}{
	"allowed": false,
}

//Service can be used to verify a token with SAND
type Service struct {
	Client

	//The default resource name that this Service will check the token against.
	Resource string

	//Default context
	Context map[string]interface{}

	//The URL of the token verification endpoint, e.g., "https://oauth.example.com/warden/token/allowed"
	TokenVerifyURL string

	//The default expiry time for cache for invalid tokens and also valid tokens without expiry times
	//Default value is 3599 seconds
	//Only services need this because client tokens will always give expiry time
	DefaultExpDuration time.Duration

	//The scopes required for the service to access the token verification endpoint
	Scopes []string
}

// VerificationOption affects how tokens are verified
type VerificationOption struct {
	TargetScopes    []string
	Resource        string
	Action          string
	Context         map[string]interface{}
	NumRetry        *int
	SkipCachedToken bool
	Timeout         time.Duration
}

//NewService returns a Service struct.
func NewService(id, secret, tokenURL, resource, verifyURL string, scopes []string) (service *Service, err error) {
	client, err := NewClient(id, secret, tokenURL)
	if err != nil || resource == "" || verifyURL == "" {
		err = errors.New("NewService: missing required argument(s)")
		return
	}
	client.cacheType = "service"
	service = &Service{
		Client:             *client,
		Resource:           resource,
		Context:            map[string]interface{}{},
		TokenVerifyURL:     verifyURL,
		Scopes:             scopes,
		DefaultExpDuration: defaultExpiryTime,
	}
	return
}

//CheckRequest checks the bearer token of an incoming HTTP request and return response with 'allowed' true/false field.
//If the error is of type sand.ConnectionError, the service should respond with
//HTTP status code 502. Otherwise the client would perform unnecessary retries.
//Example with Gin:
//  func(c *gin.Context) {
//    response, err := sandService.CheckRequest(c.Request, []string{"scope1", "scope2"}, "action")
//    if err != nil || response["allowed"] != true {
//      c.JSON(sandService.ErrorCode(err), err)    //This would set 502 on ConnectionError
//    }
//  }
func (s *Service) CheckRequest(r *http.Request, targetScopes []string, action string) (map[string]interface{}, error) {
	return s.CheckRequestWithCustomRetry(r, targetScopes, action, s.DefaultRetryCount)
}

//CheckRequestWithCustomRetry allows specifying a positive number as number of retries to
//use instead of using DefaultRetryCount on a per-request basis.
//Using a negative number for numRetry is equivalent to the "Request" function
func (s *Service) CheckRequestWithCustomRetry(r *http.Request, targetScopes []string, action string, numRetry int) (map[string]interface{}, error) {
	return s.VerifyRequest(r, &VerificationOption{TargetScopes: targetScopes, Action: action, NumRetry: &numRetry})
}

//VerifyRequest takes the token in a request and verifies with SAND
//Remember to set a reasonable NumRetry value (>= 0) for the VerificationOption
func (s *Service) VerifyRequest(r *http.Request, opt *VerificationOption) (map[string]interface{}, error) {
	token := ExtractToken(r.Header.Get("Authorization"))
	rv, err := s.VerifyTokenWithCache(token, opt)
	if err != nil {
		log.Error(err)
	}
	return rv, err
}

//ErrorCode gets the HTTP error code based on the error type. By default it is
//401 unauthorized; if the error is connection error, then it returns 502
func (s *Service) ErrorCode(err error) int {
	if err != nil {
		//Return 502 on error
		return http.StatusBadGateway
	}
	return http.StatusUnauthorized
}

//VerifyTokenWithCache tries to get the result for this token from the cache first.
//If not found in cache, if will make a token verification request with Sand.
func (s *Service) VerifyTokenWithCache(token string, opt *VerificationOption) (map[string]interface{}, error) {
	s.buildOption(opt)
	if token == "" || opt.Resource == "" {
		return notAllowedResponse, nil
	}

	var ckey string
	if theCache != nil {
		//Calculate cache key for use later
		ckey = s.cacheKey(token, opt.TargetScopes, opt.Resource)
		//Read from cache
		result := theCache.Read(ckey, 0)
		response, ok := result.(map[string]interface{})
		if ok {
			return response, nil
		}
	}
	resp, err := s.verifyToken(token, opt)
	if err != nil {
		switch err.(type) {
		case ExpiredTokenError:
			log.Infof("(%s) Service was unauthorized (401) to access Sand's token verification endpoint. Retry once by getting a new service token.", s.ClientID)
			option := *opt
			option.SkipCachedToken = true
			if resp, err = s.verifyToken(token, &option); err != nil {
				return notAllowedResponse, err
			}
		default:
			return notAllowedResponse, err
		}
	}
	if resp == nil {
		log.Warnf("(%s) Service getting nil token verification response.", s.ClientID)
		return notAllowedResponse, err
	}

	if theCache != nil {
		//Write to cache
		if resp["allowed"] == true {
			exp := s.DefaultExpDuration
			if resp["exp"] != nil {
				expTime, ok := resp["exp"].(string)
				if ok {
					exp = s.expirationTime(expTime)
				}
			} else {
				log.Warnf("(%s) Token verification result missing 'exp' field. Use default expiration time.", s.ClientID)
			}
			theCache.Write(ckey, resp, exp)
		} else {
			theCache.Write(ckey, notAllowedResponse, s.DefaultExpDuration)
		}
	}
	return resp, nil
}

//verifyToken verifies with SAND to see if the token is allowed to access this service.
//When returned error is nil and the map return value is nil or {"allowed":false},
//VerifyTokenWithCache will return notAllowedResponse, so the service will respond
//to the client with 401.
func (s *Service) verifyToken(token string, opt *VerificationOption) (map[string]interface{}, error) {
	if token == "" || opt.Resource == "" {
		return nil, nil
	}
	s.buildNumRetry(opt)
	numRetry := *opt.NumRetry
	if numRetry < 1 {
		numRetry = 1
	}
	accessToken, err := s.Token(serviceCacheKey, s.Scopes, &tokenOption{NumRetry: numRetry, Timeout: opt.Timeout})
	if err != nil {
		return nil, err
	}

	transport := http.DefaultTransport.(*http.Transport).Clone()
	transport.TLSClientConfig.MinVersion = s.SSLMinVersion
	client := &http.Client{
		Transport: transport,
		Timeout:   getTimeout(opt.Timeout, s.Timeout),
	}

	data := map[string]interface{}{
		"scopes":   opt.TargetScopes,
		"token":    token,
		"resource": opt.Resource,
		"action":   opt.Action,
		"context":  opt.Context,
	}
	dBytes, err := json.Marshal(data)
	if err != nil {
		return nil, fmt.Errorf("(%s) Service failed to marshal json data: %v", s.ClientID, err)
	}

	var body []byte
	for retryCount := 0; retryCount <= *opt.NumRetry; retryCount++ {
		req, err := http.NewRequest("POST", s.TokenVerifyURL, bytes.NewBuffer(dBytes))
		if err != nil {
			return nil, fmt.Errorf("(%s) Service failed to create token verification request: %v", s.ClientID, err)
		}
		req.Header.Add("Authorization", "Bearer "+accessToken)
		resp, err := client.Do(req)
		if err != nil {
			if retryCount < *opt.NumRetry {
				sleep := ExponentialRetryDuration(retryCount)
				time.Sleep(sleep)
				continue
			}
			return nil, InternalServerError{SandError{
				fmt.Sprintf("(%s) Service failed to verify the token: %v", s.ClientID, err),
			}}
		}
		defer resp.Body.Close()
		body, _ = ioutil.ReadAll(resp.Body)
		if resp.StatusCode != http.StatusOK {
			str := fmt.Sprintf("(%s) Error response from the authentication service: %d - %s", s.ClientID, resp.StatusCode, body)
			log.Warn(str)
			switch resp.StatusCode {
			case http.StatusInternalServerError:
				//When the response is 500, the client's token may be expired. So let the client retry
				//and return 401 by returning nil, so that the result is not cached.
				return nil, nil
			case http.StatusUnauthorized:
				return nil, ExpiredTokenError{SandError{
					fmt.Sprintf("(%s) Service cannot access the token verification endpoint at this time with the current service token", s.ClientID),
				}}
			case http.StatusForbidden:
				return nil, PermissionDeniedError{SandError{
					fmt.Sprintf("(%s) Service is denied access to the token verification endpoint", s.ClientID),
				}}
			default:
				if retryCount < *opt.NumRetry {
					sleep := ExponentialRetryDuration(retryCount)
					time.Sleep(sleep)
					continue
				}
				if resp.StatusCode == http.StatusTooManyRequests {
					return nil, TooManyRequestsError{SandError{
						fmt.Sprintf("(%s) Service cannot verify a token due to rate limiting", s.ClientID),
					}}
				}
				return nil, InternalServerError{SandError{fmt.Sprintf("(%s) %s", s.ClientID, str)}}
			}
		}
	}
	var result map[string]interface{}
	err = json.Unmarshal(body, &result)
	return result, err
}

//Set the defaults for values that are not given.
func (s *Service) buildOption(opt *VerificationOption) {
	if opt.Resource == "" {
		opt.Resource = s.Resource
	}
	if len(opt.Context) == 0 {
		opt.Context = s.Context
	}
	if len(opt.TargetScopes) == 0 {
		opt.TargetScopes = []string{}
	}
}

func (s *Service) buildNumRetry(opt *VerificationOption) {
	retry := s.DefaultRetryCount
	if opt.NumRetry != nil {
		retry = *opt.NumRetry
	}
	retry = s.tokenRequestRetryCount(retry)
	opt.NumRetry = &retry
}

//expirationTime computes the expiry time given the expiry time as a string
//Example time returned by SAND: {"exp":"2021-08-26T05:33:27.107035237Z"}
func (s *Service) expirationTime(expTime string) time.Duration {
	if expTime == "" {
		log.Warnf("(%s) 'exp' field is empty from the token verification result. Use default expiration time.", s.ClientID)
		return s.DefaultExpDuration
	}
	t, err := time.Parse(iso8601, expTime)
	if err != nil {
		log.Warnf("(%s) Error parsing the token expiration time: %s. Use default expiration time: %v", s.ClientID, expTime, err)
		return s.DefaultExpDuration
	}
	return t.Sub(time.Now())
}
<|endoftext|>


<commit_before>package metrics

import (
	"fmt"
	"testing"
	"time"

	"github.com/NVIDIA/gpu-monitoring-tools/bindings/go/nvml"
	"github.com/prometheus/client_golang/prometheus/testutil"
)

func uint64Ptr(u uint64) *uint64 {
	return &u
}

func stringPtr(u string) *string {
	return &u
}

type MockDevice struct {
	DeviceInfo *nvml.Device
	Status     *nvml.DeviceStatus
}

type mockDeviceWrapper struct {
	device MockDevice
}

func (d *mockDeviceWrapper) giveDevice() *nvml.Device {
	return d.device.DeviceInfo
}

func (d *mockDeviceWrapper) giveStatus() (status *nvml.DeviceStatus, err error) {
	status = d.device.Status
	return status, nil
}

type MockGather struct{}

func (t *MockGather) gatherDevice(deviceName string) (deviceWrapper, error) {
	device, ok := gpuDevicesMock[deviceName]
	if !ok {
		return &mockDeviceWrapper{}, fmt.Errorf("device %s not found", deviceName)
	}

	return &mockDeviceWrapper{*device}, nil
}

func (t *MockGather) gatherStatus(d deviceWrapper) (status *nvml.DeviceStatus, err error) {
	return d.giveStatus()
}

func (t *MockGather) gatherDutyCycle(uuid string, since time.Duration) (uint, error) {
	dutyCycle, ok := dutyCycleMock[uuid]
	if !ok {
		return 0, fmt.Errorf("duty cycle for %s not found", uuid)
	}
	return dutyCycle, nil
}

var (
	containerDevicesMock = map[ContainerID][]string{
		{
			namespace: "default",
			pod:       "pod1",
			container: "container1",
		}: {
			"q759757",
		},
		{
			namespace: "non-default",
			pod:       "pod2",
			container: "container2",
		}: {
			"afjodaj",
			"7v89zhi",
		},
	}

	gpuDevicesMock = map[string]*MockDevice{
		"q759757": {
			DeviceInfo: &nvml.Device{
				UUID:   "656547758",
				Model:  stringPtr("model1"),
				Memory: uint64Ptr(200),
			},
			Status: &nvml.DeviceStatus{
				Memory: nvml.MemoryInfo{
					Global: nvml.DeviceMemory{
						Used: uint64Ptr(50),
					},
				},
			},
		},
		"afjodaj": {
			DeviceInfo: &nvml.Device{
				UUID:   "850729563",
				Model:  stringPtr("model2"),
				Memory: uint64Ptr(200),
			},
			Status: &nvml.DeviceStatus{
				Memory: nvml.MemoryInfo{
					Global: nvml.DeviceMemory{
						Used: uint64Ptr(150),
					},
				},
			},
		},
		"7v89zhi": {
			DeviceInfo: &nvml.Device{
				UUID:   "3572375710",
				Model:  stringPtr("model1"),
				Memory: uint64Ptr(350),
			},
			Status: &nvml.DeviceStatus{
				Memory: nvml.MemoryInfo{
					Global: nvml.DeviceMemory{
						Used: uint64Ptr(100),
					},
				},
			},
		},
	}

	dutyCycleMock = map[string]uint{
		"656547758":  78,
		"850729563":  32,
		"3572375710": 13,
	}
)

func TestMetricsUpdate(t *testing.T) {
	g = &MockGather{}
	ms := MetricServer{}
	ms.updateMetrics(containerDevicesMock)

	if testutil.ToFloat64(
		AcceleratorRequests.WithLabelValues(
			"default", "pod1", "container1", gpuResourceName)) != 1 ||
		testutil.ToFloat64(
			AcceleratorRequests.WithLabelValues(
				"non-default", "pod2", "container2", gpuResourceName)) != 2 {
		t.Fatalf("Wrong Result in AcceleratorRequsets")
	}

	if testutil.ToFloat64(
		DutyCycle.WithLabelValues(
			"default", "pod1", "container1", "nvidia", "656547758", "model1")) != 78 ||
		testutil.ToFloat64(
			DutyCycle.WithLabelValues(
				"non-default", "pod2", "container2", "nvidia", "850729563", "model2")) != 32 ||
		testutil.ToFloat64(
			DutyCycle.WithLabelValues(
				"non-default", "pod2", "container2", "nvidia", "3572375710", "model1")) != 13 {
		t.Fatalf("Wrong Result in DutyCycle")
	}

	if testutil.ToFloat64(
		MemoryTotal.WithLabelValues(
			"default", "pod1", "container1", "nvidia", "656547758", "model1")) != 200*1024*1024 ||
		testutil.ToFloat64(
			MemoryTotal.WithLabelValues(
				"non-default", "pod2", "container2", "nvidia", "850729563", "model2")) != 200*1024*1024 ||
		testutil.ToFloat64(
			MemoryTotal.WithLabelValues(
				"non-default", "pod2", "container2", "nvidia", "3572375710", "model1")) != 350*1024*1024 {
		t.Fatalf("Wrong Result in MemoryTotal")
	}

	if testutil.ToFloat64(
		MemoryUsed.WithLabelValues(
			"default", "pod1", "container1", "nvidia", "656547758", "model1")) != 50*1024*1024 ||
		testutil.ToFloat64(
			MemoryUsed.WithLabelValues(
				"non-default", "pod2", "container2", "nvidia", "850729563", "model2")) != 150*1024*1024 ||
		testutil.ToFloat64(
			MemoryUsed.WithLabelValues(
				"non-default", "pod2", "container2", "nvidia", "3572375710", "model1")) != 100*1024*1024 {
		t.Fatalf("Wrong Result in MemoryTotal")
	}

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

package metrics

import (
	"fmt"
	"testing"
	"time"

	"github.com/NVIDIA/gpu-monitoring-tools/bindings/go/nvml"
	"github.com/prometheus/client_golang/prometheus/testutil"
)

func uint64Ptr(u uint64) *uint64 {
	return &u
}

func stringPtr(u string) *string {
	return &u
}

type MockDevice struct {
	DeviceInfo *nvml.Device
	Status     *nvml.DeviceStatus
}

type mockDeviceWrapper struct {
	device MockDevice
}

func (d *mockDeviceWrapper) giveDevice() *nvml.Device {
	return d.device.DeviceInfo
}

func (d *mockDeviceWrapper) giveStatus() (status *nvml.DeviceStatus, err error) {
	status = d.device.Status
	return status, nil
}

type MockGather struct{}

func (t *MockGather) gatherDevice(deviceName string) (deviceWrapper, error) {
	device, ok := gpuDevicesMock[deviceName]
	if !ok {
		return &mockDeviceWrapper{}, fmt.Errorf("device %s not found", deviceName)
	}

	return &mockDeviceWrapper{*device}, nil
}

func (t *MockGather) gatherStatus(d deviceWrapper) (status *nvml.DeviceStatus, err error) {
	return d.giveStatus()
}

func (t *MockGather) gatherDutyCycle(uuid string, since time.Duration) (uint, error) {
	dutyCycle, ok := dutyCycleMock[uuid]
	if !ok {
		return 0, fmt.Errorf("duty cycle for %s not found", uuid)
	}
	return dutyCycle, nil
}

var (
	containerDevicesMock = map[ContainerID][]string{
		{
			namespace: "default",
			pod:       "pod1",
			container: "container1",
		}: {
			"q759757",
		},
		{
			namespace: "non-default",
			pod:       "pod2",
			container: "container2",
		}: {
			"afjodaj",
			"7v89zhi",
		},
	}

	gpuDevicesMock = map[string]*MockDevice{
		"q759757": {
			DeviceInfo: &nvml.Device{
				UUID:   "656547758",
				Model:  stringPtr("model1"),
				Memory: uint64Ptr(200),
			},
			Status: &nvml.DeviceStatus{
				Memory: nvml.MemoryInfo{
					Global: nvml.DeviceMemory{
						Used: uint64Ptr(50),
					},
				},
			},
		},
		"afjodaj": {
			DeviceInfo: &nvml.Device{
				UUID:   "850729563",
				Model:  stringPtr("model2"),
				Memory: uint64Ptr(200),
			},
			Status: &nvml.DeviceStatus{
				Memory: nvml.MemoryInfo{
					Global: nvml.DeviceMemory{
						Used: uint64Ptr(150),
					},
				},
			},
		},
		"7v89zhi": {
			DeviceInfo: &nvml.Device{
				UUID:   "3572375710",
				Model:  stringPtr("model1"),
				Memory: uint64Ptr(350),
			},
			Status: &nvml.DeviceStatus{
				Memory: nvml.MemoryInfo{
					Global: nvml.DeviceMemory{
						Used: uint64Ptr(100),
					},
				},
			},
		},
	}

	dutyCycleMock = map[string]uint{
		"656547758":  78,
		"850729563":  32,
		"3572375710": 13,
	}
)

func TestMetricsUpdate(t *testing.T) {
	g = &MockGather{}
	ms := MetricServer{}
	ms.updateMetrics(containerDevicesMock)

	if testutil.ToFloat64(
		AcceleratorRequests.WithLabelValues(
			"default", "pod1", "container1", gpuResourceName)) != 1 ||
		testutil.ToFloat64(
			AcceleratorRequests.WithLabelValues(
				"non-default", "pod2", "container2", gpuResourceName)) != 2 {
		t.Fatalf("Wrong Result in AcceleratorRequsets")
	}

	if testutil.ToFloat64(
		DutyCycle.WithLabelValues(
			"default", "pod1", "container1", "nvidia", "656547758", "model1")) != 78 ||
		testutil.ToFloat64(
			DutyCycle.WithLabelValues(
				"non-default", "pod2", "container2", "nvidia", "850729563", "model2")) != 32 ||
		testutil.ToFloat64(
			DutyCycle.WithLabelValues(
				"non-default", "pod2", "container2", "nvidia", "3572375710", "model1")) != 13 {
		t.Fatalf("Wrong Result in DutyCycle")
	}

	if testutil.ToFloat64(
		MemoryTotal.WithLabelValues(
			"default", "pod1", "container1", "nvidia", "656547758", "model1")) != 200*1024*1024 ||
		testutil.ToFloat64(
			MemoryTotal.WithLabelValues(
				"non-default", "pod2", "container2", "nvidia", "850729563", "model2")) != 200*1024*1024 ||
		testutil.ToFloat64(
			MemoryTotal.WithLabelValues(
				"non-default", "pod2", "container2", "nvidia", "3572375710", "model1")) != 350*1024*1024 {
		t.Fatalf("Wrong Result in MemoryTotal")
	}

	if testutil.ToFloat64(
		MemoryUsed.WithLabelValues(
			"default", "pod1", "container1", "nvidia", "656547758", "model1")) != 50*1024*1024 ||
		testutil.ToFloat64(
			MemoryUsed.WithLabelValues(
				"non-default", "pod2", "container2", "nvidia", "850729563", "model2")) != 150*1024*1024 ||
		testutil.ToFloat64(
			MemoryUsed.WithLabelValues(
				"non-default", "pod2", "container2", "nvidia", "3572375710", "model1")) != 100*1024*1024 {
		t.Fatalf("Wrong Result in MemoryTotal")
	}

}
<|endoftext|>


<commit_before>package opts

import (
	"fmt"
	"math/big"
	"net"
	"path"
	"regexp"
	"strings"

	"github.com/docker/docker/api/types/filters"
	units "github.com/docker/go-units"
)

var (
	alphaRegexp  = regexp.MustCompile(`[a-zA-Z]`)
	domainRegexp = regexp.MustCompile(`^(:?(:?[a-zA-Z0-9]|(:?[a-zA-Z0-9][a-zA-Z0-9\-]*[a-zA-Z0-9]))(:?\.(:?[a-zA-Z0-9]|(:?[a-zA-Z0-9][a-zA-Z0-9\-]*[a-zA-Z0-9])))*)\.?\s*$`)
)

// ListOpts holds a list of values and a validation function.
type ListOpts struct {
	values    *[]string
	validator ValidatorFctType
}

// NewListOpts creates a new ListOpts with the specified validator.
func NewListOpts(validator ValidatorFctType) ListOpts {
	var values []string
	return *NewListOptsRef(&values, validator)
}

// NewListOptsRef creates a new ListOpts with the specified values and validator.
func NewListOptsRef(values *[]string, validator ValidatorFctType) *ListOpts {
	return &ListOpts{
		values:    values,
		validator: validator,
	}
}

func (opts *ListOpts) String() string {
	return fmt.Sprintf("%v", []string((*opts.values)))
}

// Set validates if needed the input value and adds it to the
// internal slice.
func (opts *ListOpts) Set(value string) error {
	if opts.validator != nil {
		v, err := opts.validator(value)
		if err != nil {
			return err
		}
		value = v
	}
	(*opts.values) = append((*opts.values), value)
	return nil
}

// Delete removes the specified element from the slice.
func (opts *ListOpts) Delete(key string) {
	for i, k := range *opts.values {
		if k == key {
			(*opts.values) = append((*opts.values)[:i], (*opts.values)[i+1:]...)
			return
		}
	}
}

// GetMap returns the content of values in a map in order to avoid
// duplicates.
func (opts *ListOpts) GetMap() map[string]struct{} {
	ret := make(map[string]struct{})
	for _, k := range *opts.values {
		ret[k] = struct{}{}
	}
	return ret
}

// GetAll returns the values of slice.
func (opts *ListOpts) GetAll() []string {
	return (*opts.values)
}

// GetAllOrEmpty returns the values of the slice
// or an empty slice when there are no values.
func (opts *ListOpts) GetAllOrEmpty() []string {
	v := *opts.values
	if v == nil {
		return make([]string, 0)
	}
	return v
}

// Get checks the existence of the specified key.
func (opts *ListOpts) Get(key string) bool {
	for _, k := range *opts.values {
		if k == key {
			return true
		}
	}
	return false
}

// Len returns the amount of element in the slice.
func (opts *ListOpts) Len() int {
	return len((*opts.values))
}

// Type returns a string name for this Option type
func (opts *ListOpts) Type() string {
	return "list"
}

// WithValidator returns the ListOpts with validator set.
func (opts *ListOpts) WithValidator(validator ValidatorFctType) *ListOpts {
	opts.validator = validator
	return opts
}

// NamedOption is an interface that list and map options
// with names implement.
type NamedOption interface {
	Name() string
}

// NamedListOpts is a ListOpts with a configuration name.
// This struct is useful to keep reference to the assigned
// field name in the internal configuration struct.
type NamedListOpts struct {
	name string
	ListOpts
}

var _ NamedOption = &NamedListOpts{}

// NewNamedListOptsRef creates a reference to a new NamedListOpts struct.
func NewNamedListOptsRef(name string, values *[]string, validator ValidatorFctType) *NamedListOpts {
	return &NamedListOpts{
		name:     name,
		ListOpts: *NewListOptsRef(values, validator),
	}
}

// Name returns the name of the NamedListOpts in the configuration.
func (o *NamedListOpts) Name() string {
	return o.name
}

// MapOpts holds a map of values and a validation function.
type MapOpts struct {
	values    map[string]string
	validator ValidatorFctType
}

// Set validates if needed the input value and add it to the
// internal map, by splitting on '='.
func (opts *MapOpts) Set(value string) error {
	if opts.validator != nil {
		v, err := opts.validator(value)
		if err != nil {
			return err
		}
		value = v
	}
	vals := strings.SplitN(value, "=", 2)
	if len(vals) == 1 {
		(opts.values)[vals[0]] = ""
	} else {
		(opts.values)[vals[0]] = vals[1]
	}
	return nil
}

// GetAll returns the values of MapOpts as a map.
func (opts *MapOpts) GetAll() map[string]string {
	return opts.values
}

func (opts *MapOpts) String() string {
	return fmt.Sprintf("%v", map[string]string((opts.values)))
}

// Type returns a string name for this Option type
func (opts *MapOpts) Type() string {
	return "map"
}

// NewMapOpts creates a new MapOpts with the specified map of values and a validator.
func NewMapOpts(values map[string]string, validator ValidatorFctType) *MapOpts {
	if values == nil {
		values = make(map[string]string)
	}
	return &MapOpts{
		values:    values,
		validator: validator,
	}
}

// NamedMapOpts is a MapOpts struct with a configuration name.
// This struct is useful to keep reference to the assigned
// field name in the internal configuration struct.
type NamedMapOpts struct {
	name string
	MapOpts
}

var _ NamedOption = &NamedMapOpts{}

// NewNamedMapOpts creates a reference to a new NamedMapOpts struct.
func NewNamedMapOpts(name string, values map[string]string, validator ValidatorFctType) *NamedMapOpts {
	return &NamedMapOpts{
		name:    name,
		MapOpts: *NewMapOpts(values, validator),
	}
}

// Name returns the name of the NamedMapOpts in the configuration.
func (o *NamedMapOpts) Name() string {
	return o.name
}

// ValidatorFctType defines a validator function that returns a validated string and/or an error.
type ValidatorFctType func(val string) (string, error)

// ValidatorFctListType defines a validator function that returns a validated list of string and/or an error
type ValidatorFctListType func(val string) ([]string, error)

// ValidateIPAddress validates an Ip address.
func ValidateIPAddress(val string) (string, error) {
	var ip = net.ParseIP(strings.TrimSpace(val))
	if ip != nil {
		return ip.String(), nil
	}
	return "", fmt.Errorf("%s is not an ip address", val)
}

// ValidateMACAddress validates a MAC address.
func ValidateMACAddress(val string) (string, error) {
	_, err := net.ParseMAC(strings.TrimSpace(val))
	if err != nil {
		return "", err
	}
	return val, nil
}

// ValidateDNSSearch validates domain for resolvconf search configuration.
// A zero length domain is represented by a dot (.).
func ValidateDNSSearch(val string) (string, error) {
	if val = strings.Trim(val, " "); val == "." {
		return val, nil
	}
	return validateDomain(val)
}

func validateDomain(val string) (string, error) {
	if alphaRegexp.FindString(val) == "" {
		return "", fmt.Errorf("%s is not a valid domain", val)
	}
	ns := domainRegexp.FindSubmatch([]byte(val))
	if len(ns) > 0 && len(ns[1]) < 255 {
		return string(ns[1]), nil
	}
	return "", fmt.Errorf("%s is not a valid domain", val)
}

// ValidateLabel validates that the specified string is a valid label, and returns it.
// Labels are in the form on key=value.
func ValidateLabel(val string) (string, error) {
	if strings.Count(val, "=") < 1 {
		return "", fmt.Errorf("bad attribute format: %s", val)
	}
	return val, nil
}

// ValidateSysctl validates a sysctl and returns it.
func ValidateSysctl(val string) (string, error) {
	validSysctlMap := map[string]bool{
		"kernel.msgmax":          true,
		"kernel.msgmnb":          true,
		"kernel.msgmni":          true,
		"kernel.sem":             true,
		"kernel.shmall":          true,
		"kernel.shmmax":          true,
		"kernel.shmmni":          true,
		"kernel.shm_rmid_forced": true,
	}
	validSysctlPrefixes := []string{
		"net.",
		"fs.mqueue.",
	}
	arr := strings.Split(val, "=")
	if len(arr) < 2 {
		return "", fmt.Errorf("sysctl '%s' is not whitelisted", val)
	}
	if validSysctlMap[arr[0]] {
		return val, nil
	}

	for _, vp := range validSysctlPrefixes {
		if strings.HasPrefix(arr[0], vp) {
			return val, nil
		}
	}
	return "", fmt.Errorf("sysctl '%s' is not whitelisted", val)
}

// FilterOpt is a flag type for validating filters
type FilterOpt struct {
	filter filters.Args
}

// NewFilterOpt returns a new FilterOpt
func NewFilterOpt() FilterOpt {
	return FilterOpt{filter: filters.NewArgs()}
}

func (o *FilterOpt) String() string {
	repr, err := filters.ToParam(o.filter)
	if err != nil {
		return "invalid filters"
	}
	return repr
}

// Set sets the value of the opt by parsing the command line value
func (o *FilterOpt) Set(value string) error {
	var err error
	o.filter, err = filters.ParseFlag(value, o.filter)
	return err
}

// Type returns the option type
func (o *FilterOpt) Type() string {
	return "filter"
}

// Value returns the value of this option
func (o *FilterOpt) Value() filters.Args {
	return o.filter
}

// NanoCPUs is a type for fixed point fractional number.
type NanoCPUs int64

// String returns the string format of the number
func (c *NanoCPUs) String() string {
	return big.NewRat(c.Value(), 1e9).FloatString(3)
}

// Set sets the value of the NanoCPU by passing a string
func (c *NanoCPUs) Set(value string) error {
	cpus, err := ParseCPUs(value)
	*c = NanoCPUs(cpus)
	return err
}

// Type returns the type
func (c *NanoCPUs) Type() string {
	return "decimal"
}

// Value returns the value in int64
func (c *NanoCPUs) Value() int64 {
	return int64(*c)
}

// ParseCPUs takes a string ratio and returns an integer value of nano cpus
func ParseCPUs(value string) (int64, error) {
	cpu, ok := new(big.Rat).SetString(value)
	if !ok {
		return 0, fmt.Errorf("failed to parse %v as a rational number", value)
	}
	nano := cpu.Mul(cpu, big.NewRat(1e9, 1))
	if !nano.IsInt() {
		return 0, fmt.Errorf("value is too precise")
	}
	return nano.Num().Int64(), nil
}

// ParseLink parses and validates the specified string as a link format (name:alias)
func ParseLink(val string) (string, string, error) {
	if val == "" {
		return "", "", fmt.Errorf("empty string specified for links")
	}
	arr := strings.Split(val, ":")
	if len(arr) > 2 {
		return "", "", fmt.Errorf("bad format for links: %s", val)
	}
	if len(arr) == 1 {
		return val, val, nil
	}
	// This is kept because we can actually get a HostConfig with links
	// from an already created container and the format is not `foo:bar`
	// but `/foo:/c1/bar`
	if strings.HasPrefix(arr[0], "/") {
		_, alias := path.Split(arr[1])
		return arr[0][1:], alias, nil
	}
	return arr[0], arr[1], nil
}

// ValidateLink validates that the specified string has a valid link format (containerName:alias).
func ValidateLink(val string) (string, error) {
	_, _, err := ParseLink(val)
	return val, err
}

// MemBytes is a type for human readable memory bytes (like 128M, 2g, etc)
type MemBytes int64

// String returns the string format of the human readable memory bytes
func (m *MemBytes) String() string {
	// NOTE: In spf13/pflag/flag.go, "0" is considered as "zero value" while "0 B" is not.
	// We return "0" in case value is 0 here so that the default value is hidden.
	// (Sometimes "default 0 B" is actually misleading)
	if m.Value() != 0 {
		return units.BytesSize(float64(m.Value()))
	}
	return "0"
}

// Set sets the value of the MemBytes by passing a string
func (m *MemBytes) Set(value string) error {
	val, err := units.RAMInBytes(value)
	*m = MemBytes(val)
	return err
}

// Type returns the type
func (m *MemBytes) Type() string {
	return "bytes"
}

// Value returns the value in int64
func (m *MemBytes) Value() int64 {
	return int64(*m)
}

// UnmarshalJSON is the customized unmarshaler for MemBytes
func (m *MemBytes) UnmarshalJSON(s []byte) error {
	if len(s) <= 2 || s[0] != '"' || s[len(s)-1] != '"' {
		return fmt.Errorf("invalid size: %q", s)
	}
	val, err := units.RAMInBytes(string(s[1 : len(s)-1]))
	*m = MemBytes(val)
	return err
}
<commit_msg>Use opts.MemBytes for flags.<commit_after>package opts

import (
	"fmt"
	"math/big"
	"net"
	"path"
	"regexp"
	"strings"

	"github.com/docker/docker/api/types/filters"
	units "github.com/docker/go-units"
)

var (
	alphaRegexp  = regexp.MustCompile(`[a-zA-Z]`)
	domainRegexp = regexp.MustCompile(`^(:?(:?[a-zA-Z0-9]|(:?[a-zA-Z0-9][a-zA-Z0-9\-]*[a-zA-Z0-9]))(:?\.(:?[a-zA-Z0-9]|(:?[a-zA-Z0-9][a-zA-Z0-9\-]*[a-zA-Z0-9])))*)\.?\s*$`)
)

// ListOpts holds a list of values and a validation function.
type ListOpts struct {
	values    *[]string
	validator ValidatorFctType
}

// NewListOpts creates a new ListOpts with the specified validator.
func NewListOpts(validator ValidatorFctType) ListOpts {
	var values []string
	return *NewListOptsRef(&values, validator)
}

// NewListOptsRef creates a new ListOpts with the specified values and validator.
func NewListOptsRef(values *[]string, validator ValidatorFctType) *ListOpts {
	return &ListOpts{
		values:    values,
		validator: validator,
	}
}

func (opts *ListOpts) String() string {
	return fmt.Sprintf("%v", []string((*opts.values)))
}

// Set validates if needed the input value and adds it to the
// internal slice.
func (opts *ListOpts) Set(value string) error {
	if opts.validator != nil {
		v, err := opts.validator(value)
		if err != nil {
			return err
		}
		value = v
	}
	(*opts.values) = append((*opts.values), value)
	return nil
}

// Delete removes the specified element from the slice.
func (opts *ListOpts) Delete(key string) {
	for i, k := range *opts.values {
		if k == key {
			(*opts.values) = append((*opts.values)[:i], (*opts.values)[i+1:]...)
			return
		}
	}
}

// GetMap returns the content of values in a map in order to avoid
// duplicates.
func (opts *ListOpts) GetMap() map[string]struct{} {
	ret := make(map[string]struct{})
	for _, k := range *opts.values {
		ret[k] = struct{}{}
	}
	return ret
}

// GetAll returns the values of slice.
func (opts *ListOpts) GetAll() []string {
	return (*opts.values)
}

// GetAllOrEmpty returns the values of the slice
// or an empty slice when there are no values.
func (opts *ListOpts) GetAllOrEmpty() []string {
	v := *opts.values
	if v == nil {
		return make([]string, 0)
	}
	return v
}

// Get checks the existence of the specified key.
func (opts *ListOpts) Get(key string) bool {
	for _, k := range *opts.values {
		if k == key {
			return true
		}
	}
	return false
}

// Len returns the amount of element in the slice.
func (opts *ListOpts) Len() int {
	return len((*opts.values))
}

// Type returns a string name for this Option type
func (opts *ListOpts) Type() string {
	return "list"
}

// WithValidator returns the ListOpts with validator set.
func (opts *ListOpts) WithValidator(validator ValidatorFctType) *ListOpts {
	opts.validator = validator
	return opts
}

// NamedOption is an interface that list and map options
// with names implement.
type NamedOption interface {
	Name() string
}

// NamedListOpts is a ListOpts with a configuration name.
// This struct is useful to keep reference to the assigned
// field name in the internal configuration struct.
type NamedListOpts struct {
	name string
	ListOpts
}

var _ NamedOption = &NamedListOpts{}

// NewNamedListOptsRef creates a reference to a new NamedListOpts struct.
func NewNamedListOptsRef(name string, values *[]string, validator ValidatorFctType) *NamedListOpts {
	return &NamedListOpts{
		name:     name,
		ListOpts: *NewListOptsRef(values, validator),
	}
}

// Name returns the name of the NamedListOpts in the configuration.
func (o *NamedListOpts) Name() string {
	return o.name
}

// MapOpts holds a map of values and a validation function.
type MapOpts struct {
	values    map[string]string
	validator ValidatorFctType
}

// Set validates if needed the input value and add it to the
// internal map, by splitting on '='.
func (opts *MapOpts) Set(value string) error {
	if opts.validator != nil {
		v, err := opts.validator(value)
		if err != nil {
			return err
		}
		value = v
	}
	vals := strings.SplitN(value, "=", 2)
	if len(vals) == 1 {
		(opts.values)[vals[0]] = ""
	} else {
		(opts.values)[vals[0]] = vals[1]
	}
	return nil
}

// GetAll returns the values of MapOpts as a map.
func (opts *MapOpts) GetAll() map[string]string {
	return opts.values
}

func (opts *MapOpts) String() string {
	return fmt.Sprintf("%v", map[string]string((opts.values)))
}

// Type returns a string name for this Option type
func (opts *MapOpts) Type() string {
	return "map"
}

// NewMapOpts creates a new MapOpts with the specified map of values and a validator.
func NewMapOpts(values map[string]string, validator ValidatorFctType) *MapOpts {
	if values == nil {
		values = make(map[string]string)
	}
	return &MapOpts{
		values:    values,
		validator: validator,
	}
}

// NamedMapOpts is a MapOpts struct with a configuration name.
// This struct is useful to keep reference to the assigned
// field name in the internal configuration struct.
type NamedMapOpts struct {
	name string
	MapOpts
}

var _ NamedOption = &NamedMapOpts{}

// NewNamedMapOpts creates a reference to a new NamedMapOpts struct.
func NewNamedMapOpts(name string, values map[string]string, validator ValidatorFctType) *NamedMapOpts {
	return &NamedMapOpts{
		name:    name,
		MapOpts: *NewMapOpts(values, validator),
	}
}

// Name returns the name of the NamedMapOpts in the configuration.
func (o *NamedMapOpts) Name() string {
	return o.name
}

// ValidatorFctType defines a validator function that returns a validated string and/or an error.
type ValidatorFctType func(val string) (string, error)

// ValidatorFctListType defines a validator function that returns a validated list of string and/or an error
type ValidatorFctListType func(val string) ([]string, error)

// ValidateIPAddress validates an Ip address.
func ValidateIPAddress(val string) (string, error) {
	var ip = net.ParseIP(strings.TrimSpace(val))
	if ip != nil {
		return ip.String(), nil
	}
	return "", fmt.Errorf("%s is not an ip address", val)
}

// ValidateMACAddress validates a MAC address.
func ValidateMACAddress(val string) (string, error) {
	_, err := net.ParseMAC(strings.TrimSpace(val))
	if err != nil {
		return "", err
	}
	return val, nil
}

// ValidateDNSSearch validates domain for resolvconf search configuration.
// A zero length domain is represented by a dot (.).
func ValidateDNSSearch(val string) (string, error) {
	if val = strings.Trim(val, " "); val == "." {
		return val, nil
	}
	return validateDomain(val)
}

func validateDomain(val string) (string, error) {
	if alphaRegexp.FindString(val) == "" {
		return "", fmt.Errorf("%s is not a valid domain", val)
	}
	ns := domainRegexp.FindSubmatch([]byte(val))
	if len(ns) > 0 && len(ns[1]) < 255 {
		return string(ns[1]), nil
	}
	return "", fmt.Errorf("%s is not a valid domain", val)
}

// ValidateLabel validates that the specified string is a valid label, and returns it.
// Labels are in the form on key=value.
func ValidateLabel(val string) (string, error) {
	if strings.Count(val, "=") < 1 {
		return "", fmt.Errorf("bad attribute format: %s", val)
	}
	return val, nil
}

// ValidateSysctl validates a sysctl and returns it.
func ValidateSysctl(val string) (string, error) {
	validSysctlMap := map[string]bool{
		"kernel.msgmax":          true,
		"kernel.msgmnb":          true,
		"kernel.msgmni":          true,
		"kernel.sem":             true,
		"kernel.shmall":          true,
		"kernel.shmmax":          true,
		"kernel.shmmni":          true,
		"kernel.shm_rmid_forced": true,
	}
	validSysctlPrefixes := []string{
		"net.",
		"fs.mqueue.",
	}
	arr := strings.Split(val, "=")
	if len(arr) < 2 {
		return "", fmt.Errorf("sysctl '%s' is not whitelisted", val)
	}
	if validSysctlMap[arr[0]] {
		return val, nil
	}

	for _, vp := range validSysctlPrefixes {
		if strings.HasPrefix(arr[0], vp) {
			return val, nil
		}
	}
	return "", fmt.Errorf("sysctl '%s' is not whitelisted", val)
}

// FilterOpt is a flag type for validating filters
type FilterOpt struct {
	filter filters.Args
}

// NewFilterOpt returns a new FilterOpt
func NewFilterOpt() FilterOpt {
	return FilterOpt{filter: filters.NewArgs()}
}

func (o *FilterOpt) String() string {
	repr, err := filters.ToParam(o.filter)
	if err != nil {
		return "invalid filters"
	}
	return repr
}

// Set sets the value of the opt by parsing the command line value
func (o *FilterOpt) Set(value string) error {
	var err error
	o.filter, err = filters.ParseFlag(value, o.filter)
	return err
}

// Type returns the option type
func (o *FilterOpt) Type() string {
	return "filter"
}

// Value returns the value of this option
func (o *FilterOpt) Value() filters.Args {
	return o.filter
}

// NanoCPUs is a type for fixed point fractional number.
type NanoCPUs int64

// String returns the string format of the number
func (c *NanoCPUs) String() string {
	return big.NewRat(c.Value(), 1e9).FloatString(3)
}

// Set sets the value of the NanoCPU by passing a string
func (c *NanoCPUs) Set(value string) error {
	cpus, err := ParseCPUs(value)
	*c = NanoCPUs(cpus)
	return err
}

// Type returns the type
func (c *NanoCPUs) Type() string {
	return "decimal"
}

// Value returns the value in int64
func (c *NanoCPUs) Value() int64 {
	return int64(*c)
}

// ParseCPUs takes a string ratio and returns an integer value of nano cpus
func ParseCPUs(value string) (int64, error) {
	cpu, ok := new(big.Rat).SetString(value)
	if !ok {
		return 0, fmt.Errorf("failed to parse %v as a rational number", value)
	}
	nano := cpu.Mul(cpu, big.NewRat(1e9, 1))
	if !nano.IsInt() {
		return 0, fmt.Errorf("value is too precise")
	}
	return nano.Num().Int64(), nil
}

// ParseLink parses and validates the specified string as a link format (name:alias)
func ParseLink(val string) (string, string, error) {
	if val == "" {
		return "", "", fmt.Errorf("empty string specified for links")
	}
	arr := strings.Split(val, ":")
	if len(arr) > 2 {
		return "", "", fmt.Errorf("bad format for links: %s", val)
	}
	if len(arr) == 1 {
		return val, val, nil
	}
	// This is kept because we can actually get a HostConfig with links
	// from an already created container and the format is not `foo:bar`
	// but `/foo:/c1/bar`
	if strings.HasPrefix(arr[0], "/") {
		_, alias := path.Split(arr[1])
		return arr[0][1:], alias, nil
	}
	return arr[0], arr[1], nil
}

// ValidateLink validates that the specified string has a valid link format (containerName:alias).
func ValidateLink(val string) (string, error) {
	_, _, err := ParseLink(val)
	return val, err
}

// MemBytes is a type for human readable memory bytes (like 128M, 2g, etc)
type MemBytes int64

// String returns the string format of the human readable memory bytes
func (m *MemBytes) String() string {
	// NOTE: In spf13/pflag/flag.go, "0" is considered as "zero value" while "0 B" is not.
	// We return "0" in case value is 0 here so that the default value is hidden.
	// (Sometimes "default 0 B" is actually misleading)
	if m.Value() != 0 {
		return units.BytesSize(float64(m.Value()))
	}
	return "0"
}

// Set sets the value of the MemBytes by passing a string
func (m *MemBytes) Set(value string) error {
	val, err := units.RAMInBytes(value)
	*m = MemBytes(val)
	return err
}

// Type returns the type
func (m *MemBytes) Type() string {
	return "bytes"
}

// Value returns the value in int64
func (m *MemBytes) Value() int64 {
	return int64(*m)
}

// UnmarshalJSON is the customized unmarshaler for MemBytes
func (m *MemBytes) UnmarshalJSON(s []byte) error {
	if len(s) <= 2 || s[0] != '"' || s[len(s)-1] != '"' {
		return fmt.Errorf("invalid size: %q", s)
	}
	val, err := units.RAMInBytes(string(s[1 : len(s)-1]))
	*m = MemBytes(val)
	return err
}

// MemSwapBytes is a type for human readable memory bytes (like 128M, 2g, etc).
// It differs from MemBytes in that -1 is valid and the default.
type MemSwapBytes int64

// Set sets the value of the MemSwapBytes by passing a string
func (m *MemSwapBytes) Set(value string) error {
	if value == "-1" {
		*m = MemSwapBytes(-1)
		return nil
	}
	val, err := units.RAMInBytes(value)
	*m = MemSwapBytes(val)
	return err
}

// Type returns the type
func (m *MemSwapBytes) Type() string {
	return "bytes"
}

// Value returns the value in int64
func (m *MemSwapBytes) Value() int64 {
	return int64(*m)
}

func (m *MemSwapBytes) String() string {
	b := MemBytes(*m)
	return b.String()
}

// UnmarshalJSON is the customized unmarshaler for MemSwapBytes
func (m *MemSwapBytes) UnmarshalJSON(s []byte) error {
	b := MemBytes(*m)
	return b.UnmarshalJSON(s)
}
<|endoftext|>