<commit_before>/*
Copyright 2014 The Kubernetes 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 cert

import (
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
)

// EncodePublicKeyPEM returns PEM-endcode public data
func EncodePublicKeyPEM(key *rsa.PublicKey) ([]byte, error) {
	der, err := x509.MarshalPKIXPublicKey(key)
	if err != nil {
		return []byte{}, err
	}
	block := pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: der,
	}
	return pem.EncodeToMemory(&block), nil
}

// EncodePrivateKeyPEM returns PEM-encoded private key data
func EncodePrivateKeyPEM(key *rsa.PrivateKey) []byte {
	block := pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(key),
	}
	return pem.EncodeToMemory(&block)
}

// EncodeCertPEM returns PEM-endcoded certificate data
func EncodeCertPEM(cert *x509.Certificate) []byte {
	block := pem.Block{
		Type:  "CERTIFICATE",
		Bytes: cert.Raw,
	}
	return pem.EncodeToMemory(&block)
}

// ParsePrivateKeyPEM returns a private key parsed from a PEM block in the supplied data.
// Recognizes PEM blocks for "EC PRIVATE KEY" and "RSA PRIVATE KEY"
func ParsePrivateKeyPEM(keyData []byte) (interface{}, error) {
	for {
		var privateKeyPemBlock *pem.Block
		privateKeyPemBlock, keyData = pem.Decode(keyData)
		if privateKeyPemBlock == nil {
			// we read all the PEM blocks and didn't recognize one
			return nil, fmt.Errorf("no private key PEM block found")
		}

		switch privateKeyPemBlock.Type {
		case "EC PRIVATE KEY":
			return x509.ParseECPrivateKey(privateKeyPemBlock.Bytes)
		case "RSA PRIVATE KEY":
			return x509.ParsePKCS1PrivateKey(privateKeyPemBlock.Bytes)
		}
	}
}

// ParseCertsPEM returns the x509.Certificates contained in the given PEM-encoded byte array
// Returns an error if a certificate could not be parsed, or if the data does not contain any certificates
func ParseCertsPEM(pemCerts []byte) ([]*x509.Certificate, error) {
	ok := false
	certs := []*x509.Certificate{}
	for len(pemCerts) > 0 {
		var block *pem.Block
		block, pemCerts = pem.Decode(pemCerts)
		if block == nil {
			break
		}
		// Only use PEM "CERTIFICATE" blocks without extra headers
		if block.Type != "CERTIFICATE" || len(block.Headers) != 0 {
			continue
		}

		cert, err := x509.ParseCertificate(block.Bytes)
		if err != nil {
			return certs, err
		}

		certs = append(certs, cert)
		ok = true
	}

	if !ok {
		return certs, errors.New("could not read any certificates")
	}
	return certs, nil
}
<commit_msg>Add PKCS#8 support to ParsePrivateKeyPEM<commit_after>/*
Copyright 2014 The Kubernetes 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 cert

import (
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
)

// EncodePublicKeyPEM returns PEM-endcode public data
func EncodePublicKeyPEM(key *rsa.PublicKey) ([]byte, error) {
	der, err := x509.MarshalPKIXPublicKey(key)
	if err != nil {
		return []byte{}, err
	}
	block := pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: der,
	}
	return pem.EncodeToMemory(&block), nil
}

// EncodePrivateKeyPEM returns PEM-encoded private key data
func EncodePrivateKeyPEM(key *rsa.PrivateKey) []byte {
	block := pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(key),
	}
	return pem.EncodeToMemory(&block)
}

// EncodeCertPEM returns PEM-endcoded certificate data
func EncodeCertPEM(cert *x509.Certificate) []byte {
	block := pem.Block{
		Type:  "CERTIFICATE",
		Bytes: cert.Raw,
	}
	return pem.EncodeToMemory(&block)
}

// ParsePrivateKeyPEM returns a private key parsed from a PEM block in the supplied data.
// Recognizes PEM blocks for "EC PRIVATE KEY", "RSA PRIVATE KEY", or "PRIVATE KEY"
func ParsePrivateKeyPEM(keyData []byte) (interface{}, error) {
	var privateKeyPemBlock *pem.Block
	for {
		privateKeyPemBlock, keyData = pem.Decode(keyData)
		if privateKeyPemBlock == nil {
			break
		}

		switch privateKeyPemBlock.Type {
		case "EC PRIVATE KEY":
			// ECDSA Private Key in ASN.1 format
			if key, err := x509.ParseECPrivateKey(privateKeyPemBlock.Bytes); err == nil {
				return key, nil
			}
		case "RSA PRIVATE KEY":
			// RSA Private Key in PKCS#1 format
			if key, err := x509.ParsePKCS1PrivateKey(privateKeyPemBlock.Bytes); err == nil {
				return key, nil
			}
		case "PRIVATE KEY":
			// RSA or ECDSA Private Key in unencrypted PKCS#8 format
			if key, err := x509.ParsePKCS8PrivateKey(privateKeyPemBlock.Bytes); err == nil {
				return key, nil
			}
		}

		// tolerate non-key PEM blocks for compatibility with things like "EC PARAMETERS" blocks
		// originally, only the first PEM block was parsed and expected to be a key block
	}

	// we read all the PEM blocks and didn't recognize one
	return nil, fmt.Errorf("data does not contain a valid RSA or ECDSA private key")
}

// ParseCertsPEM returns the x509.Certificates contained in the given PEM-encoded byte array
// Returns an error if a certificate could not be parsed, or if the data does not contain any certificates
func ParseCertsPEM(pemCerts []byte) ([]*x509.Certificate, error) {
	ok := false
	certs := []*x509.Certificate{}
	for len(pemCerts) > 0 {
		var block *pem.Block
		block, pemCerts = pem.Decode(pemCerts)
		if block == nil {
			break
		}
		// Only use PEM "CERTIFICATE" blocks without extra headers
		if block.Type != "CERTIFICATE" || len(block.Headers) != 0 {
			continue
		}

		cert, err := x509.ParseCertificate(block.Bytes)
		if err != nil {
			return certs, err
		}

		certs = append(certs, cert)
		ok = true
	}

	if !ok {
		return certs, errors.New("could not read any certificates")
	}
	return certs, nil
}
<|endoftext|>


<commit_before>package metadata

import (
	"testing"

	"github.com/jacksontj/dataman/datamantype"
)

type fieldValidationCase struct {
	field      *CollectionField
	goodValues []interface{}
	badValues  []interface{}
}

func (f *fieldValidationCase) Test(t *testing.T) {
	// Check the positive cases
	for _, val := range f.goodValues {
		if _, validationResult := f.field.Normalize(val); !validationResult.IsValid() {
			t.Errorf("Error validating value %v: %v", val, validationResult)
		}
	}

	// Check the negative cases
	for _, val := range f.badValues {
		if _, validationResult := f.field.Normalize(val); validationResult.IsValid() {
			t.Errorf("No error validating a bad value: %v\n%v", val, validationResult)
		}
	}
}

func TestFieldValidation_Document(t *testing.T) {
	testCase := &fieldValidationCase{
		field: &CollectionField{
			Type:      "_document",
			FieldType: DatamanTypeToFieldType(datamantype.Document),
			SubFields: map[string]*CollectionField{
				"name": {
					Name:      "name",
					Type:      "_string",
					FieldType: DatamanTypeToFieldType(datamantype.String),
					NotNull:   true,
				},
				"number": {
					Name:      "number",
					Type:      "_int",
					FieldType: DatamanTypeToFieldType(datamantype.Int),
				},
				"subDoc": {
					Type:      "_document",
					FieldType: DatamanTypeToFieldType(datamantype.Document),
					SubFields: map[string]*CollectionField{
						"name": {
							Name:      "name",
							Type:      datamantype.String,
							FieldType: DatamanTypeToFieldType(datamantype.String),
							NotNull:   true,
						},
					},
				},
			},
		},
		goodValues: []interface{}{
			map[string]interface{}{
				"name":   "someone",
				"number": 10,
			},
			map[string]interface{}{
				"name":   "someone",
				"number": 10,
				"subDoc": map[string]interface{}{
					"name":              "subname",
					"somethingelse":     10,
					"somethingelsemore": "yea",
				},
			},
		},
		badValues: []interface{}{
			"aString", // A string
			1,         // a number
			nil,       // nil
			map[int]interface{}{1: "foo"}, // wrong map type
			map[string]interface{}{
				"number": 10,
			},
			map[string]interface{}{
				"name":   "someone",
				"number": 10,
				"subDoc": map[string]interface{}{},
			},
		},
	}

	testCase.Test(t)
}

func TestFieldValidation_String(t *testing.T) {
	testCase := &fieldValidationCase{
		field: &CollectionField{
			Type:      datamantype.String,
			FieldType: DatamanTypeToFieldType(datamantype.String),
		},
		goodValues: []interface{}{
			"foo",
			"f",
			"",
			nil, // nil
		},
		badValues: []interface{}{
			1, // a number
		},
	}
	testCase.Test(t)
}

func TestFieldValidation_Int(t *testing.T) {
	testCase := &fieldValidationCase{
		field: &CollectionField{
			Type:      datamantype.Int,
			FieldType: DatamanTypeToFieldType(datamantype.Int),
		},
		goodValues: []interface{}{
			0,
			-10,
			100,
			0.0, // float
			nil, // nil
		},
		badValues: []interface{}{
			"string", // string
		},
	}
	testCase.Test(t)
}

func TestFieldValidation_Bool(t *testing.T) {
	testCase := &fieldValidationCase{
		field: &CollectionField{
			Type:      datamantype.Bool,
			FieldType: DatamanTypeToFieldType(datamantype.Bool),
		},
		goodValues: []interface{}{
			true,
			false,
			nil, // nil
		},
		badValues: []interface{}{
			"string", // string
			0.0,      // float
		},
	}
	testCase.Test(t)
}

// TODO: do this one
/*
func TestFieldValidation_DateTime(t *testing.T) {
	testCase := &fieldValidationCase{
		field: &CollectionField{
			Type: DateTime,
		},
		goodValues: []interface{}{
			true,
			false,
		},
		badValues: []interface{}{
			"string", // string
			nil, // nil
			0.0, // float
		},
	}
	testCase.Test(t)
}
*/
<commit_msg>Fix test case<commit_after>package metadata

import (
	"testing"

	"github.com/jacksontj/dataman/datamantype"
)

type fieldValidationCase struct {
	field      *CollectionField
	goodValues []interface{}
	badValues  []interface{}
}

func (f *fieldValidationCase) Test(t *testing.T) {
	// Check the positive cases
	for _, val := range f.goodValues {
		if _, validationResult := f.field.Normalize(val); !validationResult.IsValid() {
			t.Errorf("Error validating value %v: %v", val, validationResult)
		}
	}

	// Check the negative cases
	for _, val := range f.badValues {
		if _, validationResult := f.field.Normalize(val); validationResult.IsValid() {
			t.Errorf("No error validating a bad value: %v\n%v", val, validationResult)
		}
	}
}

func TestFieldValidation_Document(t *testing.T) {
	testCase := &fieldValidationCase{
		field: &CollectionField{
			Type:      "_document",
			FieldType: DatamanTypeToFieldType(datamantype.Document),
			SubFields: map[string]*CollectionField{
				"name": {
					Name:      "name",
					Type:      "_string",
					FieldType: DatamanTypeToFieldType(datamantype.String),
					NotNull:   true,
				},
				"number": {
					Name:      "number",
					Type:      "_int",
					FieldType: DatamanTypeToFieldType(datamantype.Int),
				},
				"subDoc": {
					Type:      "_document",
					FieldType: DatamanTypeToFieldType(datamantype.Document),
					SubFields: map[string]*CollectionField{
						"name": {
							Name:      "name",
							Type:      "_string",
							FieldType: DatamanTypeToFieldType(datamantype.String),
							NotNull:   true,
						},
					},
				},
			},
		},
		goodValues: []interface{}{
			map[string]interface{}{
				"name":   "someone",
				"number": 10,
			},
			map[string]interface{}{
				"name":   "someone",
				"number": 10,
				"subDoc": map[string]interface{}{
					"name":              "subname",
					"somethingelse":     10,
					"somethingelsemore": "yea",
				},
			},
		},
		badValues: []interface{}{
			"aString", // A string
			1,         // a number
			nil,       // nil
			map[int]interface{}{1: "foo"}, // wrong map type
			map[string]interface{}{
				"number": 10,
			},
			map[string]interface{}{
				"name":   "someone",
				"number": 10,
				"subDoc": map[string]interface{}{},
			},
		},
	}

	testCase.Test(t)
}

func TestFieldValidation_String(t *testing.T) {
	testCase := &fieldValidationCase{
		field: &CollectionField{
			Type:      "_string",
			FieldType: DatamanTypeToFieldType(datamantype.String),
		},
		goodValues: []interface{}{
			"foo",
			"f",
			"",
			1,   // a number
			nil, // nil
		},
		badValues: []interface{}{
			struct{}{},
		},
	}
	testCase.Test(t)
}

func TestFieldValidation_Int(t *testing.T) {
	testCase := &fieldValidationCase{
		field: &CollectionField{
			Type:      "_int",
			FieldType: DatamanTypeToFieldType(datamantype.Int),
		},
		goodValues: []interface{}{
			0,
			-10,
			100,
			0.0, // float
			nil, // nil
		},
		badValues: []interface{}{
			"string", // string
		},
	}
	testCase.Test(t)
}

func TestFieldValidation_Bool(t *testing.T) {
	testCase := &fieldValidationCase{
		field: &CollectionField{
			Type:      "_bool",
			FieldType: DatamanTypeToFieldType(datamantype.Bool),
		},
		goodValues: []interface{}{
			true,
			false,
			nil, // nil
		},
		badValues: []interface{}{
			"string", // string
			0.0,      // float
		},
	}
	testCase.Test(t)
}

// TODO: do this one
/*
func TestFieldValidation_DateTime(t *testing.T) {
	testCase := &fieldValidationCase{
		field: &CollectionField{
			Type: DateTime,
		},
		goodValues: []interface{}{
			true,
			false,
		},
		badValues: []interface{}{
			"string", // string
			nil, // nil
			0.0, // float
		},
	}
	testCase.Test(t)
}
*/
<|endoftext|>


<commit_before>package sqlconstants

import "errors"
import "strings"

const CURRENT_VENDOR = SQLITE3
const SQLITE3_FILE_URL = "throttle.db"
const SQLITE3_CONNECTION_STRING = "file:" + SQLITE3_FILE_URL
const SQLITE3_CHECK_IF_TABLE_EXISTS = "SELECT NAME FROM SQLITE_MASTER WHERE TYPE=? AND NAME=?;"
const SQLITE3_CREATE_BUCKET_SCHEMA = "CREATE TABLE BUCKETS(NAME TEXT NOT NULL UNIQUE PRIMARY KEY, VOLUME INTEGER DEFAULT 0, CAPACITY INTEGER, TIMEFRAME INTEGER, CREATED_TIMESTAMP TIMESTAMP DEFAULT CURRENT_TIMESTAMP, UPDATED_TIMESTAMP INTEGER DEFAULT (strftime('%s','now')))"
const SQLITE3_INSERT_INTO_BUCKET = "INSERT INTO BUCKETS (NAME, VOLUME, CAPACITY, TIMEFRAME) VALUES (?, ?, ?, ?)"
const SQLITE3_UPDATE_BUCKET = "UPDATE BUCKETS SET (VOLUME,CAPACITY,TIMEFRAME, UPDATED_TIMESTAMP) VALUES (?, ?, ?, ?)"

const SQLITE3_REFRESH_BUCKET = "UPDATE BUCKETS SET VOLUME = (SELECT CASE WHEN expr > 0 THEN expr ELSE 0 END FROM (SELECT CAST(ROUND(VOLUME - ((CAPACITY*1.0)/TIMEFRAME)*(STRFTIME('%s','now') - UPDATED_TIMESTAMP)) AS INT) AS expr FROM BUCKETS WHERE NAME = ?) b), UPDATED_TIMESTAMP = STRFTIME('%s', 'now') WHERE NAME = ?;"
const SQLITE3_LOAD_BUCKET = "SELECT VOLUME, CAPACITY, TIMEFRAME FROM BUCKETS WHERE NAME = ?"
const SQLITE3_DELETE_BUCKET = "DELETE BUCKETS WHERE NAME = ?"
const SQLITE3_SELECT_ALL_BUCKETS = "SELECT * FROM BUCKETS where NAME = ?"
const SQLITE3_BUCKET_NAME = "BUCKETS"
const SQLITE3_GET_SCHEMA = "SELECT SQL FROM SQLITE_MASTER WHERE TYPE=? AND NAME=?"
const SQLITE3_FILL_BUCKET = "UPDATE BUCKETS SET VOLUME = VOLUME + ? WHERE NAME = ?"
const SQLITE3_DROP_BUCKETS = "DROP TABLE BUCKETS"

const (
	SQLITE3_TYPE_NULL      = ""
	SQLITE3_TYPE_INTEGER   = "INTEGER"
	SQLITE3_TYPE_TEXT      = "TEXT"
	SQLITE3_TYPE_REAL      = "REAL"
	SQLITE3_TYPE_BLOB      = "BLOG"
	SQLITE3_TYPE_BOOL      = "BOOL"
	SQLITE3_TYPE_TIMESTAMP = "TIMESTAMP"
	//add postgres, plus other types here
)

const (
	SQLITE3 = iota
	POSTGRESQL
	MYSQL
	MONGODB
)

func CreateStatementFunctionsToReplace() ([]string, error) {
	if CURRENT_VENDOR == SQLITE3 {
		return []string{"(strftime('%s','now'))"}, nil
	}

	return nil, errors.New("Error: Current Vendor is not support in CreateStatementFunctionsToReplace")
}

func CurrentVendor() string {
	if CURRENT_VENDOR == SQLITE3 {
		return "sqlite3"
	} else if CURRENT_VENDOR == POSTGRESQL {
		return "postgres"
	} else if CURRENT_VENDOR == MYSQL {
		return "mysql"
	} else if CURRENT_VENDOR == MONGODB {
		return "mongodb"
	}

	return ""
}

func IsSQLConstraint(constraint string) (bool, error) {
	c := strings.ToUpper(constraint)
	if CURRENT_VENDOR == SQLITE3 {
		switch c {
		case "PRIMARY":
			return true, nil
		case "DEFAULT":
			return true, nil
		case "NOT":
			return true, nil
		case "UNIQUE":
			return true, nil
		case "CHECK":
			return true, nil
		case "REFERENCES":
			return true, nil
		case "COLLATE":
			return true, nil
		default:
			return false, nil
		}

	}

	return false, errors.New("Go type no recongnized for current db vendor => " + CurrentVendor())

}

func GoType(sqlType string) (string, error) {
	if CURRENT_VENDOR == SQLITE3 {
		switch sqlType {
		case SQLITE3_TYPE_NULL:
			return "nil", nil
		case SQLITE3_TYPE_INTEGER:
			return "int", nil
		case SQLITE3_TYPE_REAL:
			return "float64", nil
		case SQLITE3_TYPE_BOOL:
			return "bool", nil
		case SQLITE3_TYPE_TEXT:
			return "string", nil
		case SQLITE3_TYPE_TIMESTAMP:
			return "time.Time", nil
		default:
			return "", nil

		}
	}

	return "", errors.New("Go type no recongnized for current db vendor => " + CurrentVendor())

}

func SQLType(goType string) (string, error) {
	if CURRENT_VENDOR == SQLITE3 {
		switch goType {
		case "nil":
			return SQLITE3_TYPE_NULL, nil
		case "int":
			return SQLITE3_TYPE_INTEGER, nil
		case "float64":
			return SQLITE3_TYPE_REAL, nil
		case "bool":
			return SQLITE3_TYPE_BOOL, nil
		case "string":
			return SQLITE3_TYPE_TEXT, nil
		case "time.Time":
			return SQLITE3_TYPE_TIMESTAMP, nil

		}
	}

	return "", errors.New("Go type no recongnized for current db vendor => " + CurrentVendor())

}
<commit_msg>Added functions to create the equivalent sqlite functions to get x (weeks,days,months) from the startdate.<commit_after>package sqlconstants

import "errors"
import "strings"
import "strconv"
import "time"
import "../../timeutil"

const CURRENT_VENDOR = SQLITE3
const SQLITE3_FILE_URL = "throttle.db"
const SQLITE3_CONNECTION_STRING = "file:" + SQLITE3_FILE_URL
const SQLITE3_CHECK_IF_TABLE_EXISTS = "SELECT NAME FROM SQLITE_MASTER WHERE TYPE=? AND NAME=?;"
const SQLITE3_CREATE_BUCKET_SCHEMA = "CREATE TABLE BUCKETS(NAME TEXT NOT NULL UNIQUE PRIMARY KEY, VOLUME INTEGER DEFAULT 0, CAPACITY INTEGER, TIMEFRAME INTEGER, CREATED_TIMESTAMP TIMESTAMP DEFAULT CURRENT_TIMESTAMP, UPDATED_TIMESTAMP INTEGER DEFAULT (strftime('%s','now')))"
const SQLITE3_INSERT_INTO_BUCKET = "INSERT INTO BUCKETS (NAME, VOLUME, CAPACITY, TIMEFRAME) VALUES (?, ?, ?, ?)"
const SQLITE3_UPDATE_BUCKET = "UPDATE BUCKETS SET (VOLUME,CAPACITY,TIMEFRAME, UPDATED_TIMESTAMP) VALUES (?, ?, ?, ?)"

const SQLITE3_REFRESH_BUCKET = "UPDATE BUCKETS SET VOLUME = (SELECT CASE WHEN expr > 0 THEN expr ELSE 0 END FROM (SELECT CAST(ROUND(VOLUME - ((CAPACITY*1.0)/TIMEFRAME)*(STRFTIME('%s','now') - UPDATED_TIMESTAMP)) AS INT) AS expr FROM BUCKETS WHERE NAME = ?) b), UPDATED_TIMESTAMP = STRFTIME('%s', 'now') WHERE NAME = ?;"
const SQLITE3_LOAD_BUCKET = "SELECT VOLUME, CAPACITY, TIMEFRAME FROM BUCKETS WHERE NAME = ?"
const SQLITE3_DELETE_BUCKET = "DELETE BUCKETS WHERE NAME = ?"
const SQLITE3_SELECT_ALL_BUCKETS = "SELECT * FROM BUCKETS where NAME = ?"
const SQLITE3_BUCKET_NAME = "BUCKETS"
const SQLITE3_GET_SCHEMA = "SELECT SQL FROM SQLITE_MASTER WHERE TYPE=? AND NAME=?"
const SQLITE3_FILL_BUCKET = "UPDATE BUCKETS SET VOLUME = VOLUME + ? WHERE NAME = ?"
const SQLITE3_DROP_BUCKETS = "DROP TABLE BUCKETS"

const (
	SQLITE3_TYPE_NULL      = ""
	SQLITE3_TYPE_INTEGER   = "INTEGER"
	SQLITE3_TYPE_TEXT      = "TEXT"
	SQLITE3_TYPE_REAL      = "REAL"
	SQLITE3_TYPE_BLOB      = "BLOG"
	SQLITE3_TYPE_BOOL      = "BOOL"
	SQLITE3_TYPE_TIMESTAMP = "TIMESTAMP"
	//add postgres, plus other types here
)

const (
	SQLITE3 = iota
	POSTGRESQL
	MYSQL
	MONGODB
)

func FireTriggerEveryXDaysFrom(triggerDay int, startTimeDate time.Time) (string, error) {

	if CurrentVendor() != "sqlite3" {
		return "", errors.New("Vendor not recognized ")
	}
	baseSQL := "DATETIME('" + timeutil.NormalFormat(startTimeDate) + "',"
	baseSQL = baseSQL + "'+" + strconv.Itoa(triggerDay) + " days')"

	return baseSQL, nil
}

func FireTriggerEveryXMonthsFrom(triggerMonth int, startTimeDate time.Time) (string, error) {

	if CurrentVendor() != "sqlite3" {
		return "", errors.New("Vendor not recognized ")
	}
	baseSQL := "DATETIME('" + timeutil.NormalFormat(startTimeDate) + "',"
	baseSQL = baseSQL + "'+" + strconv.Itoa(triggerMonth) + " months')"

	return baseSQL, nil
}

func FireTriggerEveryXWeeksFrom(triggerWeek int, startTimeDate time.Time) (string, error) {

	if CurrentVendor() != "sqlite3" {
		return "", errors.New("Vendor not recognized ")
	}
	triggerWeek = triggerWeek * 7

	baseSQL := "DATETIME('" + timeutil.NormalFormat(startTimeDate) + "',"
	baseSQL = baseSQL + "'+" + strconv.Itoa(triggerWeek) + " days')"

	return baseSQL, nil
}

func CreateStatementFunctionsToReplace() ([]string, error) {
	if CURRENT_VENDOR == SQLITE3 {
		return []string{"(strftime('%s','now'))"}, nil
	}

	return nil, errors.New("Error: Current Vendor is not support in CreateStatementFunctionsToReplace")
}

func CurrentVendor() string {
	if CURRENT_VENDOR == SQLITE3 {
		return "sqlite3"
	} else if CURRENT_VENDOR == POSTGRESQL {
		return "postgres"
	} else if CURRENT_VENDOR == MYSQL {
		return "mysql"
	} else if CURRENT_VENDOR == MONGODB {
		return "mongodb"
	}

	return ""
}

func IsSQLConstraint(constraint string) (bool, error) {
	c := strings.ToUpper(constraint)
	if CURRENT_VENDOR == SQLITE3 {
		switch c {
		case "PRIMARY":
			return true, nil
		case "DEFAULT":
			return true, nil
		case "NOT":
			return true, nil
		case "UNIQUE":
			return true, nil
		case "CHECK":
			return true, nil
		case "REFERENCES":
			return true, nil
		case "COLLATE":
			return true, nil
		default:
			return false, nil
		}

	}

	return false, errors.New("Go type no recongnized for current db vendor => " + CurrentVendor())

}

func GoType(sqlType string) (string, error) {
	if CURRENT_VENDOR == SQLITE3 {
		switch sqlType {
		case SQLITE3_TYPE_NULL:
			return "nil", nil
		case SQLITE3_TYPE_INTEGER:
			return "int", nil
		case SQLITE3_TYPE_REAL:
			return "float64", nil
		case SQLITE3_TYPE_BOOL:
			return "bool", nil
		case SQLITE3_TYPE_TEXT:
			return "string", nil
		case SQLITE3_TYPE_TIMESTAMP:
			return "time.Time", nil
		default:
			return "", nil

		}
	}

	return "", errors.New("Go type no recongnized for current db vendor => " + CurrentVendor())

}

func SQLType(goType string) (string, error) {
	if CURRENT_VENDOR == SQLITE3 {
		switch goType {
		case "nil":
			return SQLITE3_TYPE_NULL, nil
		case "int":
			return SQLITE3_TYPE_INTEGER, nil
		case "float64":
			return SQLITE3_TYPE_REAL, nil
		case "bool":
			return SQLITE3_TYPE_BOOL, nil
		case "string":
			return SQLITE3_TYPE_TEXT, nil
		case "time.Time":
			return SQLITE3_TYPE_TIMESTAMP, nil

		}
	}

	return "", errors.New("Go type no recongnized for current db vendor => " + CurrentVendor())

}
<|endoftext|>


<commit_before>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"sync"
)

// Service names used in ConnectionStatus.
const (
	KeybaseServiceName     = "keybase-service"
	MDServiceName          = "md-server"
	LoginStatusUpdateName  = "login"
	LogoutStatusUpdateName = "logout"
)

type errDisconnected struct{}

func (errDisconnected) Error() string { return "Disconnected" }

type kbfsCurrentStatus struct {
	lock            sync.Mutex
	failingServices map[string]error
	invalidateChan  chan StatusUpdate
}

// Init inits the kbfsCurrentStatus.
func (kcs *kbfsCurrentStatus) Init() {
	kcs.failingServices = map[string]error{}
	kcs.invalidateChan = make(chan StatusUpdate)
}

// CurrentStatus returns a copy of the current status.
func (kcs *kbfsCurrentStatus) CurrentStatus() (map[string]error, chan StatusUpdate) {
	kcs.lock.Lock()
	defer kcs.lock.Unlock()

	res := map[string]error{}
	for k, v := range kcs.failingServices {
		res[k] = v
	}
	return res, kcs.invalidateChan
}

// PushConnectionStatusChange pushes a change to the connection status of one of the services.
func (kcs *kbfsCurrentStatus) PushConnectionStatusChange(service string, err error) {
	kcs.lock.Lock()
	defer kcs.lock.Unlock()

	if err != nil {
		kcs.failingServices[service] = err
	} else {
		// Potentially exit early if nothing changes.
		_, exist := kcs.failingServices[service]
		if !exist {
			return
		}
		delete(kcs.failingServices, service)
	}

	close(kcs.invalidateChan)
	kcs.invalidateChan = make(chan StatusUpdate)
}

// PushStatusChange forces a new status be fetched by status listeners.
func (kcs *kbfsCurrentStatus) PushStatusChange() {
	kcs.lock.Lock()
	defer kcs.lock.Unlock()
	close(kcs.invalidateChan)
	kcs.invalidateChan = make(chan StatusUpdate)
}
<commit_msg>connection_status: don't invalidate if service is already errored<commit_after>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"sync"
)

// Service names used in ConnectionStatus.
const (
	KeybaseServiceName     = "keybase-service"
	MDServiceName          = "md-server"
	LoginStatusUpdateName  = "login"
	LogoutStatusUpdateName = "logout"
)

type errDisconnected struct{}

func (errDisconnected) Error() string { return "Disconnected" }

type kbfsCurrentStatus struct {
	lock            sync.Mutex
	failingServices map[string]error
	invalidateChan  chan StatusUpdate
}

// Init inits the kbfsCurrentStatus.
func (kcs *kbfsCurrentStatus) Init() {
	kcs.failingServices = map[string]error{}
	kcs.invalidateChan = make(chan StatusUpdate)
}

// CurrentStatus returns a copy of the current status.
func (kcs *kbfsCurrentStatus) CurrentStatus() (map[string]error, chan StatusUpdate) {
	kcs.lock.Lock()
	defer kcs.lock.Unlock()

	res := map[string]error{}
	for k, v := range kcs.failingServices {
		res[k] = v
	}
	return res, kcs.invalidateChan
}

// PushConnectionStatusChange pushes a change to the connection status of one of the services.
func (kcs *kbfsCurrentStatus) PushConnectionStatusChange(service string, err error) {
	kcs.lock.Lock()
	defer kcs.lock.Unlock()

	if err != nil {
		// Exit early if the service is already failed, to avoid an
		// invalidation.
		_, errExisted := kcs.failingServices[service]
		kcs.failingServices[service] = err
		if errExisted {
			return
		}
	} else {
		// Potentially exit early if nothing changes.
		_, exist := kcs.failingServices[service]
		if !exist {
			return
		}
		delete(kcs.failingServices, service)
	}

	close(kcs.invalidateChan)
	kcs.invalidateChan = make(chan StatusUpdate)
}

// PushStatusChange forces a new status be fetched by status listeners.
func (kcs *kbfsCurrentStatus) PushStatusChange() {
	kcs.lock.Lock()
	defer kcs.lock.Unlock()
	close(kcs.invalidateChan)
	kcs.invalidateChan = make(chan StatusUpdate)
}
<|endoftext|>


<commit_before>package currency

import (
	"encoding/xml"
	"github.com/mbanzon/simplehttp"
	"time"
)

const (
	ecbResourceUrl = "http://www.ecb.europa.eu/stats/eurofxref/eurofxref-daily.xml"
	dateFormat     = "2006-01-02"
)

type currencyEnvelope struct {
	Sender string `xml:"Sender>name"`
	Cube   []cube `xml:"Cube>Cube>Cube"`
}

type timeEnvelope struct {
	Time timeCube `xml:"Cube>Cube"`
}

type timeCube struct {
	Time string `xml:"time,attr"`
}

// structure used for parsing the currencies.
type cube struct {
	Name string  `xml:"currency,attr"`
	Rate float64 `xml:"rate,attr"`
}

func parseEcbData() (time.Time, map[string]float64, error) {
	r := simplehttp.NewGetRequest(ecbResourceUrl)
	data, err := r.MakeRequest()
	if err != nil {
		return time.Time{}, nil, err
	}

	var c currencyEnvelope
	err = xml.Unmarshal(data, &c) // first unmarshal to get currencies
	if err != nil {
		return time.Time{}, nil, err
	}

	var t timeEnvelope
	err = xml.Unmarshal(data, &t) // second unmarshal to get time
	if err != nil {
		return time.Time{}, nil, err
	}

	currencyTime, err := time.Parse(dateFormat, t.Time.Time)
	if err != nil {
		return time.Time{}, nil, err
	}

	currencies := make(map[string]float64)

	currencies[eur] = 1

	for _, currency := range c.Cube {
		currencies[currency.Name] = currency.Rate
	}

	return currencyTime, currencies, nil
}
<commit_msg>Comment.<commit_after>package currency

import (
	"encoding/xml"
	"github.com/mbanzon/simplehttp"
	"time"
)

const (
	ecbResourceUrl = "http://www.ecb.europa.eu/stats/eurofxref/eurofxref-daily.xml"
	dateFormat     = "2006-01-02"
)

type currencyEnvelope struct {
	Sender string `xml:"Sender>name"`
	Cube   []cube `xml:"Cube>Cube>Cube"`
}

type timeEnvelope struct {
	Time timeCube `xml:"Cube>Cube"`
}

// structure used for parsing the time.
type timeCube struct {
	Time string `xml:"time,attr"`
}

// structure used for parsing the currencies.
type cube struct {
	Name string  `xml:"currency,attr"`
	Rate float64 `xml:"rate,attr"`
}

func parseEcbData() (time.Time, map[string]float64, error) {
	r := simplehttp.NewGetRequest(ecbResourceUrl)
	data, err := r.MakeRequest()
	if err != nil {
		return time.Time{}, nil, err
	}

	var c currencyEnvelope
	err = xml.Unmarshal(data, &c) // first unmarshal to get currencies
	if err != nil {
		return time.Time{}, nil, err
	}

	var t timeEnvelope
	err = xml.Unmarshal(data, &t) // second unmarshal to get time
	if err != nil {
		return time.Time{}, nil, err
	}

	currencyTime, err := time.Parse(dateFormat, t.Time.Time)
	if err != nil {
		return time.Time{}, nil, err
	}

	currencies := make(map[string]float64)

	currencies[eur] = 1

	for _, currency := range c.Cube {
		currencies[currency.Name] = currency.Rate
	}

	return currencyTime, currencies, nil
}
<|endoftext|>


<commit_before>package librariesio

import (
	"context"
	"fmt"
	"net/http"
	"reflect"
	"strings"
	"testing"
)

func TestProject(t *testing.T) {
	server, mux, url := startNewServer()
	client := NewClient(APIKey)
	client.BaseURL = url
	defer server.Close()

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if method := "GET"; method != r.Method {
			t.Errorf("expected HTTP %v request, got %v", method, r.Method)
		}
		fmt.Fprintf(w, `{"name":"cookiecutter"}`)
	})

	project, _, err := client.Project(context.Background(), "pypi", "cookiecutter")
	if err != nil {
		t.Fatalf("Project returned unexpected error: %v", err)
	}

	name := "cookiecutter"
	want := &Project{Name: &name}

	if !reflect.DeepEqual(project, want) {
		t.Errorf("\nExpected %#v\nGot %#v", want, project)
	}
}

func TestProjectDeps(t *testing.T) {
	server, mux, url := startNewServer()
	client := NewClient(APIKey)
	client.BaseURL = url
	defer server.Close()

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if method := "GET"; method != r.Method {
			t.Errorf("expected HTTP %v request, got %v", method, r.Method)
		}

		if url := r.URL.String(); !strings.Contains(url, "/dependencies") {
			t.Errorf("unexpected URL, got %v", url)
		}

		fmt.Fprintf(w, `{
			"name":"ava",
			"dependencies": [
				{"name": "mocha", "latest": "3.2.0", "outdated": false},
				{"name": "chalk", "latest": "1.1.3", "outdated": true}
			]
		}`)
	})

	project, _, err := client.ProjectDeps(context.Background(), "npm", "ava", "latest")

	if err != nil {
		t.Fatalf("ProjectDeps returned unexpected error: %v", err)
	}

	want := &Project{
		Name: String("ava"),
		Dependencies: []*ProjectDependency{
			&ProjectDependency{
				Latest:   String("3.2.0"),
				Name:     String("mocha"),
				Outdated: Bool(false),
			},
			&ProjectDependency{
				Latest:   String("1.1.3"),
				Name:     String("chalk"),
				Outdated: Bool(true),
			},
		},
	}

	if !reflect.DeepEqual(project, want) {
		t.Errorf("\nExpected %#v\nGot %#v", want, project)
	}
}
<commit_msg>Simplify code as reported with 'gofmt -s'<commit_after>package librariesio

import (
	"context"
	"fmt"
	"net/http"
	"reflect"
	"strings"
	"testing"
)

func TestProject(t *testing.T) {
	server, mux, url := startNewServer()
	client := NewClient(APIKey)
	client.BaseURL = url
	defer server.Close()

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if method := "GET"; method != r.Method {
			t.Errorf("expected HTTP %v request, got %v", method, r.Method)
		}
		fmt.Fprintf(w, `{"name":"cookiecutter"}`)
	})

	project, _, err := client.Project(context.Background(), "pypi", "cookiecutter")
	if err != nil {
		t.Fatalf("Project returned unexpected error: %v", err)
	}

	name := "cookiecutter"
	want := &Project{Name: &name}

	if !reflect.DeepEqual(project, want) {
		t.Errorf("\nExpected %#v\nGot %#v", want, project)
	}
}

func TestProjectDeps(t *testing.T) {
	server, mux, url := startNewServer()
	client := NewClient(APIKey)
	client.BaseURL = url
	defer server.Close()

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if method := "GET"; method != r.Method {
			t.Errorf("expected HTTP %v request, got %v", method, r.Method)
		}

		if url := r.URL.String(); !strings.Contains(url, "/dependencies") {
			t.Errorf("unexpected URL, got %v", url)
		}

		fmt.Fprintf(w, `{
			"name":"ava",
			"dependencies": [
				{"name": "mocha", "latest": "3.2.0", "outdated": false},
				{"name": "chalk", "latest": "1.1.3", "outdated": true}
			]
		}`)
	})

	project, _, err := client.ProjectDeps(context.Background(), "npm", "ava", "latest")

	if err != nil {
		t.Fatalf("ProjectDeps returned unexpected error: %v", err)
	}

	want := &Project{
		Name: String("ava"),
		Dependencies: []*ProjectDependency{
			{
				Latest:   String("3.2.0"),
				Name:     String("mocha"),
				Outdated: Bool(false),
			},
			{
				Latest:   String("1.1.3"),
				Name:     String("chalk"),
				Outdated: Bool(true),
			},
		},
	}

	if !reflect.DeepEqual(project, want) {
		t.Errorf("\nExpected %#v\nGot %#v", want, project)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	dockerlib "github.com/fsouza/go-dockerclient"
	"github.com/newrelic/go_nagios"
	"github.com/shenwei356/util/bytesize"
	"strings"
	"sync"
)

func NewCheckDocker(endpoint string) (*CheckDocker, error) {
	var err error
	cd := &CheckDocker{}
	cd.WarnMetaSpace = 100 // defaults
	cd.CritMetaSpace = 100
	cd.WarnDataSpace = 100
	cd.CritDataSpace = 100

	cd.dockerclient, err = dockerlib.NewClient(endpoint)
	if err != nil {
		return nil, err
	}

	return cd, err
}

type CheckDocker struct {
	WarnMetaSpace        float64
	CritMetaSpace        float64
	WarnDataSpace        float64
	CritDataSpace        float64
	ImageId              string
	dockerclient         *dockerlib.Client
	dockerInfoData       *dockerlib.Env
	dockerContainersData []dockerlib.APIContainers
}

func (cd *CheckDocker) GetData() error {
	errChan := make(chan error)
	var err error
	var wg sync.WaitGroup

	wg.Add(2)

	go func(cd *CheckDocker, errChan chan error) {
		defer wg.Done()

		cd.dockerInfoData, err = cd.dockerclient.Info()
		if err != nil {
			errChan <- err
		}
	}(cd, errChan)

	go func(cd *CheckDocker, errChan chan error) {
		defer wg.Done()

		cd.dockerContainersData, err = cd.dockerclient.ListContainers(dockerlib.ListContainersOptions{})
		if err != nil {
			errChan <- err
		}
	}(cd, errChan)

	go func() {
		wg.Wait()
		close(errChan)
	}()

	err = <-errChan

	return err
}

func (cd *CheckDocker) getByteSizeDriverStatus(key string) (bytesize.ByteSize, error) {
	var statusInArray [][]string

	err := json.Unmarshal([]byte(cd.dockerInfoData.Get("DriverStatus")), &statusInArray)

	if err != nil {
		return -1, errors.New("Unable to extract DriverStatus info.")
	}

	for _, status := range statusInArray {
		if status[0] == key {
			return bytesize.Parse([]byte(status[1]))
		}
	}

	return -1, errors.New(fmt.Sprintf("DriverStatus does not contain \"%v\"", key))
}

func (cd *CheckDocker) GetDataSpaceUsed() (bytesize.ByteSize, error) {
	return cd.getByteSizeDriverStatus("Data Space Used")
}

func (cd *CheckDocker) GetDataSpaceTotal() (bytesize.ByteSize, error) {
	return cd.getByteSizeDriverStatus("Data Space Total")
}

func (cd *CheckDocker) GetMetaSpaceUsed() (bytesize.ByteSize, error) {
	return cd.getByteSizeDriverStatus("Metadata Space Used")
}

func (cd *CheckDocker) GetMetaSpaceTotal() (bytesize.ByteSize, error) {
	return cd.getByteSizeDriverStatus("Metadata Space Total")
}

func (cd *CheckDocker) IsContainerRunning(imageId string) (dockerlib.APIContainers, bool) {
	for _, container := range cd.dockerContainersData {
		if strings.HasPrefix(container.Image, imageId) && strings.HasPrefix(container.Status, "Up") {
			return container, true
		}
	}
	return dockerlib.APIContainers{}, false
}

func (cd *CheckDocker) IsContainerAGhost(imageId string) (dockerlib.APIContainers, bool) {
	for _, container := range cd.dockerContainersData {
		if strings.HasPrefix(container.Image, imageId) && strings.Contains(container.Status, "Ghost") {
			return container, true
		}
	}
	return dockerlib.APIContainers{}, false
}

func (cd *CheckDocker) CheckMetaSpace(warnThreshold, criticalThreshold float64) *nagios.NagiosStatus {
	usedByteSize, err := cd.GetMetaSpaceUsed()
	if err != nil {
		return &nagios.NagiosStatus{err.Error(), nagios.NAGIOS_CRITICAL}
	}

	totalByteSize, err := cd.GetMetaSpaceTotal()
	if err != nil {
		return &nagios.NagiosStatus{err.Error(), nagios.NAGIOS_CRITICAL}
	}

	percentUsed := float64(usedByteSize/totalByteSize) * 100

	status := &nagios.NagiosStatus{fmt.Sprintf("Metadata Space Usage: %f", percentUsed) + "%", nagios.NAGIOS_OK}

	if percentUsed >= warnThreshold {
		status.Value = nagios.NAGIOS_WARNING
	}
	if percentUsed >= criticalThreshold {
		status.Value = nagios.NAGIOS_CRITICAL
	}

	return status
}

func (cd *CheckDocker) CheckDataSpace(warnThreshold, criticalThreshold float64) *nagios.NagiosStatus {
	usedByteSize, err := cd.GetDataSpaceUsed()
	if err != nil {
		return &nagios.NagiosStatus{err.Error(), nagios.NAGIOS_CRITICAL}
	}

	totalByteSize, err := cd.GetDataSpaceTotal()
	if err != nil {
		return &nagios.NagiosStatus{err.Error(), nagios.NAGIOS_CRITICAL}
	}

	percentUsed := float64(usedByteSize/totalByteSize) * 100

	status := &nagios.NagiosStatus{fmt.Sprintf("Data Space Usage: %f", percentUsed) + "%", nagios.NAGIOS_OK}

	if percentUsed >= warnThreshold {
		status.Value = nagios.NAGIOS_WARNING
	}
	if percentUsed >= criticalThreshold {
		status.Value = nagios.NAGIOS_CRITICAL
	}

	return status
}

func (cd *CheckDocker) CheckImageContainerIsInGoodShape(imageId string) *nagios.NagiosStatus {
	containerRunning, isRunning := cd.IsContainerRunning(imageId)
	containerGhost, isGhost := cd.IsContainerAGhost(imageId)

	if !isRunning {
		return &nagios.NagiosStatus{fmt.Sprintf("Container of image: %v is not running.", imageId), nagios.NAGIOS_CRITICAL}
	}
	if isGhost {
		return &nagios.NagiosStatus{fmt.Sprintf("Container(ID: %v) of image: %v is in ghost state.", containerGhost.ID, imageId), nagios.NAGIOS_CRITICAL}
	}

	return &nagios.NagiosStatus{fmt.Sprintf("Container(ID: %v) of image: %v is in top shape.", containerRunning.ID, imageId), nagios.NAGIOS_OK}
}

func main() {
	dockerEndpoint := *flag.String("base-url", "http://localhost:2375", "The Base URL for the Docker server")

	cd, err := NewCheckDocker(dockerEndpoint)
	if err != nil {
		nagios.Critical(err)
	}

	flag.Float64Var(&cd.WarnMetaSpace, "warn-meta-space", 100, "Warning threshold for Metadata Space")
	flag.Float64Var(&cd.CritMetaSpace, "crit-meta-space", 100, "Critical threshold for Metadata Space")
	flag.Float64Var(&cd.WarnDataSpace, "warn-data-space", 100, "Warning threshold for Data Space")
	flag.Float64Var(&cd.CritDataSpace, "crit-data-space", 100, "Critical threshold for Data Space")
	flag.StringVar(&cd.ImageId, "image-id", "", "An image ID that must be running on the Docker server")

	flag.Parse()

	err = cd.GetData()
	if err != nil {
		nagios.Critical(err)
	}

	baseStatus := &nagios.NagiosStatus{fmt.Sprintf("Total Containers: %v", len(cd.dockerContainersData)), nagios.NAGIOS_OK}

	statuses := make([]*nagios.NagiosStatus, 0)
	statuses = append(statuses, cd.CheckMetaSpace(cd.WarnMetaSpace, cd.CritMetaSpace))
	statuses = append(statuses, cd.CheckDataSpace(cd.WarnDataSpace, cd.CritDataSpace))

	if cd.ImageId != "" {
		statuses = append(statuses, cd.CheckImageContainerIsInGoodShape(cd.ImageId))
	}

	baseStatus.Aggregate(statuses)
	nagios.ExitWithStatus(baseStatus)
}
<commit_msg>only devicemapper contains metadata and data space info.<commit_after>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	dockerlib "github.com/fsouza/go-dockerclient"
	"github.com/newrelic/go_nagios"
	"github.com/shenwei356/util/bytesize"
	"strings"
	"sync"
)

func NewCheckDocker(endpoint string) (*CheckDocker, error) {
	var err error
	cd := &CheckDocker{}
	cd.WarnMetaSpace = 100 // defaults
	cd.CritMetaSpace = 100
	cd.WarnDataSpace = 100
	cd.CritDataSpace = 100

	cd.dockerclient, err = dockerlib.NewClient(endpoint)
	if err != nil {
		return nil, err
	}

	return cd, err
}

type CheckDocker struct {
	WarnMetaSpace        float64
	CritMetaSpace        float64
	WarnDataSpace        float64
	CritDataSpace        float64
	ImageId              string
	dockerclient         *dockerlib.Client
	dockerInfoData       *dockerlib.Env
	dockerContainersData []dockerlib.APIContainers
}

func (cd *CheckDocker) GetData() error {
	errChan := make(chan error)
	var err error
	var wg sync.WaitGroup

	wg.Add(2)

	go func(cd *CheckDocker, errChan chan error) {
		defer wg.Done()

		cd.dockerInfoData, err = cd.dockerclient.Info()
		if err != nil {
			errChan <- err
		}
	}(cd, errChan)

	go func(cd *CheckDocker, errChan chan error) {
		defer wg.Done()

		cd.dockerContainersData, err = cd.dockerclient.ListContainers(dockerlib.ListContainersOptions{})
		if err != nil {
			errChan <- err
		}
	}(cd, errChan)

	go func() {
		wg.Wait()
		close(errChan)
	}()

	err = <-errChan

	return err
}

func (cd *CheckDocker) getByteSizeDriverStatus(key string) (bytesize.ByteSize, error) {
	var statusInArray [][]string

	err := json.Unmarshal([]byte(cd.dockerInfoData.Get("DriverStatus")), &statusInArray)

	if err != nil {
		return -1, errors.New("Unable to extract DriverStatus info.")
	}

	for _, status := range statusInArray {
		if status[0] == key {
			return bytesize.Parse([]byte(status[1]))
		}
	}

	return -1, errors.New(fmt.Sprintf("DriverStatus does not contain \"%v\"", key))
}

func (cd *CheckDocker) GetDataSpaceUsed() (bytesize.ByteSize, error) {
	return cd.getByteSizeDriverStatus("Data Space Used")
}

func (cd *CheckDocker) GetDataSpaceTotal() (bytesize.ByteSize, error) {
	return cd.getByteSizeDriverStatus("Data Space Total")
}

func (cd *CheckDocker) GetMetaSpaceUsed() (bytesize.ByteSize, error) {
	return cd.getByteSizeDriverStatus("Metadata Space Used")
}

func (cd *CheckDocker) GetMetaSpaceTotal() (bytesize.ByteSize, error) {
	return cd.getByteSizeDriverStatus("Metadata Space Total")
}

func (cd *CheckDocker) IsContainerRunning(imageId string) (dockerlib.APIContainers, bool) {
	for _, container := range cd.dockerContainersData {
		if strings.HasPrefix(container.Image, imageId) && strings.HasPrefix(container.Status, "Up") {
			return container, true
		}
	}
	return dockerlib.APIContainers{}, false
}

func (cd *CheckDocker) IsContainerAGhost(imageId string) (dockerlib.APIContainers, bool) {
	for _, container := range cd.dockerContainersData {
		if strings.HasPrefix(container.Image, imageId) && strings.Contains(container.Status, "Ghost") {
			return container, true
		}
	}
	return dockerlib.APIContainers{}, false
}

func (cd *CheckDocker) CheckMetaSpace(warnThreshold, criticalThreshold float64) *nagios.NagiosStatus {
	usedByteSize, err := cd.GetMetaSpaceUsed()
	if err != nil {
		return &nagios.NagiosStatus{err.Error(), nagios.NAGIOS_CRITICAL}
	}

	totalByteSize, err := cd.GetMetaSpaceTotal()
	if err != nil {
		return &nagios.NagiosStatus{err.Error(), nagios.NAGIOS_CRITICAL}
	}

	percentUsed := float64(usedByteSize/totalByteSize) * 100

	status := &nagios.NagiosStatus{fmt.Sprintf("Metadata Space Usage: %f", percentUsed) + "%", nagios.NAGIOS_OK}

	if percentUsed >= warnThreshold {
		status.Value = nagios.NAGIOS_WARNING
	}
	if percentUsed >= criticalThreshold {
		status.Value = nagios.NAGIOS_CRITICAL
	}

	return status
}

func (cd *CheckDocker) CheckDataSpace(warnThreshold, criticalThreshold float64) *nagios.NagiosStatus {
	usedByteSize, err := cd.GetDataSpaceUsed()
	if err != nil {
		return &nagios.NagiosStatus{err.Error(), nagios.NAGIOS_CRITICAL}
	}

	totalByteSize, err := cd.GetDataSpaceTotal()
	if err != nil {
		return &nagios.NagiosStatus{err.Error(), nagios.NAGIOS_CRITICAL}
	}

	percentUsed := float64(usedByteSize/totalByteSize) * 100

	status := &nagios.NagiosStatus{fmt.Sprintf("Data Space Usage: %f", percentUsed) + "%", nagios.NAGIOS_OK}

	if percentUsed >= warnThreshold {
		status.Value = nagios.NAGIOS_WARNING
	}
	if percentUsed >= criticalThreshold {
		status.Value = nagios.NAGIOS_CRITICAL
	}

	return status
}

func (cd *CheckDocker) CheckImageContainerIsInGoodShape(imageId string) *nagios.NagiosStatus {
	containerRunning, isRunning := cd.IsContainerRunning(imageId)
	containerGhost, isGhost := cd.IsContainerAGhost(imageId)

	if !isRunning {
		return &nagios.NagiosStatus{fmt.Sprintf("Container of image: %v is not running.", imageId), nagios.NAGIOS_CRITICAL}
	}
	if isGhost {
		return &nagios.NagiosStatus{fmt.Sprintf("Container(ID: %v) of image: %v is in ghost state.", containerGhost.ID, imageId), nagios.NAGIOS_CRITICAL}
	}

	return &nagios.NagiosStatus{fmt.Sprintf("Container(ID: %v) of image: %v is in top shape.", containerRunning.ID, imageId), nagios.NAGIOS_OK}
}

func main() {
	dockerEndpoint := *flag.String("base-url", "http://localhost:2375", "The Base URL for the Docker server")

	cd, err := NewCheckDocker(dockerEndpoint)
	if err != nil {
		nagios.Critical(err)
	}

	flag.Float64Var(&cd.WarnMetaSpace, "warn-meta-space", 100, "Warning threshold for Metadata Space")
	flag.Float64Var(&cd.CritMetaSpace, "crit-meta-space", 100, "Critical threshold for Metadata Space")
	flag.Float64Var(&cd.WarnDataSpace, "warn-data-space", 100, "Warning threshold for Data Space")
	flag.Float64Var(&cd.CritDataSpace, "crit-data-space", 100, "Critical threshold for Data Space")
	flag.StringVar(&cd.ImageId, "image-id", "", "An image ID that must be running on the Docker server")

	flag.Parse()

	err = cd.GetData()
	if err != nil {
		nagios.Critical(err)
	}

	baseStatus := &nagios.NagiosStatus{fmt.Sprintf("Total Containers: %v", len(cd.dockerContainersData)), nagios.NAGIOS_OK}

	statuses := make([]*nagios.NagiosStatus, 0)

	driver := cd.dockerInfoData.Get("Driver")

	// Unfortunately, Metadata Space and Data Space information is only available on devicemapper
	if driver == "devicemapper" {
		statuses = append(statuses, cd.CheckMetaSpace(cd.WarnMetaSpace, cd.CritMetaSpace))
		statuses = append(statuses, cd.CheckDataSpace(cd.WarnDataSpace, cd.CritDataSpace))
	}

	if cd.ImageId != "" {
		statuses = append(statuses, cd.CheckImageContainerIsInGoodShape(cd.ImageId))
	}

	baseStatus.Aggregate(statuses)
	nagios.ExitWithStatus(baseStatus)
}
<|endoftext|>


<commit_before>package test

import (
	"io/ioutil"
	"os"
	"path/filepath"
)

import (
	tp "tritium/proto"
	yaml "goyaml"
)

type TritiumTest_Hash struct {
	Key              string
	Value            string
}

type TritiumTest struct {
	Script           string
	Input            string
	Output           string
	EnvProto         []*TritiumTest_Hash
	ExportsProto     []*TritiumTest_Hash
	Logs             []string
	Transformer      *tp.Transform
}

func (test *TritiumTest) Env() (env map[string]string) {
	env = map[string]string{}
	for _, hash := range test.EnvProto {
		env[hash.Key] = hash.Value
	}
	return
}

func (test *TritiumTest) Exports() (exports [][]string) {
	exports = make([][]string, len(test.ExportsProto))
	for n, hash := range test.ExportsProto {
		exports[n] = []string{hash.Key, hash.Value}
	}
	return
}

func (test *TritiumTest) SetEnv(env map[string]string) {
	th := make([]*TritiumTest_Hash, len(env))
	i := 0
	for key, value := range env {
		th[i] = &TritiumTest_Hash{Key: key, Value: value}
		i++
	}
	test.EnvProto = th
}

func (test *TritiumTest) SetExports(exports [][]string) {
	th := make([]*TritiumTest_Hash, len(exports))
	for n, export := range exports {
		th[n] = &TritiumTest_Hash{Key: export[0], Value: export[1]}
	}
	test.ExportsProto = th
}

func NewTritiumTestFromFolder(path string) (test *TritiumTest, err error) {
	script, err := ioutil.ReadFile(filepath.Join(path, "input.ts"))
	if err != nil {
		script = []byte("")
	}

	input, err := ioutil.ReadFile(filepath.Join(path, "input.http"))
	if err != nil {
		input = []byte("")
	}

	output, err := ioutil.ReadFile(filepath.Join(path, "output.http"))
	if err != nil {
		output = []byte("")
	}

	env := make(map[string]string, 0)
	data, err := ioutil.ReadFile(filepath.Join(path, "vars.yml"))
	if err == nil {
		err = yaml.Unmarshal(data, &env)
		if err != nil {
			return
		}
	}

	exports := make([][]string, 0)
	data, err = ioutil.ReadFile(filepath.Join(path, "exports.yml"))
	if err == nil {
		err = yaml.Unmarshal(data, &exports)
		if err != nil {
			return
		}
	}

	test = &TritiumTest{
		Script: string(script),
		Input:  string(input),
		Output: string(output),
		Logs:   []string{},
	}

	test.SetEnv(env)
	test.SetExports(exports)

	return test, nil
}

func (test *TritiumTest) WriteFolder(path string) (err error) {
	err = os.MkdirAll(path, 0755)
	if err != nil {
		return
	}

	input_script := filepath.Join(path, "input.ts")
	err = ioutil.WriteFile(input_script, []byte(test.Script), 0644)
	if err != nil {
		return
	}

	input_file := filepath.Join(path, "input.txt")
	err = ioutil.WriteFile(input_file, []byte(test.Input), 0644)
	if err != nil {
		return
	}

	output_file := filepath.Join(path, "output.txt")
	err = ioutil.WriteFile(output_file, []byte(test.Output), 0644)
	if err != nil {
		return
	}

	return
}
<commit_msg>add rewrite rules to TritiumTest<commit_after>package test

import (
	"io/ioutil"
	"os"
	"path/filepath"
)

import (
	tp "tritium/proto"
	yaml "goyaml"
)

type TritiumTest_Hash struct {
	Key              string
	Value            string
}

type TritiumTest struct {
	Script           string
	Input            string
	Output           string
	EnvProto         []*TritiumTest_Hash
	ExportsProto     []*TritiumTest_Hash
	Logs             []string
	Transformer      *tp.Transform
	Rrules           []*tp.RewriteRule
}

func (test *TritiumTest) Env() (env map[string]string) {
	env = map[string]string{}
	for _, hash := range test.EnvProto {
		env[hash.Key] = hash.Value
	}
	return
}

func (test *TritiumTest) Exports() (exports [][]string) {
	exports = make([][]string, len(test.ExportsProto))
	for n, hash := range test.ExportsProto {
		exports[n] = []string{hash.Key, hash.Value}
	}
	return
}

func (test *TritiumTest) SetEnv(env map[string]string) {
	th := make([]*TritiumTest_Hash, len(env))
	i := 0
	for key, value := range env {
		th[i] = &TritiumTest_Hash{Key: key, Value: value}
		i++
	}
	test.EnvProto = th
}

func (test *TritiumTest) SetExports(exports [][]string) {
	th := make([]*TritiumTest_Hash, len(exports))
	for n, export := range exports {
		th[n] = &TritiumTest_Hash{Key: export[0], Value: export[1]}
	}
	test.ExportsProto = th
}

func NewTritiumTestFromFolder(path string) (test *TritiumTest, err error) {
	script, err := ioutil.ReadFile(filepath.Join(path, "input.ts"))
	if err != nil {
		script = []byte("")
	}

	input, err := ioutil.ReadFile(filepath.Join(path, "input.http"))
	if err != nil {
		input = []byte("")
	}

	output, err := ioutil.ReadFile(filepath.Join(path, "output.http"))
	if err != nil {
		output = []byte("")
	}

	env := make(map[string]string, 0)
	data, err := ioutil.ReadFile(filepath.Join(path, "vars.yml"))
	if err == nil {
		err = yaml.Unmarshal(data, &env)
		if err != nil {
			return
		}
	}

	exports := make([][]string, 0)
	data, err = ioutil.ReadFile(filepath.Join(path, "exports.yml"))
	if err == nil {
		err = yaml.Unmarshal(data, &exports)
		if err != nil {
			return
		}
	}

	test = &TritiumTest{
		Script: string(script),
		Input:  string(input),
		Output: string(output),
		Logs:   []string{},
	}

	test.SetEnv(env)
	test.SetExports(exports)

	return test, nil
}

func (test *TritiumTest) WriteFolder(path string) (err error) {
	err = os.MkdirAll(path, 0755)
	if err != nil {
		return
	}

	input_script := filepath.Join(path, "input.ts")
	err = ioutil.WriteFile(input_script, []byte(test.Script), 0644)
	if err != nil {
		return
	}

	input_file := filepath.Join(path, "input.txt")
	err = ioutil.WriteFile(input_file, []byte(test.Input), 0644)
	if err != nil {
		return
	}

	output_file := filepath.Join(path, "output.txt")
	err = ioutil.WriteFile(output_file, []byte(test.Output), 0644)
	if err != nil {
		return
	}

	return
}
<|endoftext|>


<commit_before>package keycloak

func (as *ActionSuite) Test_GetToken() {
	err := GetToken()
	as.NoError(err)
}
<commit_msg>Refactor keycloak test.<commit_after>package keycloak

import (
	"testing"
)

func TestGetToken(t *testing.T) {
	err := GetToken()
	if err != nil {
		t.Errorf("Failed test because: %s", err.Error())
	}
}
<|endoftext|>


<commit_before>package wsutil

import (
	"bytes"
	"runtime"
	"testing"

	"github.com/gobwas/ws"
)

func TestControlHandler(t *testing.T) {
	for _, test := range []struct {
		name  string
		state ws.State
		in    ws.Frame
		out   ws.Frame
		noOut bool
		err   error
	}{
		{
			name: "ping",
			in:   ws.NewPingFrame(nil),
			out:  ws.NewPongFrame(nil),
		},
		{
			name: "ping",
			in:   ws.NewPingFrame([]byte("catch the ball")),
			out:  ws.NewPongFrame([]byte("catch the ball")),
		},
		{
			name:  "ping",
			state: ws.StateServerSide,
			in:    ws.MaskFrame(ws.NewPingFrame([]byte("catch the ball"))),
			out:   ws.NewPongFrame([]byte("catch the ball")),
		},
		{
			name:  "pong",
			in:    ws.NewPongFrame(nil),
			noOut: true,
		},
		{
			name:  "pong",
			in:    ws.NewPongFrame([]byte("catched")),
			noOut: true,
		},
		{
			name: "close",
			in:   ws.NewCloseFrame(nil),
			out:  ws.NewCloseFrame(nil),
			err: ClosedError{
				Code: ws.StatusNoStatusRcvd,
			},
		},
		{
			name: "close",
			in: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "goodbye!",
			)),
			out: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "",
			)),
			err: ClosedError{
				Code:   ws.StatusGoingAway,
				Reason: "goodbye!",
			},
		},
		{
			name:  "close",
			state: ws.StateServerSide,
			in: ws.MaskFrame(ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "goodbye!",
			))),
			out: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "",
			)),
			err: ClosedError{
				Code:   ws.StatusGoingAway,
				Reason: "goodbye!",
			},
		},
		{
			name: "close",
			in: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusNormalClosure, string([]byte{0, 200}),
			)),
			out: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusProtocolError, ws.ErrProtocolInvalidUTF8.Error(),
			)),
			err: ws.ErrProtocolInvalidUTF8,
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			defer func() {
				if err := recover(); err != nil {
					stack := make([]byte, 4096)
					n := runtime.Stack(stack, true)
					t.Fatalf(
						"panic recovered: %v\n%s",
						err, stack[:n],
					)
				}
			}()
			var (
				out = bytes.NewBuffer(nil)
				in  = bytes.NewReader(test.in.Payload)
			)
			c := ControlHandler{
				Src:   in,
				Dst:   out,
				State: test.state,
			}

			err := c.Handle(test.in.Header)
			if err != test.err {
				t.Errorf("unexpected error: %v; want %v", err, test.err)
			}

			if in.Len() != 0 {
				t.Errorf("handler did not drained the input")
			}

			act := out.Bytes()
			switch {
			case len(act) == 0 && test.noOut:
				return
			case len(act) == 0 && !test.noOut:
				t.Errorf("unexpected silence")
			case len(act) > 0 && test.noOut:
				t.Errorf("unexpected sent frame")
			default:
				exp := ws.MustCompileFrame(test.out)
				if !bytes.Equal(act, exp) {
					fa := ws.MustReadFrame(bytes.NewReader(act))
					fe := ws.MustReadFrame(bytes.NewReader(exp))
					t.Errorf(
						"unexpected sent frame:\n\tact: %+v\n\texp: %+v\nbytes:\n\tact: %v\n\texp: %v",
						fa, fe, act, exp,
					)
				}
			}
		})
	}
}
<commit_msg>testing: to fail "set the reserved bytes for the header to the max value"<commit_after>package wsutil

import (
	"bytes"
	"runtime"
	"testing"

	"github.com/gobwas/ws"
)

func TestControlHandler(t *testing.T) {
	for _, test := range []struct {
		name  string
		state ws.State
		in    ws.Frame
		out   ws.Frame
		noOut bool
		err   error
	}{
		{
			name: "ping",
			in:   ws.NewPingFrame(nil),
			out:  ws.NewPongFrame(nil),
		},
		{
			name: "ping",
			in:   ws.NewPingFrame([]byte("catch the ball")),
			out:  ws.NewPongFrame([]byte("catch the ball")),
		},
		{
			name:  "ping",
			state: ws.StateServerSide,
			in:    ws.MaskFrame(ws.NewPingFrame([]byte("catch the ball"))),
			out:   ws.NewPongFrame([]byte("catch the ball")),
		},
		{
			name: "ping",
			in:   ws.NewPingFrame(bytes.Repeat([]byte{0xfe}, 125)),
			out:  ws.NewPongFrame(bytes.Repeat([]byte{0xfe}, 125)),
		},
		{
			name:  "pong",
			in:    ws.NewPongFrame(nil),
			noOut: true,
		},
		{
			name:  "pong",
			in:    ws.NewPongFrame([]byte("catched")),
			noOut: true,
		},
		{
			name: "close",
			in:   ws.NewCloseFrame(nil),
			out:  ws.NewCloseFrame(nil),
			err: ClosedError{
				Code: ws.StatusNoStatusRcvd,
			},
		},
		{
			name: "close",
			in: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "goodbye!",
			)),
			out: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "",
			)),
			err: ClosedError{
				Code:   ws.StatusGoingAway,
				Reason: "goodbye!",
			},
		},
		{
			name: "close",
			in: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "bye",
			)),
			out: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "",
			)),
			err: ClosedError{
				Code:   ws.StatusGoingAway,
				Reason: "bye",
			},
		},
		{
			name:  "close",
			state: ws.StateServerSide,
			in: ws.MaskFrame(ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "goodbye!",
			))),
			out: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "",
			)),
			err: ClosedError{
				Code:   ws.StatusGoingAway,
				Reason: "goodbye!",
			},
		},
		{
			name: "close",
			in: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusNormalClosure, string([]byte{0, 200}),
			)),
			out: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusProtocolError, ws.ErrProtocolInvalidUTF8.Error(),
			)),
			err: ws.ErrProtocolInvalidUTF8,
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			defer func() {
				if err := recover(); err != nil {
					stack := make([]byte, 4096)
					n := runtime.Stack(stack, true)
					t.Fatalf(
						"panic recovered: %v\n%s",
						err, stack[:n],
					)
				}
			}()
			var (
				out = bytes.NewBuffer(nil)
				in  = bytes.NewReader(test.in.Payload)
			)
			c := ControlHandler{
				Src:   in,
				Dst:   out,
				State: test.state,
			}

			err := c.Handle(test.in.Header)
			if err != test.err {
				t.Errorf("unexpected error: %v; want %v", err, test.err)
			}

			if in.Len() != 0 {
				t.Errorf("handler did not drained the input")
			}

			act := out.Bytes()
			switch {
			case len(act) == 0 && test.noOut:
				return
			case len(act) == 0 && !test.noOut:
				t.Errorf("unexpected silence")
			case len(act) > 0 && test.noOut:
				t.Errorf("unexpected sent frame")
			default:
				exp := ws.MustCompileFrame(test.out)
				if !bytes.Equal(act, exp) {
					fa := ws.MustReadFrame(bytes.NewReader(act))
					fe := ws.MustReadFrame(bytes.NewReader(exp))
					t.Errorf(
						"unexpected sent frame:\n\tact: %+v\n\texp: %+v\nbytes:\n\tact: %v\n\texp: %v",
						fa, fe, act, exp,
					)
				}
			}
		})
	}
}
<|endoftext|>


<commit_before>package api

import (
	"koding/db/mongodb/modelhelper"
	"socialapi/config"
	"socialapi/workers/common/handler"
	"socialapi/workers/common/mux"

	"github.com/RangelReale/osin"
	"gopkg.in/mgo.v2"
)

const (
	OauthAuthorization = "oauth-authorization"
	GenerateToken      = "generate-token"
)

// AddHandlers add oauth for koding api
func AddHandlers(m *mux.Mux, config *config.Config) {
	session := modelhelper.Mongo.Session

	oauth := NewOAuthHandler(session)
	m.AddUnscopedHandler(
		handler.Request{
			Handler:  oauth.AuthorizeClient,
			Name:     OauthAuthorization,
			Type:     handler.GetRequest,
			Endpoint: "/oauth/authorize",
		},
	)
	m.AddUnscopedHandler(
		handler.Request{
			Handler:  oauth.GenerateToken,
			Name:     GenerateToken,
			Type:     handler.GetRequest,
			Endpoint: "/oauth/token",
		},
	)

}

func NewOAuthHandler(session *mgo.Session) *Oauth {
	sconfig := osin.NewServerConfig()
	sconfig.AllowedAuthorizeTypes = osin.AllowedAuthorizeType{osin.CODE, osin.TOKEN}
	sconfig.AllowedAccessTypes = osin.AllowedAccessType{osin.AUTHORIZATION_CODE,
		osin.REFRESH_TOKEN, osin.PASSWORD, osin.CLIENT_CREDENTIALS, osin.ASSERTION}
	sconfig.AllowGetAccessRequest = true
	storage := modelhelper.NewOauthStore(session)
	server := osin.NewServer(sconfig, storage)

	return &Oauth{
		sconfig: sconfig,
		server:  server,
		Storage: storage,
	}
}
<commit_msg>go/oauth: add comment for constants<commit_after>package api

import (
	"koding/db/mongodb/modelhelper"
	"socialapi/config"
	"socialapi/workers/common/handler"
	"socialapi/workers/common/mux"

	"github.com/RangelReale/osin"
	"gopkg.in/mgo.v2"
)

const (
	OauthAuthorization = "oauth-authorization"
	GenerateToken      = "generate-token"
)

// AddHandlers add oauth for koding api
func AddHandlers(m *mux.Mux, config *config.Config) {
	session := modelhelper.Mongo.Session

	oauth := NewOAuthHandler(session)
	m.AddUnscopedHandler(
		handler.Request{
			Handler:  oauth.AuthorizeClient,
			Name:     OauthAuthorization,
			Type:     handler.GetRequest,
			Endpoint: "/oauth/authorize",
		},
	)
	m.AddUnscopedHandler(
		handler.Request{
			Handler:  oauth.GenerateToken,
			Name:     GenerateToken,
			Type:     handler.GetRequest,
			Endpoint: "/oauth/token",
		},
	)

}

func NewOAuthHandler(session *mgo.Session) *Oauth {
	sconfig := osin.NewServerConfig()
	// AllowedAccessType is a collection of allowed access request types
	sconfig.AllowedAuthorizeTypes = osin.AllowedAuthorizeType{osin.CODE, osin.TOKEN}
	// AccessRequestType is the type for OAuth param `grant_type`
	sconfig.AllowedAccessTypes = osin.AllowedAccessType{osin.AUTHORIZATION_CODE,
		osin.REFRESH_TOKEN, osin.PASSWORD, osin.CLIENT_CREDENTIALS, osin.ASSERTION}
	// If true allows access request using GET, else only POST - default false
	sconfig.AllowGetAccessRequest = true
	storage := modelhelper.NewOauthStore(session)
	server := osin.NewServer(sconfig, storage)

	return &Oauth{
		sconfig: sconfig,
		server:  server,
		Storage: storage,
	}
}
<|endoftext|>


<commit_before>// Package limitedstore implements a simple wrapper around existing
// datastores (tusd.DataStore) while limiting the used storage size.
// It will start terminating existing uploads if not enough space is left in
// order to create a new upload.
// The order in which the uploads will be terminated is defined by their size,
// whereas the biggest ones are deleted first.
// This package's functionality is very limited and naive. It will terminate
// uploads whether they are finished yet or not. Only one datastore is allowed to
// access the underlying storage else the limited store will not function
// properly. Two tusd.FileStore instances using the same directory, for example.
// In addition the limited store will keep a list of the uploads' IDs in memory
// which may create a growing memory leak.
package limitedstore

import (
	"github.com/tus/tusd"
	"sync"
  "sort"
)

type LimitedStore struct {
	StoreSize int64
	tusd.DataStore

	uploads  map[string]int64
	usedSize int64

	mutex *sync.Mutex
}

// pair structure to perform map-sorting
type pair struct {
  key string
  value int64
}

type pairlist []pair

func (p pairlist) Len() int           { return len(p) }
func (p pairlist) Swap(i, j int)       { p[i], p[j] = p[j], p[i] }
func (p pairlist) Less(i, j int) bool  { return p[i].value < p[j].value }

// Create a new limited store with the given size as the maximum storage size
func New(storeSize int64, dataStore tusd.DataStore) *LimitedStore {
	return &LimitedStore{
		StoreSize: storeSize,
		DataStore: dataStore,
		uploads:   make(map[string]int64),
		mutex:     new(sync.Mutex),
	}
}

func (store *LimitedStore) NewUpload(info tusd.FileInfo) (string, error) {
	store.mutex.Lock()
	defer store.mutex.Unlock()

	if err := store.ensureSpace(info.Size); err != nil {
		return "", err
	}

	id, err := store.DataStore.NewUpload(info)
	if err != nil {
		return "", err
	}

	store.usedSize += info.Size
	store.uploads[id] = info.Size

	return id, nil
}

func (store *LimitedStore) Terminate(id string) error {
	store.mutex.Lock()
	defer store.mutex.Unlock()

	return store.terminate(id)
}

func (store *LimitedStore) terminate(id string) error {
	err := store.DataStore.Terminate(id)
	if err != nil {
		return err
	}

	size := store.uploads[id]
	delete(store.uploads, id)
	store.usedSize -= size

	return nil
}

// Ensure enough space is available to store an upload of the specified size.
// It will terminate uploads until enough space is freed.
func (store *LimitedStore) ensureSpace(size int64) error {
	if (store.usedSize + size) <= store.StoreSize {
		// Enough space is available to store the new upload
		return nil
	}
  sorted_uploads := make(pairlist, len(store.uploads))
  i := 0
  for u,h := range store.uploads {
    sorted_uploads[i] = pair{u, h}
    i++
  }
  sort.Sort(sort.Reverse(sorted_uploads))

  // Forward traversal through the
  // uploads in terms of size, biggest upload first
  for _,k := range sorted_uploads {
    id := k.key
		if err := store.terminate(id); err != nil {
			return err
		}

		if (store.usedSize + size) <= store.StoreSize {
			// Enough space has been freed to store the new upload
			return nil
		}
	}

	return nil
}
<commit_msg>Adjust code style<commit_after>// Package limitedstore implements a simple wrapper around existing
// datastores (tusd.DataStore) while limiting the used storage size.
// It will start terminating existing uploads if not enough space is left in
// order to create a new upload.
// The order in which the uploads will be terminated is defined by their size,
// whereas the biggest ones are deleted first.
// This package's functionality is very limited and naive. It will terminate
// uploads whether they are finished yet or not. Only one datastore is allowed to
// access the underlying storage else the limited store will not function
// properly. Two tusd.FileStore instances using the same directory, for example.
// In addition the limited store will keep a list of the uploads' IDs in memory
// which may create a growing memory leak.
package limitedstore

import (
	"github.com/tus/tusd"
	"sort"
	"sync"
)

type LimitedStore struct {
	StoreSize int64
	tusd.DataStore

	uploads  map[string]int64
	usedSize int64

	mutex *sync.Mutex
}

// pair structure to perform map-sorting
type pair struct {
	key   string
	value int64
}

type pairlist []pair

func (p pairlist) Len() int           { return len(p) }
func (p pairlist) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p pairlist) Less(i, j int) bool { return p[i].value < p[j].value }

// Create a new limited store with the given size as the maximum storage size
func New(storeSize int64, dataStore tusd.DataStore) *LimitedStore {
	return &LimitedStore{
		StoreSize: storeSize,
		DataStore: dataStore,
		uploads:   make(map[string]int64),
		mutex:     new(sync.Mutex),
	}
}

func (store *LimitedStore) NewUpload(info tusd.FileInfo) (string, error) {
	store.mutex.Lock()
	defer store.mutex.Unlock()

	if err := store.ensureSpace(info.Size); err != nil {
		return "", err
	}

	id, err := store.DataStore.NewUpload(info)
	if err != nil {
		return "", err
	}

	store.usedSize += info.Size
	store.uploads[id] = info.Size

	return id, nil
}

func (store *LimitedStore) Terminate(id string) error {
	store.mutex.Lock()
	defer store.mutex.Unlock()

	return store.terminate(id)
}

func (store *LimitedStore) terminate(id string) error {
	err := store.DataStore.Terminate(id)
	if err != nil {
		return err
	}

	size := store.uploads[id]
	delete(store.uploads, id)
	store.usedSize -= size

	return nil
}

// Ensure enough space is available to store an upload of the specified size.
// It will terminate uploads until enough space is freed.
func (store *LimitedStore) ensureSpace(size int64) error {
	if (store.usedSize + size) <= store.StoreSize {
		// Enough space is available to store the new upload
		return nil
	}

	sortedUploads := make(pairlist, len(store.uploads))
	i := 0
	for u, h := range store.uploads {
		sortedUploads[i] = pair{u, h}
		i++
	}
	sort.Sort(sort.Reverse(sortedUploads))

	// Forward traversal through the uploads in terms of size, biggest upload first
	for _, k := range sortedUploads {
		id := k.key

		if err := store.terminate(id); err != nil {
			return err
		}

		if (store.usedSize + size) <= store.StoreSize {
			// Enough space has been freed to store the new upload
			return nil
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"regexp"
	"time"

	"github.com/rcrowley/go-metrics"
	"github.com/skia-dev/glog"
	"github.com/skia-dev/influxdb/client"
	"go.skia.org/infra/alertserver/go/alerting"
	"go.skia.org/infra/go/buildbot"
)

/*
	This file contains goroutines which trigger more complex alerts than
	can be expressed using the rule format in alerts.cfg.
*/

func StartAlertRoutines(am *alerting.AlertManager, tickInterval time.Duration, c *client.Client) {
	// Disconnected buildslaves.
	go func() {
		seriesTmpl := "buildbot.buildslaves.%s.connected"
		re := regexp.MustCompile("[^A-Za-z0-9]+")
		emailAction, err := alerting.ParseAction("Email(infra-alerts@skia.org)")
		if err != nil {
			glog.Fatal(err)
		}
		actions := []alerting.Action{emailAction}
		for _ = range time.Tick(tickInterval) {
			glog.Info("Loading buildslave data.")
			slaves, err := buildbot.GetBuildSlaves()
			if err != nil {
				glog.Error(err)
				continue
			}
			for _, m := range slaves {
				for _, s := range m {
					v := int64(0)
					if s.Connected {
						v = int64(1)
					}
					metric := fmt.Sprintf(seriesTmpl, re.ReplaceAllString(s.Name, "_"))
					metrics.GetOrRegisterGauge(metric, metrics.DefaultRegistry).Update(v)
					if !s.Connected {
						// This buildslave is offline. Figure out which one it is.
						if err := am.AddAlert(&alerting.Alert{
							Name:        fmt.Sprintf("Buildslave %s offline", s.Name),
							Category:    "Infra",
							Message:     fmt.Sprintf("Buildslave %s is offline.", s.Name),
							Nag:         int64(time.Hour),
							AutoDismiss: int64(2 * tickInterval),
							Actions:     actions,
						}); err != nil {
							glog.Error(err)
						}
					}
				}
			}
		}
	}()
}
<commit_msg>More informative buildslave disconnect alerts<commit_after>package main

import (
	"fmt"
	"regexp"
	"time"

	"github.com/rcrowley/go-metrics"
	"github.com/skia-dev/glog"
	"github.com/skia-dev/influxdb/client"
	"go.skia.org/infra/alertserver/go/alerting"
	"go.skia.org/infra/go/buildbot"
)

/*
	This file contains goroutines which trigger more complex alerts than
	can be expressed using the rule format in alerts.cfg.
*/

func StartAlertRoutines(am *alerting.AlertManager, tickInterval time.Duration, c *client.Client) {
	// Disconnected buildslaves.
	go func() {
		seriesTmpl := "buildbot.buildslaves.%s.connected"
		re := regexp.MustCompile("[^A-Za-z0-9]+")
		emailAction, err := alerting.ParseAction("Email(infra-alerts@skia.org)")
		if err != nil {
			glog.Fatal(err)
		}
		actions := []alerting.Action{emailAction}
		for _ = range time.Tick(tickInterval) {
			glog.Info("Loading buildslave data.")
			slaves, err := buildbot.GetBuildSlaves()
			if err != nil {
				glog.Error(err)
				continue
			}
			for masterName, m := range slaves {
				for _, s := range m {
					v := int64(0)
					if s.Connected {
						v = int64(1)
					}
					metric := fmt.Sprintf(seriesTmpl, re.ReplaceAllString(s.Name, "_"))
					metrics.GetOrRegisterGauge(metric, metrics.DefaultRegistry).Update(v)
					if !s.Connected {
						// This buildslave is offline. Figure out which one it is.
						if err := am.AddAlert(&alerting.Alert{
							Name:        fmt.Sprintf("Buildslave %s offline", s.Name),
							Category:    "Infra",
							Message:     fmt.Sprintf("Buildslave %s is not connected to https://uberchromegw.corp.gogle.com/i/%s/buildslaves/%s.", s.Name, masterName, s.Name),
							Nag:         int64(time.Hour),
							AutoDismiss: int64(2 * tickInterval),
							Actions:     actions,
						}); err != nil {
							glog.Error(err)
						}
					}
				}
			}
		}
	}()
}
<|endoftext|>


<commit_before>package maas

import (
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/testing"
)

type ConfigSuite struct{}

var _ = Suite(new(ConfigSuite))

// copyAttrs copies values from src into dest.  If src contains a key that was
// already in dest, its value in dest will still be updated to the one from
// src.
func copyAttrs(src, dest map[string]interface{}) {
	for k, v := range src {
		dest[k] = v
	}
}

// newConfig creates a MAAS environment config from attributes.
func newConfig(values map[string]interface{}) (*maasEnvironConfig, error) {
	defaults := map[string]interface{}{
		"name":           "testenv",
		"type":           "maas",
		"ca-cert":        testing.CACert,
		"ca-private-key": testing.CAKey,
	}
	cfg := make(map[string]interface{})
	copyAttrs(defaults, cfg)
	copyAttrs(values, cfg)
	env, err := environs.NewFromAttrs(cfg)
	if err != nil {
		return nil, err
	}
	return env.(*maasEnviron).ecfg(), nil
}

func (ConfigSuite) TestParsesMAASSettings(c *C) {
	server := "maas.example.com"
	oauth := "consumer-key:resource-token:resource-secret"
	secret := "ssssssht"
	ecfg, err := newConfig(map[string]interface{}{
		"maas-server":  server,
		"maas-oauth":   oauth,
		"admin-secret": secret,
	})
	c.Assert(err, IsNil)
	c.Check(ecfg.MAASServer(), Equals, server)
	c.Check(ecfg.MAASOAuth(), DeepEquals, oauth)
	c.Check(ecfg.AdminSecret(), Equals, secret)
}

func (ConfigSuite) TestRequiresMaasServer(c *C) {
	oauth := "consumer-key:resource-token:resource-secret"
	_, err := newConfig(map[string]interface{}{
		"maas-oauth":   oauth,
		"admin-secret": "secret",
	})
	c.Check(err, NotNil)
}

func (ConfigSuite) TestRequiresOAuth(c *C) {
	_, err := newConfig(map[string]interface{}{
		"maas-server":  "maas.example.com",
		"admin-secret": "secret",
	})
	c.Check(err, NotNil)
}

func (ConfigSuite) TestChecksWellFormedOAuth(c *C) {
	_, err := newConfig(map[string]interface{}{
		"maas-server":  "maas.example.com",
		"maas-oauth":   "This should have been a 3-part token.",
		"admin-secret": "secret",
	})
	c.Check(err, NotNil)
}
<commit_msg>No need to test for required parameters; leave it to schema.  Check error message for malformed maas-oauth.<commit_after>package maas

import (
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/testing"
)

type ConfigSuite struct{}

var _ = Suite(new(ConfigSuite))

// copyAttrs copies values from src into dest.  If src contains a key that was
// already in dest, its value in dest will still be updated to the one from
// src.
func copyAttrs(src, dest map[string]interface{}) {
	for k, v := range src {
		dest[k] = v
	}
}

// newConfig creates a MAAS environment config from attributes.
func newConfig(values map[string]interface{}) (*maasEnvironConfig, error) {
	defaults := map[string]interface{}{
		"name":           "testenv",
		"type":           "maas",
		"ca-cert":        testing.CACert,
		"ca-private-key": testing.CAKey,
	}
	cfg := make(map[string]interface{})
	copyAttrs(defaults, cfg)
	copyAttrs(values, cfg)
	env, err := environs.NewFromAttrs(cfg)
	if err != nil {
		return nil, err
	}
	return env.(*maasEnviron).ecfg(), nil
}

func (ConfigSuite) TestParsesMAASSettings(c *C) {
	server := "maas.example.com"
	oauth := "consumer-key:resource-token:resource-secret"
	secret := "ssssssht"
	ecfg, err := newConfig(map[string]interface{}{
		"maas-server":  server,
		"maas-oauth":   oauth,
		"admin-secret": secret,
	})
	c.Assert(err, IsNil)
	c.Check(ecfg.MAASServer(), Equals, server)
	c.Check(ecfg.MAASOAuth(), DeepEquals, oauth)
	c.Check(ecfg.AdminSecret(), Equals, secret)
}

func (ConfigSuite) TestChecksWellFormedMaasOAuth(c *C) {
	_, err := newConfig(map[string]interface{}{
		"maas-server":  "maas.example.com",
		"maas-oauth":   "This should have been a 3-part token.",
		"admin-secret": "secret",
	})
	c.Assert(err, NotNil)
	c.Check(err, ErrorMatches, ".*malformed maas-oauth.*")
}
<|endoftext|>


<commit_before>// Copyright Project Harbor 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 adapter

import (
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"

	"github.com/docker/distribution"
	"github.com/docker/distribution/manifest/schema1"
	"github.com/goharbor/harbor/src/common/http/modifier"
	common_http_auth "github.com/goharbor/harbor/src/common/http/modifier/auth"
	"github.com/goharbor/harbor/src/common/utils/log"
	registry_pkg "github.com/goharbor/harbor/src/common/utils/registry"
	"github.com/goharbor/harbor/src/common/utils/registry/auth"
	"github.com/goharbor/harbor/src/replication/model"
	"github.com/goharbor/harbor/src/replication/util"
)

// const definition
const (
	UserAgentReplication = "harbor-replication-service"
)

// ImageRegistry defines the capabilities that an image registry should have
type ImageRegistry interface {
	FetchImages(namespaces []string, filters []*model.Filter) ([]*model.Resource, error)
	ManifestExist(repository, reference string) (exist bool, digest string, err error)
	PullManifest(repository, reference string, accepttedMediaTypes []string) (manifest distribution.Manifest, digest string, err error)
	PushManifest(repository, reference, mediaType string, payload []byte) error
	// the "reference" can be "tag" or "digest", the function needs to handle both
	DeleteManifest(repository, reference string) error
	BlobExist(repository, digest string) (exist bool, err error)
	PullBlob(repository, digest string) (size int64, blob io.ReadCloser, err error)
	PushBlob(repository, digest string, size int64, blob io.Reader) error
}

// DefaultImageRegistry provides a default implementation for interface ImageRegistry
type DefaultImageRegistry struct {
	sync.RWMutex
	*registry_pkg.Registry
	registry *model.Registry
	client   *http.Client
	clients  map[string]*registry_pkg.Repository
}

// NewDefaultImageRegistry returns an instance of DefaultImageRegistry
func NewDefaultImageRegistry(registry *model.Registry) (*DefaultImageRegistry, error) {
	transport := util.GetHTTPTransport(registry.Insecure)
	modifiers := []modifier.Modifier{
		&auth.UserAgentModifier{
			UserAgent: UserAgentReplication,
		},
	}
	if registry.Credential != nil {
		var cred modifier.Modifier
		if registry.Credential.Type == model.CredentialTypeSecret {
			cred = common_http_auth.NewSecretAuthorizer(registry.Credential.AccessSecret)
		} else {
			cred = auth.NewBasicAuthCredential(
				registry.Credential.AccessKey,
				registry.Credential.AccessSecret)
		}
		tokenServiceURL := ""
		// the registry is a local Harbor instance if the core URL is specified,
		// use the internal token service URL instead
		if len(registry.CoreURL) > 0 {
			tokenServiceURL = fmt.Sprintf("%s/service/token", registry.CoreURL)
		}
		authorizer := auth.NewStandardTokenAuthorizer(&http.Client{
			Transport: transport,
		}, cred, tokenServiceURL)

		modifiers = append(modifiers, authorizer)
	}
	client := &http.Client{
		Transport: registry_pkg.NewTransport(transport, modifiers...),
	}
	reg, err := registry_pkg.NewRegistry(registry.URL, client)
	if err != nil {
		return nil, err
	}
	return &DefaultImageRegistry{
		Registry: reg,
		client:   client,
		registry: registry,
		clients:  map[string]*registry_pkg.Repository{},
	}, nil
}

func (d *DefaultImageRegistry) getClient(repository string) (*registry_pkg.Repository, error) {
	client, exist := d.clients[repository]
	if exist {
		return client, nil
	}

	return d.create(repository)
}

func (d *DefaultImageRegistry) create(repository string) (*registry_pkg.Repository, error) {
	d.Lock()
	defer d.Unlock()
	// double check
	client, exist := d.clients[repository]
	if exist {
		return client, nil
	}

	client, err := registry_pkg.NewRepository(repository, d.registry.URL, d.client)
	if err != nil {
		return nil, err
	}
	d.clients[repository] = client
	return client, nil
}

// HealthCheck checks health status of a registry
func (d *DefaultImageRegistry) HealthCheck() (model.HealthStatus, error) {
	var err error
	if d.registry.Credential == nil {
		err = d.PingSimple()
	} else {
		err = d.Ping()
	}
	if err != nil {
		log.Errorf("failed to ping registry %s: %v", d.registry.URL, err)
		return model.Unhealthy, nil
	}
	return model.Healthy, nil
}

// FetchImages ...
func (d *DefaultImageRegistry) FetchImages(namespaces []string, filters []*model.Filter) ([]*model.Resource, error) {
	return nil, errors.New("not implemented")
}

// ManifestExist ...
func (d *DefaultImageRegistry) ManifestExist(repository, reference string) (bool, string, error) {
	client, err := d.getClient(repository)
	if err != nil {
		return false, "", err
	}
	digest, exist, err := client.ManifestExist(reference)
	return exist, digest, err
}

// PullManifest ...
func (d *DefaultImageRegistry) PullManifest(repository, reference string, accepttedMediaTypes []string) (distribution.Manifest, string, error) {
	client, err := d.getClient(repository)
	if err != nil {
		return nil, "", err
	}
	digest, mediaType, payload, err := client.PullManifest(reference, accepttedMediaTypes)
	if err != nil {
		return nil, "", err
	}
	if strings.Contains(mediaType, "application/json") {
		mediaType = schema1.MediaTypeManifest
	}
	manifest, _, err := registry_pkg.UnMarshal(mediaType, payload)
	if err != nil {
		return nil, "", err
	}
	return manifest, digest, nil
}

// PushManifest ...
func (d *DefaultImageRegistry) PushManifest(repository, reference, mediaType string, payload []byte) error {
	client, err := d.getClient(repository)
	if err != nil {
		return err
	}
	_, err = client.PushManifest(reference, mediaType, payload)
	return err
}

// DeleteManifest ...
func (d *DefaultImageRegistry) DeleteManifest(repository, reference string) error {
	client, err := d.getClient(repository)
	if err != nil {
		return err
	}
	digest := reference
	if !isDigest(digest) {
		dgt, exist, err := client.ManifestExist(reference)
		if err != nil {
			return err
		}
		if !exist {
			log.Debugf("the manifest of %s:%s doesn't exist", repository, reference)
			return nil
		}
		digest = dgt
	}
	return client.DeleteManifest(digest)
}

// BlobExist ...
func (d *DefaultImageRegistry) BlobExist(repository, digest string) (bool, error) {
	client, err := d.getClient(repository)
	if err != nil {
		return false, err
	}
	return client.BlobExist(digest)
}

// PullBlob ...
func (d *DefaultImageRegistry) PullBlob(repository, digest string) (int64, io.ReadCloser, error) {
	client, err := d.getClient(repository)
	if err != nil {
		return 0, nil, err
	}
	return client.PullBlob(digest)
}

// PushBlob ...
func (d *DefaultImageRegistry) PushBlob(repository, digest string, size int64, blob io.Reader) error {
	client, err := d.getClient(repository)
	if err != nil {
		return err
	}
	return client.PushBlob(digest, size, blob)
}

func isDigest(str string) bool {
	return strings.Contains(str, ":")
}

// ListTag ...
func (d *DefaultImageRegistry) ListTag(repository string) ([]string, error) {
	client, err := d.getClient(repository)
	if err != nil {
		return []string{}, err
	}
	return client.ListTag()
}
<commit_msg>fix deadlock when client not exist.<commit_after>// Copyright Project Harbor 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 adapter

import (
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"

	"github.com/docker/distribution"
	"github.com/docker/distribution/manifest/schema1"
	"github.com/goharbor/harbor/src/common/http/modifier"
	common_http_auth "github.com/goharbor/harbor/src/common/http/modifier/auth"
	"github.com/goharbor/harbor/src/common/utils/log"
	registry_pkg "github.com/goharbor/harbor/src/common/utils/registry"
	"github.com/goharbor/harbor/src/common/utils/registry/auth"
	"github.com/goharbor/harbor/src/replication/model"
	"github.com/goharbor/harbor/src/replication/util"
)

// const definition
const (
	UserAgentReplication = "harbor-replication-service"
)

// ImageRegistry defines the capabilities that an image registry should have
type ImageRegistry interface {
	FetchImages(namespaces []string, filters []*model.Filter) ([]*model.Resource, error)
	ManifestExist(repository, reference string) (exist bool, digest string, err error)
	PullManifest(repository, reference string, accepttedMediaTypes []string) (manifest distribution.Manifest, digest string, err error)
	PushManifest(repository, reference, mediaType string, payload []byte) error
	// the "reference" can be "tag" or "digest", the function needs to handle both
	DeleteManifest(repository, reference string) error
	BlobExist(repository, digest string) (exist bool, err error)
	PullBlob(repository, digest string) (size int64, blob io.ReadCloser, err error)
	PushBlob(repository, digest string, size int64, blob io.Reader) error
}

// DefaultImageRegistry provides a default implementation for interface ImageRegistry
type DefaultImageRegistry struct {
	sync.RWMutex
	*registry_pkg.Registry
	registry *model.Registry
	client   *http.Client
	clients  map[string]*registry_pkg.Repository
}

// NewDefaultImageRegistry returns an instance of DefaultImageRegistry
func NewDefaultImageRegistry(registry *model.Registry) (*DefaultImageRegistry, error) {
	transport := util.GetHTTPTransport(registry.Insecure)
	modifiers := []modifier.Modifier{
		&auth.UserAgentModifier{
			UserAgent: UserAgentReplication,
		},
	}
	if registry.Credential != nil {
		var cred modifier.Modifier
		if registry.Credential.Type == model.CredentialTypeSecret {
			cred = common_http_auth.NewSecretAuthorizer(registry.Credential.AccessSecret)
		} else {
			cred = auth.NewBasicAuthCredential(
				registry.Credential.AccessKey,
				registry.Credential.AccessSecret)
		}
		tokenServiceURL := ""
		// the registry is a local Harbor instance if the core URL is specified,
		// use the internal token service URL instead
		if len(registry.CoreURL) > 0 {
			tokenServiceURL = fmt.Sprintf("%s/service/token", registry.CoreURL)
		}
		authorizer := auth.NewStandardTokenAuthorizer(&http.Client{
			Transport: transport,
		}, cred, tokenServiceURL)

		modifiers = append(modifiers, authorizer)
	}
	client := &http.Client{
		Transport: registry_pkg.NewTransport(transport, modifiers...),
	}
	reg, err := registry_pkg.NewRegistry(registry.URL, client)
	if err != nil {
		return nil, err
	}
	return &DefaultImageRegistry{
		Registry: reg,
		client:   client,
		registry: registry,
		clients:  map[string]*registry_pkg.Repository{},
	}, nil
}

func (d *DefaultImageRegistry) getClient(repository string) (*registry_pkg.Repository, error) {
	d.RLock()
	client, exist := d.clients[repository]
	d.RUnlock()
	if exist {
		return client, nil
	}

	return d.create(repository)
}

func (d *DefaultImageRegistry) create(repository string) (*registry_pkg.Repository, error) {
	d.Lock()
	defer d.Unlock()
	// double check
	client, exist := d.clients[repository]
	if exist {
		return client, nil
	}

	client, err := registry_pkg.NewRepository(repository, d.registry.URL, d.client)
	if err != nil {
		return nil, err
	}
	d.clients[repository] = client
	return client, nil
}

// HealthCheck checks health status of a registry
func (d *DefaultImageRegistry) HealthCheck() (model.HealthStatus, error) {
	var err error
	if d.registry.Credential == nil {
		err = d.PingSimple()
	} else {
		err = d.Ping()
	}
	if err != nil {
		log.Errorf("failed to ping registry %s: %v", d.registry.URL, err)
		return model.Unhealthy, nil
	}
	return model.Healthy, nil
}

// FetchImages ...
func (d *DefaultImageRegistry) FetchImages(namespaces []string, filters []*model.Filter) ([]*model.Resource, error) {
	return nil, errors.New("not implemented")
}

// ManifestExist ...
func (d *DefaultImageRegistry) ManifestExist(repository, reference string) (bool, string, error) {
	client, err := d.getClient(repository)
	if err != nil {
		return false, "", err
	}
	digest, exist, err := client.ManifestExist(reference)
	return exist, digest, err
}

// PullManifest ...
func (d *DefaultImageRegistry) PullManifest(repository, reference string, accepttedMediaTypes []string) (distribution.Manifest, string, error) {
	client, err := d.getClient(repository)
	if err != nil {
		return nil, "", err
	}
	digest, mediaType, payload, err := client.PullManifest(reference, accepttedMediaTypes)
	if err != nil {
		return nil, "", err
	}
	if strings.Contains(mediaType, "application/json") {
		mediaType = schema1.MediaTypeManifest
	}
	manifest, _, err := registry_pkg.UnMarshal(mediaType, payload)
	if err != nil {
		return nil, "", err
	}
	return manifest, digest, nil
}

// PushManifest ...
func (d *DefaultImageRegistry) PushManifest(repository, reference, mediaType string, payload []byte) error {
	client, err := d.getClient(repository)
	if err != nil {
		return err
	}
	_, err = client.PushManifest(reference, mediaType, payload)
	return err
}

// DeleteManifest ...
func (d *DefaultImageRegistry) DeleteManifest(repository, reference string) error {
	client, err := d.getClient(repository)
	if err != nil {
		return err
	}
	digest := reference
	if !isDigest(digest) {
		dgt, exist, err := client.ManifestExist(reference)
		if err != nil {
			return err
		}
		if !exist {
			log.Debugf("the manifest of %s:%s doesn't exist", repository, reference)
			return nil
		}
		digest = dgt
	}
	return client.DeleteManifest(digest)
}

// BlobExist ...
func (d *DefaultImageRegistry) BlobExist(repository, digest string) (bool, error) {
	client, err := d.getClient(repository)
	if err != nil {
		return false, err
	}
	return client.BlobExist(digest)
}

// PullBlob ...
func (d *DefaultImageRegistry) PullBlob(repository, digest string) (int64, io.ReadCloser, error) {
	client, err := d.getClient(repository)
	if err != nil {
		return 0, nil, err
	}
	return client.PullBlob(digest)
}

// PushBlob ...
func (d *DefaultImageRegistry) PushBlob(repository, digest string, size int64, blob io.Reader) error {
	client, err := d.getClient(repository)
	if err != nil {
		return err
	}
	return client.PushBlob(digest, size, blob)
}

func isDigest(str string) bool {
	return strings.Contains(str, ":")
}

// ListTag ...
func (d *DefaultImageRegistry) ListTag(repository string) ([]string, error) {
	client, err := d.getClient(repository)
	if err != nil {
		return []string{}, err
	}
	return client.ListTag()
}
<|endoftext|>