<commit_before>// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package x509

import (
	"crypto/dsa"
	"crypto/ecdsa"
	"crypto/rsa"
	"encoding/asn1"
	"encoding/json"

	"time"

	"github.com/asaskevich/govalidator"
	"github.com/zmap/zgrab/ztools/keys"
	"github.com/zmap/zgrab/ztools/x509/pkix"
)

type auxKeyUsage struct {
	DigitalSignature  bool   `json:"digital_signature,omitempty"`
	ContentCommitment bool   `json:"content_commitment,omitempty"`
	KeyEncipherment   bool   `json:"key_encipherment,omitempty"`
	DataEncipherment  bool   `json:"data_encipherment,omitempty"`
	KeyAgreement      bool   `json:"key_agreement,omitempty"`
	CertificateSign   bool   `json:"certificate_sign,omitempty"`
	CRLSign           bool   `json:"crl_sign,omitempty"`
	EncipherOnly      bool   `json:"encipher_only,omitempty"`
	DecipherOnly      bool   `json:"decipher_only,omitempty"`
	Value             uint32 `json:"value"`
}

// MarshalJSON implements the json.Marshaler interface
func (k KeyUsage) MarshalJSON() ([]byte, error) {
	var enc auxKeyUsage
	enc.Value = uint32(k)
	if k&KeyUsageDigitalSignature > 0 {
		enc.DigitalSignature = true
	}
	if k&KeyUsageContentCommitment > 0 {
		enc.ContentCommitment = true
	}
	if k&KeyUsageKeyEncipherment > 0 {
		enc.KeyEncipherment = true
	}
	if k&KeyUsageDataEncipherment > 0 {
		enc.DataEncipherment = true
	}
	if k&KeyUsageKeyAgreement > 0 {
		enc.KeyAgreement = true
	}
	if k&KeyUsageCertSign > 0 {
		enc.CertificateSign = true
	}
	if k&KeyUsageCRLSign > 0 {
		enc.CRLSign = true
	}
	if k&KeyUsageEncipherOnly > 0 {
		enc.EncipherOnly = true
	}
	if k&KeyUsageDecipherOnly > 0 {
		enc.DecipherOnly = true
	}
	return json.Marshal(&enc)
}

// UnmarshalJSON implements the json.Unmarshler interface
func (k *KeyUsage) UnmarshalJSON(b []byte) error {
	var aux auxKeyUsage
	if err := json.Unmarshal(b, &aux); err != nil {
		return err
	}
	// TODO: validate the flags match
	v := int(aux.Value)
	*k = KeyUsage(v)
	return nil
}

type auxSignatureAlgorithm struct {
	Name string      `json:"name,omitempty"`
	OID  pkix.AuxOID `json:"oid"`
}

// MarshalJSON implements the json.Marshaler interface
func (s *SignatureAlgorithm) MarshalJSON() ([]byte, error) {
	aux := auxSignatureAlgorithm{
		Name: s.String(),
	}
	for _, val := range signatureAlgorithmDetails {
		if val.algo == *s {
			aux.OID = make([]int, len(val.oid))
			for idx := range val.oid {
				aux.OID[idx] = val.oid[idx]
			}
		}
	}
	return json.Marshal(&aux)
}

// UnmarshalJSON implements the json.Unmarshler interface
func (s *SignatureAlgorithm) UnmarshalJSON(b []byte) error {
	var aux auxSignatureAlgorithm
	if err := json.Unmarshal(b, &aux); err != nil {
		return err
	}
	*s = UnknownSignatureAlgorithm
	oid := asn1.ObjectIdentifier(aux.OID.AsSlice())
	for _, val := range signatureAlgorithmDetails {
		if val.oid.Equal(oid) {
			*s = val.algo
			break
		}
	}
	return nil
}

type auxPublicKeyAlgorithm struct {
	Name string      `json:"name,omitempty"`
	OID  pkix.AuxOID `json:"oid"`
}

// MarshalJSON implements the json.Marshaler interface
func (p *PublicKeyAlgorithm) MarshalJSON() ([]byte, error) {
	aux := auxPublicKeyAlgorithm{
		Name: p.String(),
	}
	return json.Marshal(&aux)
}

// UnmarshalJSON implements the json.Unmarshaler interface
func (p *PublicKeyAlgorithm) UnmarshalJSON(b []byte) error {
	var aux auxPublicKeyAlgorithm
	if err := json.Unmarshal(b, &aux); err != nil {
		return err
	}
	panic("unimplemented")
}

type auxValidity struct {
	Start          string `json:"start"`
	End            string `json:"end"`
	ValidityPeriod int    `json:"length"`
}

func (v *validity) MarshalJSON() ([]byte, error) {
	aux := auxValidity{
		Start:          v.NotBefore.UTC().Format(time.RFC3339),
		End:            v.NotAfter.UTC().Format(time.RFC3339),
		ValidityPeriod: int(v.NotAfter.Sub(v.NotBefore).Seconds()),
	}
	return json.Marshal(&aux)
}

func (v *validity) UnmarshalJSON(b []byte) error {
	var aux auxValidity
	if err := json.Unmarshal(b, &aux); err != nil {
		return err
	}
	var err error
	if v.NotBefore, err = time.Parse(time.RFC3339, aux.Start); err != nil {
		return err
	}
	if v.NotAfter, err = time.Parse(time.RFC3339, aux.End); err != nil {
		return err
	}

	return nil
}

type jsonSubjectKeyInfo struct {
	KeyAlgorithm    PublicKeyAlgorithm     `json:"key_algorithm"`
	RSAPublicKey    *keys.RSAPublicKey     `json:"rsa_public_key,omitempty"`
	DSAPublicKey    interface{}            `json:"dsa_public_key,omitempty"`
	ECDSAPublicKey  interface{}            `json:"ecdsa_public_key,omitempty"`
	SPKIFingerprint CertificateFingerprint `json:"fingerprint_sha256"`
}

type jsonSignature struct {
	SignatureAlgorithm SignatureAlgorithm `json:"signature_algorithm"`
	Value              []byte             `json:"value"`
	Valid              bool               `json:"valid"`
	SelfSigned         bool               `json:"self_signed"`
}

type fullValidity struct {
	validity
	ValidityPeriod int
}

type jsonCertificate struct {
	Version                   int                          `json:"version"`
	SerialNumber              string                       `json:"serial_number"`
	SignatureAlgorithm        SignatureAlgorithm           `json:"signature_algorithm"`
	Issuer                    pkix.Name                    `json:"issuer"`
	IssuerDN                  string                       `json:"issuer_dn,omitempty"`
	Validity                  fullValidity                 `json:"validity"`
	Subject                   pkix.Name                    `json:"subject"`
	SubjectDN                 string                       `json:"subject_dn,omitempty"`
	SubjectKeyInfo            jsonSubjectKeyInfo           `json:"subject_key_info"`
	Extensions                *CertificateExtensions       `json:"extensions,omitempty"`
	UnknownExtensions         UnknownCertificateExtensions `json:"unknown_extensions,omitempty"`
	Signature                 jsonSignature                `json:"signature"`
	FingerprintMD5            CertificateFingerprint       `json:"fingerprint_md5"`
	FingerprintSHA1           CertificateFingerprint       `json:"fingerprint_sha1"`
	FingerprintSHA256         CertificateFingerprint       `json:"fingerprint_sha256"`
	SPKISubjectFingerprint    CertificateFingerprint       `json:"spki_subject_fingerprint"`
	TBSCertificateFingerprint CertificateFingerprint       `json:"tbs_fingerprint"`
	ValidationLevel           CertValidationLevel          `json:"validation_level"`
	Names                     []string                     `json:"names"`
}

func (c *Certificate) MarshalJSON() ([]byte, error) {
	// Fill out the certificate
	jc := new(jsonCertificate)
	jc.Version = c.Version
	jc.SerialNumber = c.SerialNumber.String()
	jc.SignatureAlgorithm = c.SignatureAlgorithm
	jc.Issuer = c.Issuer
	jc.IssuerDN = c.Issuer.String()
	jc.Validity.NotBefore = c.NotBefore
	jc.Validity.NotAfter = c.NotAfter
	jc.Validity.ValidityPeriod = c.ValidityPeriod
	jc.Subject = c.Subject
	jc.SubjectDN = c.Subject.String()
	jc.SubjectKeyInfo.KeyAlgorithm = c.PublicKeyAlgorithm

	// Include all subject names, DNS names there are
	for _, obj := range c.Subject.Names {

		switch name := obj.Value.(type) {
		case string:

			if name[0] == '*' {
				name = name[2:]
			}
			// Check that this is actually a url and not something else
			if govalidator.IsURL(name) {
				jc.Names = append(jc.Names, name)
			}
		}
	}

	for _, name := range c.DNSNames {

		if name[0] == '*' {
			name = name[2:]
		}

		if govalidator.IsURL(name) {
			jc.Names = append(jc.Names, name)
		}
	}

	// Pull out the key
	keyMap := make(map[string]interface{})

	switch key := c.PublicKey.(type) {
	case *rsa.PublicKey:
		rsaKey := new(keys.RSAPublicKey)
		rsaKey.PublicKey = key
		jc.SubjectKeyInfo.RSAPublicKey = rsaKey
	case *dsa.PublicKey:
		keyMap["p"] = key.P.Bytes()
		keyMap["q"] = key.Q.Bytes()
		keyMap["g"] = key.G.Bytes()
		keyMap["y"] = key.Y.Bytes()
		jc.SubjectKeyInfo.DSAPublicKey = keyMap
	case *ecdsa.PublicKey:
		params := key.Params()
		keyMap["p"] = params.P.Bytes()
		keyMap["n"] = params.N.Bytes()
		keyMap["b"] = params.B.Bytes()
		keyMap["gx"] = params.Gx.Bytes()
		keyMap["gy"] = params.Gy.Bytes()
		keyMap["x"] = key.X.Bytes()
		keyMap["y"] = key.Y.Bytes()
		jc.SubjectKeyInfo.ECDSAPublicKey = keyMap
	case *AugmentedECDSA:
		pub := key.Pub

		keyMap["pub"] = key.Raw.Bytes

		params := pub.Params()
		keyMap["p"] = params.P.Bytes()
		keyMap["n"] = params.N.Bytes()
		keyMap["b"] = params.B.Bytes()
		keyMap["gx"] = params.Gx.Bytes()
		keyMap["gy"] = params.Gy.Bytes()
		keyMap["x"] = pub.X.Bytes()
		keyMap["y"] = pub.Y.Bytes()

		keyMap["asn1_oid"] = c.SignatureAlgorithmOID

		jc.SubjectKeyInfo.ECDSAPublicKey = keyMap
	}

	jc.Extensions, jc.UnknownExtensions = c.jsonifyExtensions()

	// TODO: Handle the fact this might not match
	jc.Signature.SignatureAlgorithm = jc.SignatureAlgorithm
	jc.Signature.Value = c.Signature
	jc.Signature.Valid = c.validSignature
	if c.Subject.CommonName == c.Issuer.CommonName {
		jc.Signature.SelfSigned = true
	}
	jc.FingerprintMD5 = c.FingerprintMD5
	jc.FingerprintSHA1 = c.FingerprintSHA1
	jc.FingerprintSHA256 = c.FingerprintSHA256
	jc.SPKISubjectFingerprint = c.SPKISubjectFingerprint
	jc.TBSCertificateFingerprint = c.TBSCertificateFingerprint
	jc.ValidationLevel = c.ValidationLevel

	return json.Marshal(jc)
}
<commit_msg>Fixed bug with url parseing. #140<commit_after>// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package x509

import (
	"crypto/dsa"
	"crypto/ecdsa"
	"crypto/rsa"
	"encoding/asn1"
	"encoding/json"

	"time"

	"github.com/asaskevich/govalidator"
	"github.com/zmap/zgrab/ztools/keys"
	"github.com/zmap/zgrab/ztools/x509/pkix"
)

type auxKeyUsage struct {
	DigitalSignature  bool   `json:"digital_signature,omitempty"`
	ContentCommitment bool   `json:"content_commitment,omitempty"`
	KeyEncipherment   bool   `json:"key_encipherment,omitempty"`
	DataEncipherment  bool   `json:"data_encipherment,omitempty"`
	KeyAgreement      bool   `json:"key_agreement,omitempty"`
	CertificateSign   bool   `json:"certificate_sign,omitempty"`
	CRLSign           bool   `json:"crl_sign,omitempty"`
	EncipherOnly      bool   `json:"encipher_only,omitempty"`
	DecipherOnly      bool   `json:"decipher_only,omitempty"`
	Value             uint32 `json:"value"`
}

// MarshalJSON implements the json.Marshaler interface
func (k KeyUsage) MarshalJSON() ([]byte, error) {
	var enc auxKeyUsage
	enc.Value = uint32(k)
	if k&KeyUsageDigitalSignature > 0 {
		enc.DigitalSignature = true
	}
	if k&KeyUsageContentCommitment > 0 {
		enc.ContentCommitment = true
	}
	if k&KeyUsageKeyEncipherment > 0 {
		enc.KeyEncipherment = true
	}
	if k&KeyUsageDataEncipherment > 0 {
		enc.DataEncipherment = true
	}
	if k&KeyUsageKeyAgreement > 0 {
		enc.KeyAgreement = true
	}
	if k&KeyUsageCertSign > 0 {
		enc.CertificateSign = true
	}
	if k&KeyUsageCRLSign > 0 {
		enc.CRLSign = true
	}
	if k&KeyUsageEncipherOnly > 0 {
		enc.EncipherOnly = true
	}
	if k&KeyUsageDecipherOnly > 0 {
		enc.DecipherOnly = true
	}
	return json.Marshal(&enc)
}

// UnmarshalJSON implements the json.Unmarshler interface
func (k *KeyUsage) UnmarshalJSON(b []byte) error {
	var aux auxKeyUsage
	if err := json.Unmarshal(b, &aux); err != nil {
		return err
	}
	// TODO: validate the flags match
	v := int(aux.Value)
	*k = KeyUsage(v)
	return nil
}

type auxSignatureAlgorithm struct {
	Name string      `json:"name,omitempty"`
	OID  pkix.AuxOID `json:"oid"`
}

// MarshalJSON implements the json.Marshaler interface
func (s *SignatureAlgorithm) MarshalJSON() ([]byte, error) {
	aux := auxSignatureAlgorithm{
		Name: s.String(),
	}
	for _, val := range signatureAlgorithmDetails {
		if val.algo == *s {
			aux.OID = make([]int, len(val.oid))
			for idx := range val.oid {
				aux.OID[idx] = val.oid[idx]
			}
		}
	}
	return json.Marshal(&aux)
}

// UnmarshalJSON implements the json.Unmarshler interface
func (s *SignatureAlgorithm) UnmarshalJSON(b []byte) error {
	var aux auxSignatureAlgorithm
	if err := json.Unmarshal(b, &aux); err != nil {
		return err
	}
	*s = UnknownSignatureAlgorithm
	oid := asn1.ObjectIdentifier(aux.OID.AsSlice())
	for _, val := range signatureAlgorithmDetails {
		if val.oid.Equal(oid) {
			*s = val.algo
			break
		}
	}
	return nil
}

type auxPublicKeyAlgorithm struct {
	Name string      `json:"name,omitempty"`
	OID  pkix.AuxOID `json:"oid"`
}

// MarshalJSON implements the json.Marshaler interface
func (p *PublicKeyAlgorithm) MarshalJSON() ([]byte, error) {
	aux := auxPublicKeyAlgorithm{
		Name: p.String(),
	}
	return json.Marshal(&aux)
}

// UnmarshalJSON implements the json.Unmarshaler interface
func (p *PublicKeyAlgorithm) UnmarshalJSON(b []byte) error {
	var aux auxPublicKeyAlgorithm
	if err := json.Unmarshal(b, &aux); err != nil {
		return err
	}
	panic("unimplemented")
}

type auxValidity struct {
	Start          string `json:"start"`
	End            string `json:"end"`
	ValidityPeriod int    `json:"length"`
}

func (v *validity) MarshalJSON() ([]byte, error) {
	aux := auxValidity{
		Start:          v.NotBefore.UTC().Format(time.RFC3339),
		End:            v.NotAfter.UTC().Format(time.RFC3339),
		ValidityPeriod: int(v.NotAfter.Sub(v.NotBefore).Seconds()),
	}
	return json.Marshal(&aux)
}

func (v *validity) UnmarshalJSON(b []byte) error {
	var aux auxValidity
	if err := json.Unmarshal(b, &aux); err != nil {
		return err
	}
	var err error
	if v.NotBefore, err = time.Parse(time.RFC3339, aux.Start); err != nil {
		return err
	}
	if v.NotAfter, err = time.Parse(time.RFC3339, aux.End); err != nil {
		return err
	}

	return nil
}

type jsonSubjectKeyInfo struct {
	KeyAlgorithm    PublicKeyAlgorithm     `json:"key_algorithm"`
	RSAPublicKey    *keys.RSAPublicKey     `json:"rsa_public_key,omitempty"`
	DSAPublicKey    interface{}            `json:"dsa_public_key,omitempty"`
	ECDSAPublicKey  interface{}            `json:"ecdsa_public_key,omitempty"`
	SPKIFingerprint CertificateFingerprint `json:"fingerprint_sha256"`
}

type jsonSignature struct {
	SignatureAlgorithm SignatureAlgorithm `json:"signature_algorithm"`
	Value              []byte             `json:"value"`
	Valid              bool               `json:"valid"`
	SelfSigned         bool               `json:"self_signed"`
}

type fullValidity struct {
	validity
	ValidityPeriod int
}

type jsonCertificate struct {
	Version                   int                          `json:"version"`
	SerialNumber              string                       `json:"serial_number"`
	SignatureAlgorithm        SignatureAlgorithm           `json:"signature_algorithm"`
	Issuer                    pkix.Name                    `json:"issuer"`
	IssuerDN                  string                       `json:"issuer_dn,omitempty"`
	Validity                  fullValidity                 `json:"validity"`
	Subject                   pkix.Name                    `json:"subject"`
	SubjectDN                 string                       `json:"subject_dn,omitempty"`
	SubjectKeyInfo            jsonSubjectKeyInfo           `json:"subject_key_info"`
	Extensions                *CertificateExtensions       `json:"extensions,omitempty"`
	UnknownExtensions         UnknownCertificateExtensions `json:"unknown_extensions,omitempty"`
	Signature                 jsonSignature                `json:"signature"`
	FingerprintMD5            CertificateFingerprint       `json:"fingerprint_md5"`
	FingerprintSHA1           CertificateFingerprint       `json:"fingerprint_sha1"`
	FingerprintSHA256         CertificateFingerprint       `json:"fingerprint_sha256"`
	SPKISubjectFingerprint    CertificateFingerprint       `json:"spki_subject_fingerprint"`
	TBSCertificateFingerprint CertificateFingerprint       `json:"tbs_fingerprint"`
	ValidationLevel           CertValidationLevel          `json:"validation_level"`
	Names                     []string                     `json:"names"`
}

func (c *Certificate) MarshalJSON() ([]byte, error) {
	// Fill out the certificate
	jc := new(jsonCertificate)
	jc.Version = c.Version
	jc.SerialNumber = c.SerialNumber.String()
	jc.SignatureAlgorithm = c.SignatureAlgorithm
	jc.Issuer = c.Issuer
	jc.IssuerDN = c.Issuer.String()
	jc.Validity.NotBefore = c.NotBefore
	jc.Validity.NotAfter = c.NotAfter
	jc.Validity.ValidityPeriod = c.ValidityPeriod
	jc.Subject = c.Subject
	jc.SubjectDN = c.Subject.String()
	jc.SubjectKeyInfo.KeyAlgorithm = c.PublicKeyAlgorithm

	// Include all subject names, DNS names there are
	for _, obj := range c.Subject.Names {

		switch name := obj.Value.(type) {
		case string:

			flag := true

			if name[0] == '*' {
				flag = govalidator.IsURL(name[2:])
			} else {
				flag = govalidator.IsURL(name)
			}

			// Check that this is actually a url and not something else
			if flag {
				jc.Names = append(jc.Names, name)
			}
		}
	}

	for _, name := range c.DNSNames {

		flag := true

		if name[0] == '*' {
			flag = govalidator.IsURL(name[2:])
		} else {
			flag = govalidator.IsURL(name)
		}

		if flag {
			jc.Names = append(jc.Names, name)
		}
	}

	// Pull out the key
	keyMap := make(map[string]interface{})

	switch key := c.PublicKey.(type) {
	case *rsa.PublicKey:
		rsaKey := new(keys.RSAPublicKey)
		rsaKey.PublicKey = key
		jc.SubjectKeyInfo.RSAPublicKey = rsaKey
	case *dsa.PublicKey:
		keyMap["p"] = key.P.Bytes()
		keyMap["q"] = key.Q.Bytes()
		keyMap["g"] = key.G.Bytes()
		keyMap["y"] = key.Y.Bytes()
		jc.SubjectKeyInfo.DSAPublicKey = keyMap
	case *ecdsa.PublicKey:
		params := key.Params()
		keyMap["p"] = params.P.Bytes()
		keyMap["n"] = params.N.Bytes()
		keyMap["b"] = params.B.Bytes()
		keyMap["gx"] = params.Gx.Bytes()
		keyMap["gy"] = params.Gy.Bytes()
		keyMap["x"] = key.X.Bytes()
		keyMap["y"] = key.Y.Bytes()
		jc.SubjectKeyInfo.ECDSAPublicKey = keyMap
	case *AugmentedECDSA:
		pub := key.Pub

		keyMap["pub"] = key.Raw.Bytes

		params := pub.Params()
		keyMap["p"] = params.P.Bytes()
		keyMap["n"] = params.N.Bytes()
		keyMap["b"] = params.B.Bytes()
		keyMap["gx"] = params.Gx.Bytes()
		keyMap["gy"] = params.Gy.Bytes()
		keyMap["x"] = pub.X.Bytes()
		keyMap["y"] = pub.Y.Bytes()

		keyMap["asn1_oid"] = c.SignatureAlgorithmOID

		jc.SubjectKeyInfo.ECDSAPublicKey = keyMap
	}

	jc.Extensions, jc.UnknownExtensions = c.jsonifyExtensions()

	// TODO: Handle the fact this might not match
	jc.Signature.SignatureAlgorithm = jc.SignatureAlgorithm
	jc.Signature.Value = c.Signature
	jc.Signature.Valid = c.validSignature
	if c.Subject.CommonName == c.Issuer.CommonName {
		jc.Signature.SelfSigned = true
	}
	jc.FingerprintMD5 = c.FingerprintMD5
	jc.FingerprintSHA1 = c.FingerprintSHA1
	jc.FingerprintSHA256 = c.FingerprintSHA256
	jc.SPKISubjectFingerprint = c.SPKISubjectFingerprint
	jc.TBSCertificateFingerprint = c.TBSCertificateFingerprint
	jc.ValidationLevel = c.ValidationLevel

	return json.Marshal(jc)
}
<|endoftext|>


<commit_before>package dbng_test

import (
	"github.com/concourse/atc/dbng"

	"github.com/concourse/atc"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/types"
)

var _ = Describe("BuildFactory", func() {
	Describe("MarkNonInterceptibleBuilds", func() {
		Context("one-off builds", func() {

		})

		Context("pipeline builds", func() {

			It("[#139963615] marks builds that aren't the latest as non-interceptible, ", func() {
				build1, err := defaultPipeline.CreateJobBuild("some-job")
				Expect(err).NotTo(HaveOccurred())

				build2, err := defaultPipeline.CreateJobBuild("some-job")
				Expect(err).NotTo(HaveOccurred())

				build1.Finish(dbng.BuildStatusErrored)
				build2.Finish(dbng.BuildStatusErrored)

				p, _, err := defaultTeam.SavePipeline("other-pipeline", atc.Config{
					Jobs: atc.JobConfigs{
						{
							Name: "some-other-job",
						},
					},
				}, dbng.ConfigVersion(0), dbng.PipelineUnpaused)
				Expect(err).NotTo(HaveOccurred())

				pb1, err := p.CreateJobBuild("some-other-job")
				Expect(err).NotTo(HaveOccurred())

				pb2, err := p.CreateJobBuild("some-other-job")
				Expect(err).NotTo(HaveOccurred())

				pb1.Finish(dbng.BuildStatusErrored)
				pb2.Finish(dbng.BuildStatusErrored)

				err = buildFactory.MarkNonInterceptibleBuilds()
				Expect(err).NotTo(HaveOccurred())

				var i bool
				i, err = build1.Interceptible()
				Expect(err).NotTo(HaveOccurred())
				Expect(i).To(BeFalse())

				i, err = build2.Interceptible()
				Expect(err).NotTo(HaveOccurred())
				Expect(i).To(BeTrue())

				i, err = pb1.Interceptible()
				Expect(err).NotTo(HaveOccurred())
				Expect(i).To(BeFalse())

				i, err = pb2.Interceptible()
				Expect(err).NotTo(HaveOccurred())
				Expect(i).To(BeTrue())

			})

			DescribeTable("completed builds",
				func(status dbng.BuildStatus, matcher types.GomegaMatcher) {
					b, err := defaultPipeline.CreateJobBuild("some-job")
					Expect(err).NotTo(HaveOccurred())

					var i bool
					b.Finish(status)
					err = buildFactory.MarkNonInterceptibleBuilds()
					i, err = b.Interceptible()
					Expect(err).NotTo(HaveOccurred())
					Expect(i).To(matcher)
				},
				Entry("succeeded is non-interceptible", dbng.BuildStatusSucceeded, BeFalse()),
				Entry("aborted is interceptible", dbng.BuildStatusAborted, BeTrue()),
				Entry("errored is interceptible", dbng.BuildStatusErrored, BeTrue()),
				Entry("failed is interceptible", dbng.BuildStatusFailed, BeTrue()),
			)

			It("does not mark non-completed builds", func() {
				b, err := defaultPipeline.CreateJobBuild("some-job")
				Expect(err).NotTo(HaveOccurred())

				var i bool
				i, err = b.Interceptible()
				Expect(err).NotTo(HaveOccurred())
				Expect(i).To(BeTrue())

				err = buildFactory.MarkNonInterceptibleBuilds()
				i, err = b.Interceptible()
				Expect(err).NotTo(HaveOccurred())
				Expect(i).To(BeTrue())

				err = b.SaveStatus(dbng.BuildStatusStarted)
				Expect(err).NotTo(HaveOccurred())

				err = buildFactory.MarkNonInterceptibleBuilds()
				i, err = b.Interceptible()
				Expect(err).NotTo(HaveOccurred())
				Expect(i).To(BeTrue())
			})
		})
	})
})
<commit_msg>Add test for one-off builds interceptibility [#139972403]<commit_after>package dbng_test

import (
	"github.com/concourse/atc/dbng"

	"github.com/concourse/atc"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/types"
)

var _ = Describe("BuildFactory", func() {
	Describe("MarkNonInterceptibleBuilds", func() {
		Context("one-off builds", func() {
			DescribeTable("completed builds",
				func(status dbng.BuildStatus, matcher types.GomegaMatcher) {
					b, err := defaultTeam.CreateOneOffBuild()
					Expect(err).NotTo(HaveOccurred())

					var i bool
					b.Finish(status)
					err = buildFactory.MarkNonInterceptibleBuilds()
					i, err = b.Interceptible()
					Expect(err).NotTo(HaveOccurred())
					Expect(i).To(matcher)
				},
				Entry("succeeded is non-interceptible", dbng.BuildStatusSucceeded, BeFalse()),
				Entry("aborted is non-interceptible", dbng.BuildStatusAborted, BeFalse()),
				Entry("errored is non-interceptible", dbng.BuildStatusErrored, BeFalse()),
				Entry("failed is non-interceptible", dbng.BuildStatusFailed, BeFalse()),
			)

			It("non-completed is interceptible", func() {
				b, err := defaultTeam.CreateOneOffBuild()
				Expect(err).NotTo(HaveOccurred())

				var i bool
				err = buildFactory.MarkNonInterceptibleBuilds()
				i, err = b.Interceptible()
				Expect(err).NotTo(HaveOccurred())
				Expect(i).To(BeTrue())

			})
		})

		Context("pipeline builds", func() {

			It("[#139963615] marks builds that aren't the latest as non-interceptible, ", func() {
				build1, err := defaultPipeline.CreateJobBuild("some-job")
				Expect(err).NotTo(HaveOccurred())

				build2, err := defaultPipeline.CreateJobBuild("some-job")
				Expect(err).NotTo(HaveOccurred())

				build1.Finish(dbng.BuildStatusErrored)
				build2.Finish(dbng.BuildStatusErrored)

				p, _, err := defaultTeam.SavePipeline("other-pipeline", atc.Config{
					Jobs: atc.JobConfigs{
						{
							Name: "some-other-job",
						},
					},
				}, dbng.ConfigVersion(0), dbng.PipelineUnpaused)
				Expect(err).NotTo(HaveOccurred())

				pb1, err := p.CreateJobBuild("some-other-job")
				Expect(err).NotTo(HaveOccurred())

				pb2, err := p.CreateJobBuild("some-other-job")
				Expect(err).NotTo(HaveOccurred())

				pb1.Finish(dbng.BuildStatusErrored)
				pb2.Finish(dbng.BuildStatusErrored)

				err = buildFactory.MarkNonInterceptibleBuilds()
				Expect(err).NotTo(HaveOccurred())

				var i bool
				i, err = build1.Interceptible()
				Expect(err).NotTo(HaveOccurred())
				Expect(i).To(BeFalse())

				i, err = build2.Interceptible()
				Expect(err).NotTo(HaveOccurred())
				Expect(i).To(BeTrue())

				i, err = pb1.Interceptible()
				Expect(err).NotTo(HaveOccurred())
				Expect(i).To(BeFalse())

				i, err = pb2.Interceptible()
				Expect(err).NotTo(HaveOccurred())
				Expect(i).To(BeTrue())

			})

			DescribeTable("completed builds",
				func(status dbng.BuildStatus, matcher types.GomegaMatcher) {
					b, err := defaultPipeline.CreateJobBuild("some-job")
					Expect(err).NotTo(HaveOccurred())

					var i bool
					b.Finish(status)
					err = buildFactory.MarkNonInterceptibleBuilds()
					i, err = b.Interceptible()
					Expect(err).NotTo(HaveOccurred())
					Expect(i).To(matcher)
				},
				Entry("succeeded is non-interceptible", dbng.BuildStatusSucceeded, BeFalse()),
				Entry("aborted is interceptible", dbng.BuildStatusAborted, BeTrue()),
				Entry("errored is interceptible", dbng.BuildStatusErrored, BeTrue()),
				Entry("failed is interceptible", dbng.BuildStatusFailed, BeTrue()),
			)

			It("does not mark non-completed builds", func() {
				b, err := defaultPipeline.CreateJobBuild("some-job")
				Expect(err).NotTo(HaveOccurred())

				var i bool
				i, err = b.Interceptible()
				Expect(err).NotTo(HaveOccurred())
				Expect(i).To(BeTrue())

				err = buildFactory.MarkNonInterceptibleBuilds()
				i, err = b.Interceptible()
				Expect(err).NotTo(HaveOccurred())
				Expect(i).To(BeTrue())

				err = b.SaveStatus(dbng.BuildStatusStarted)
				Expect(err).NotTo(HaveOccurred())

				err = buildFactory.MarkNonInterceptibleBuilds()
				i, err = b.Interceptible()
				Expect(err).NotTo(HaveOccurred())
				Expect(i).To(BeTrue())
			})
		})
	})
})
<|endoftext|>


<commit_before>package subsumption_test

import (
	. "."
	gomock "github.com/golang/mock/gomock"
	"testing"
)

func TestCreate(t *testing.T) {
	agent := Agent{}

	if agent.Size() != 0 {
		t.Errorf("agent.Size() should be 0")
	}
}

func TestAddBehavior(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	behabior := NewMockBehavior(ctrl)
	agent := Agent{}

	agent.AddBehavior(behabior)

	if agent.Size() != 1 {
		t.Errorf("agent.Size() should be 1")
	}
}

func TestInit(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	behabior := NewMockBehavior(ctrl)
	behabior.EXPECT().Init()

	agent := Agent{}
	agent.AddBehavior(behabior)

	agent.Init()
}

func TestPerform(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	behabior := NewMockBehavior(ctrl)
	behabior.EXPECT().Sense().Return(true, nil)
	behabior.EXPECT().Perform().Return(true, nil)

	agent := Agent{}
	agent.AddBehavior(behabior)

	ret, err := agent.Perform()

	if err != nil {
		t.Errorf("unexpected agent.Perform() fail with : %v", err)
	}
	if ret != true {
		t.Errorf("agent.Perform() should be true because behabior is active")
	}
}
<commit_msg>add test case for failure with Behavior.Sense()<commit_after>package subsumption_test

import (
	. "."
	"errors"
	gomock "github.com/golang/mock/gomock"
	"testing"
)

func TestCreate(t *testing.T) {
	agent := Agent{}

	if agent.Size() != 0 {
		t.Errorf("agent.Size() should be 0")
	}
}

func TestAddBehavior(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	behabior := NewMockBehavior(ctrl)
	agent := Agent{}

	agent.AddBehavior(behabior)

	if agent.Size() != 1 {
		t.Errorf("agent.Size() should be 1")
	}
}

func TestInit(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	behabior := NewMockBehavior(ctrl)
	behabior.EXPECT().Init()

	agent := Agent{}
	agent.AddBehavior(behabior)

	agent.Init()
}

func TestPerform(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	behabior := NewMockBehavior(ctrl)
	behabior.EXPECT().Sense().Return(true, nil)
	behabior.EXPECT().Perform().Return(true, nil)

	agent := Agent{}
	agent.AddBehavior(behabior)

	ret, err := agent.Perform()

	if err != nil {
		t.Errorf("unexpected agent.Perform() fail with : %v", err)
	}
	if ret != true {
		t.Errorf("agent.Perform() should be true because behabior is active")
	}
}

func TestSenceError(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	behabior := NewMockBehavior(ctrl)
	behabior.EXPECT().Sense().Return(false, errors.New("unknown error"))

	agent := Agent{}
	agent.AddBehavior(behabior)

	ret, err := agent.Perform()

	if err != nil {
		t.Errorf("unexpected agent.Perform() fail with : %v", err)
	}
	if ret != false {
		t.Errorf("agent.Perform() should be false because no behavior is active")
	}
}
<|endoftext|>


<commit_before>package post

import (
	"github.com/hackform/governor"
	"github.com/hackform/governor/service/post/comment/model"
	"github.com/labstack/echo"
	"github.com/sirupsen/logrus"
	"net/http"
)

type (
	reqGetComments struct {
		Postid string `json:"-"`
		Amount int    `json:"amount"`
		Offset int    `json:"offset"`
	}

	reqGetComment struct {
		Postid    string `json:"-"`
		Commentid string `json:"-"`
		Amount    int    `json:"amount"`
		Offset    int    `json:"offset"`
	}

	resGetComments struct {
		Comments commentmodel.ModelSlice `json:"comments"`
	}

	resGetComment struct {
		commentmodel.Model
	}
)

func (r *reqGetComments) valid() *governor.Error {
	if err := hasPostid(r.Postid); err != nil {
		return err
	}
	if err := validAmount(r.Amount); err != nil {
		return err
	}
	if err := validOffset(r.Offset); err != nil {
		return err
	}
	return nil
}

func (r *reqGetComment) valid() *governor.Error {
	if err := hasPostid(r.Postid); err != nil {
		return err
	}
	if err := hasPostid(r.Commentid); err != nil {
		return err
	}
	if err := validAmount(r.Amount); err != nil {
		return err
	}
	if err := validOffset(r.Offset); err != nil {
		return err
	}
	return nil
}

func (p *Post) mountComments(conf governor.Config, r *echo.Group, l *logrus.Logger) error {
	db := p.db.DB()

	r.GET("/:postid/c", func(c echo.Context) error {
		rcomms := &reqGetComments{}
		if err := c.Bind(rcomms); err != nil {
			return governor.NewErrorUser(moduleIDComments, err.Error(), 0, http.StatusBadRequest)
		}
		rcomms.Postid = c.Param("postid")
		if err := rcomms.valid(); err != nil {
			return err
		}

		commentsSlice, err := commentmodel.GetResponses(db, rcomms.Postid, rcomms.Amount, rcomms.Offset)
		if err != nil {
			err.AddTrace(moduleIDComments)
			return err
		}

		return c.JSON(http.StatusOK, &resGetComments{
			Comments: commentsSlice,
		})
	})

	r.GET("/:postid/c/:commentid", func(c echo.Context) error {
		rcomms := &reqGetComment{}
		if err := c.Bind(rcomms); err != nil {
			return governor.NewErrorUser(moduleIDComments, err.Error(), 0, http.StatusBadRequest)
		}
		rcomms.Postid = c.Param("postid")
		rcomms.Commentid = c.Param("commentid")
		if err := rcomms.valid(); err != nil {
			return err
		}

		comment, err := commentmodel.GetByIDB64(db, rcomms.Commentid, rcomms.Postid)
		if err != nil {
			if err.Code() == 2 {
				err.SetErrorUser()
			}
			err.AddTrace(moduleIDComments)
			return err
		}

		return c.JSON(http.StatusOK, &resGetComment{
			Model: *comment,
		})
	})
	return nil
}
<commit_msg>comments children<commit_after>package post

import (
	"github.com/hackform/governor"
	"github.com/hackform/governor/service/post/comment/model"
	"github.com/labstack/echo"
	"github.com/sirupsen/logrus"
	"net/http"
)

type (
	reqGetComments struct {
		Postid string `json:"-"`
		Amount int    `json:"amount"`
		Offset int    `json:"offset"`
	}

	reqGetComment struct {
		Postid    string `json:"-"`
		Commentid string `json:"-"`
		Amount    int    `json:"amount"`
		Offset    int    `json:"offset"`
	}

	resGetComments struct {
		Comments commentmodel.ModelSlice `json:"comments"`
	}

	resGetComment struct {
		commentmodel.Model
	}
)

func (r *reqGetComments) valid() *governor.Error {
	if err := hasPostid(r.Postid); err != nil {
		return err
	}
	if err := validAmount(r.Amount); err != nil {
		return err
	}
	if err := validOffset(r.Offset); err != nil {
		return err
	}
	return nil
}

func (r *reqGetComment) valid() *governor.Error {
	if err := hasPostid(r.Postid); err != nil {
		return err
	}
	if err := hasPostid(r.Commentid); err != nil {
		return err
	}
	if err := validAmount(r.Amount); err != nil {
		return err
	}
	if err := validOffset(r.Offset); err != nil {
		return err
	}
	return nil
}

func (p *Post) mountComments(conf governor.Config, r *echo.Group, l *logrus.Logger) error {
	db := p.db.DB()

	r.GET("/:postid/c", func(c echo.Context) error {
		rcomms := &reqGetComments{}
		if err := c.Bind(rcomms); err != nil {
			return governor.NewErrorUser(moduleIDComments, err.Error(), 0, http.StatusBadRequest)
		}
		rcomms.Postid = c.Param("postid")
		if err := rcomms.valid(); err != nil {
			return err
		}

		commentsSlice, err := commentmodel.GetResponses(db, rcomms.Postid, rcomms.Amount, rcomms.Offset)
		if err != nil {
			err.AddTrace(moduleIDComments)
			return err
		}

		return c.JSON(http.StatusOK, &resGetComments{
			Comments: commentsSlice,
		})
	})

	r.GET("/:postid/c/:commentid", func(c echo.Context) error {
		rcomms := &reqGetComment{}
		if err := c.Bind(rcomms); err != nil {
			return governor.NewErrorUser(moduleIDComments, err.Error(), 0, http.StatusBadRequest)
		}
		rcomms.Postid = c.Param("postid")
		rcomms.Commentid = c.Param("commentid")
		if err := rcomms.valid(); err != nil {
			return err
		}

		comment, err := commentmodel.GetByIDB64(db, rcomms.Commentid, rcomms.Postid)
		if err != nil {
			if err.Code() == 2 {
				err.SetErrorUser()
			}
			err.AddTrace(moduleIDComments)
			return err
		}

		return c.JSON(http.StatusOK, &resGetComment{
			Model: *comment,
		})
	})

	r.GET("/:postid/c/:commentid/children", func(c echo.Context) error {
		rcomms := &reqGetComment{}
		if err := c.Bind(rcomms); err != nil {
			return governor.NewErrorUser(moduleIDComments, err.Error(), 0, http.StatusBadRequest)
		}
		rcomms.Postid = c.Param("postid")
		rcomms.Commentid = c.Param("commentid")
		if err := rcomms.valid(); err != nil {
			return err
		}

		comments, err := commentmodel.GetChildren(db, rcomms.Commentid, rcomms.Postid, rcomms.Amount, rcomms.Offset)
		if err != nil {
			if err.Code() == 2 {
				err.SetErrorUser()
			}
			err.AddTrace(moduleIDComments)
			return err
		}

		return c.JSON(http.StatusOK, &resGetComments{
			Comments: comments,
		})
	})
	return nil
}
<|endoftext|>


<commit_before>package project

import (
	"database/sql"
	"encoding/base64"
	"encoding/json"

	"github.com/go-gorp/gorp"
	yaml "gopkg.in/yaml.v2"

	"github.com/ovh/cds/engine/api/database/gorpmapping"
	"github.com/ovh/cds/engine/api/secret"
	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/log"
)

type dbProject sdk.Project
type dbVariable sdk.Variable
type dbProjectVariableAudit sdk.ProjectVariableAudit
type dbProjectKey sdk.ProjectKey

func init() {
	gorpmapping.Register(gorpmapping.New(dbProject{}, "project", true, "id"))
	gorpmapping.Register(gorpmapping.New(dbProjectVariableAudit{}, "project_variable_audit", true, "id"))
	gorpmapping.Register(gorpmapping.New(dbProjectKey{}, "project_key", false))
}

// PostGet is a db hook
func (p *dbProject) PostGet(db gorp.SqlExecutor) error {
	var fields = struct {
		Metadata   sql.NullString `db:"metadata"`
		VCSServers []byte         `db:"vcs_servers"`
	}{}

	if err := db.QueryRow("select metadata,vcs_servers from project where id = $1", p.ID).Scan(&fields.Metadata, &fields.VCSServers); err != nil {
		return err
	}

	if err := gorpmapping.JSONNullString(fields.Metadata, &p.Metadata); err != nil {
		return err
	}

	if len(fields.VCSServers) > 0 {
		clearVCSServer, err := secret.Decrypt([]byte(fields.VCSServers))
		if err != nil {
			return err
		}

		if len(clearVCSServer) > 0 {
			if err := yaml.Unmarshal(clearVCSServer, &p.VCSServers); err != nil {
				log.Error("Unable to load project %d: %v", p.ID, err)
				p.VCSServers = nil
				db.Update(p)
			}
		}
	}

	return nil
}

// PostUpdate is a db hook
func (p *dbProject) PostUpdate(db gorp.SqlExecutor) error {
	bm, err := json.Marshal(p.Metadata)
	if err != nil {
		return err
	}

	bv := []byte{}
	if len(p.VCSServers) > 0 {
		b1, err := yaml.Marshal(p.VCSServers)
		if err != nil {
			return err
		}
		encryptedVCSServerStr, err := secret.Encrypt(b1)
		if err != nil {
			return err
		}
		bv = encryptedVCSServerStr
	}

	_, err = db.Exec("update project set metadata = $2, vcs_servers = $3 where id = $1", p.ID, bm, bv)
	return err
}

// PostInsert is a db hook
func (p *dbProject) PostInsert(db gorp.SqlExecutor) error {
	return p.PostUpdate(db)
}

// PostGet is a db hook
func (pva *dbProjectVariableAudit) PostGet(db gorp.SqlExecutor) error {
	var before, after sql.NullString
	query := "SELECT variable_before, variable_after from project_variable_audit WHERE id = $1"
	if err := db.QueryRow(query, pva.ID).Scan(&before, &after); err != nil {
		return err
	}

	if before.Valid {
		vBefore := &sdk.Variable{}
		if err := json.Unmarshal([]byte(before.String), vBefore); err != nil {
			return err
		}
		if sdk.NeedPlaceholder(vBefore.Type) {
			vBefore.Value = sdk.PasswordPlaceholder
		}
		pva.VariableBefore = vBefore

	}

	if after.Valid {
		vAfter := &sdk.Variable{}
		if err := json.Unmarshal([]byte(after.String), vAfter); err != nil {
			return err
		}
		if sdk.NeedPlaceholder(vAfter.Type) {
			vAfter.Value = sdk.PasswordPlaceholder
		}
		pva.VariableAfter = vAfter
	}

	return nil
}

// PostUpdate is a db hook
func (pva *dbProjectVariableAudit) PostUpdate(db gorp.SqlExecutor) error {
	var vB, vA sql.NullString

	if pva.VariableBefore != nil {
		v, err := json.Marshal(pva.VariableBefore)
		if err != nil {
			return err
		}
		vB.Valid = true
		vB.String = string(v)
	}

	if pva.VariableAfter != nil {
		v, err := json.Marshal(pva.VariableAfter)
		if err != nil {
			return err
		}
		vA.Valid = true
		vA.String = string(v)
	}

	query := "update project_variable_audit set variable_before = $2, variable_after = $3 where id = $1"
	if _, err := db.Exec(query, pva.ID, vB, vA); err != nil {
		return err
	}
	return nil
}

// PostInsert is a db hook
func (pva *dbProjectVariableAudit) PostInsert(db gorp.SqlExecutor) error {
	return pva.PostUpdate(db)
}

// PreInsert
func (pva *dbProjectVariableAudit) PreInsert(s gorp.SqlExecutor) error {
	if pva.VariableBefore != nil {
		if sdk.NeedPlaceholder(pva.VariableBefore.Type) {
			secret, err := secret.Encrypt([]byte(pva.VariableBefore.Value))
			if err != nil {
				return err
			}
			pva.VariableBefore.Value = base64.StdEncoding.EncodeToString(secret)
		}
	}
	if pva.VariableAfter != nil {
		if sdk.NeedPlaceholder(pva.VariableAfter.Type) {
			var err error
			secret, err := secret.Encrypt([]byte(pva.VariableAfter.Value))
			if err != nil {
				return err
			}
			pva.VariableAfter.Value = base64.StdEncoding.EncodeToString(secret)
		}
	}
	return nil
}
<commit_msg>fix (api): update project shouldn't erase vcs server (#2078)<commit_after>package project

import (
	"database/sql"
	"encoding/base64"
	"encoding/json"

	"github.com/go-gorp/gorp"
	yaml "gopkg.in/yaml.v2"

	"github.com/ovh/cds/engine/api/database/gorpmapping"
	"github.com/ovh/cds/engine/api/secret"
	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/log"
)

type dbProject sdk.Project
type dbVariable sdk.Variable
type dbProjectVariableAudit sdk.ProjectVariableAudit
type dbProjectKey sdk.ProjectKey

func init() {
	gorpmapping.Register(gorpmapping.New(dbProject{}, "project", true, "id"))
	gorpmapping.Register(gorpmapping.New(dbProjectVariableAudit{}, "project_variable_audit", true, "id"))
	gorpmapping.Register(gorpmapping.New(dbProjectKey{}, "project_key", false))
}

// PostGet is a db hook
func (p *dbProject) PostGet(db gorp.SqlExecutor) error {
	var fields = struct {
		Metadata   sql.NullString `db:"metadata"`
		VCSServers []byte         `db:"vcs_servers"`
	}{}

	if err := db.QueryRow("select metadata,vcs_servers from project where id = $1", p.ID).Scan(&fields.Metadata, &fields.VCSServers); err != nil {
		return err
	}

	if err := gorpmapping.JSONNullString(fields.Metadata, &p.Metadata); err != nil {
		return err
	}

	if len(fields.VCSServers) > 0 {
		clearVCSServer, err := secret.Decrypt([]byte(fields.VCSServers))
		if err != nil {
			return err
		}

		if len(clearVCSServer) > 0 {
			if err := yaml.Unmarshal(clearVCSServer, &p.VCSServers); err != nil {
				log.Error("Unable to load project %d: %v", p.ID, err)
				p.VCSServers = nil
				db.Update(p)
			}
		}
	}

	return nil
}

// PostUpdate is a db hook
func (p *dbProject) PostUpdate(db gorp.SqlExecutor) error {
	bm, errm := json.Marshal(p.Metadata)
	if errm != nil {
		return errm
	}

	if len(p.VCSServers) > 0 {
		b1, err := yaml.Marshal(p.VCSServers)
		if err != nil {
			return err
		}
		encryptedVCSServerStr, err := secret.Encrypt(b1)
		if err != nil {
			return err
		}
		_, err = db.Exec("update project set metadata = $2, vcs_servers = $3 where id = $1", p.ID, bm, encryptedVCSServerStr)
		return err
	}

	_, err := db.Exec("update project set metadata = $2 where id = $1", p.ID, bm)
	return err
}

// PostInsert is a db hook
func (p *dbProject) PostInsert(db gorp.SqlExecutor) error {
	return p.PostUpdate(db)
}

// PostGet is a db hook
func (pva *dbProjectVariableAudit) PostGet(db gorp.SqlExecutor) error {
	var before, after sql.NullString
	query := "SELECT variable_before, variable_after from project_variable_audit WHERE id = $1"
	if err := db.QueryRow(query, pva.ID).Scan(&before, &after); err != nil {
		return err
	}

	if before.Valid {
		vBefore := &sdk.Variable{}
		if err := json.Unmarshal([]byte(before.String), vBefore); err != nil {
			return err
		}
		if sdk.NeedPlaceholder(vBefore.Type) {
			vBefore.Value = sdk.PasswordPlaceholder
		}
		pva.VariableBefore = vBefore

	}

	if after.Valid {
		vAfter := &sdk.Variable{}
		if err := json.Unmarshal([]byte(after.String), vAfter); err != nil {
			return err
		}
		if sdk.NeedPlaceholder(vAfter.Type) {
			vAfter.Value = sdk.PasswordPlaceholder
		}
		pva.VariableAfter = vAfter
	}

	return nil
}

// PostUpdate is a db hook
func (pva *dbProjectVariableAudit) PostUpdate(db gorp.SqlExecutor) error {
	var vB, vA sql.NullString

	if pva.VariableBefore != nil {
		v, err := json.Marshal(pva.VariableBefore)
		if err != nil {
			return err
		}
		vB.Valid = true
		vB.String = string(v)
	}

	if pva.VariableAfter != nil {
		v, err := json.Marshal(pva.VariableAfter)
		if err != nil {
			return err
		}
		vA.Valid = true
		vA.String = string(v)
	}

	query := "update project_variable_audit set variable_before = $2, variable_after = $3 where id = $1"
	if _, err := db.Exec(query, pva.ID, vB, vA); err != nil {
		return err
	}
	return nil
}

// PostInsert is a db hook
func (pva *dbProjectVariableAudit) PostInsert(db gorp.SqlExecutor) error {
	return pva.PostUpdate(db)
}

// PreInsert
func (pva *dbProjectVariableAudit) PreInsert(s gorp.SqlExecutor) error {
	if pva.VariableBefore != nil {
		if sdk.NeedPlaceholder(pva.VariableBefore.Type) {
			secret, err := secret.Encrypt([]byte(pva.VariableBefore.Value))
			if err != nil {
				return err
			}
			pva.VariableBefore.Value = base64.StdEncoding.EncodeToString(secret)
		}
	}
	if pva.VariableAfter != nil {
		if sdk.NeedPlaceholder(pva.VariableAfter.Type) {
			var err error
			secret, err := secret.Encrypt([]byte(pva.VariableAfter.Value))
			if err != nil {
				return err
			}
			pva.VariableAfter.Value = base64.StdEncoding.EncodeToString(secret)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 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 kube

import (
	"github.com/prometheus/client_golang/prometheus"

	prowapi "k8s.io/test-infra/prow/apis/prowjobs/v1"
)

var (
	metricLabels = []string{
		// namespace of the job
		"job_namespace",
		// name of the job
		"job_name",
		// type of the prowjob: presubmit, postsubmit, periodic, batch
		"type",
		// state of the prowjob: triggered, pending, success, failure, aborted, error
		"state",
		// the org of the prowjob's repo
		"org",
		// the prowjob's repo
		"repo",
		// the base_ref of the prowjob's repo
		"base_ref",
	}
	prowJobs = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "prowjobs",
		Help: "Number of prowjobs in the system",
	}, metricLabels)
	prowJobTransitions = prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "prowjob_state_transitions",
		Help: "Number of prowjobs transitioning states",
	}, metricLabels)
)

type jobLabel struct {
	jobNamespace string
	jobName      string
	jobType      string
	state        string
	org          string
	repo         string
	baseRef      string
}

func (jl *jobLabel) values() []string {
	return []string{jl.jobNamespace, jl.jobName, jl.jobType, jl.state, jl.org, jl.repo, jl.baseRef}
}

func init() {
	prometheus.MustRegister(prowJobs)
	prometheus.MustRegister(prowJobTransitions)
}

func getJobLabelMap(pjs []prowapi.ProwJob) map[jobLabel]float64 {
	jobLabelMap := make(map[jobLabel]float64)

	for _, pj := range pjs {
		jobLabelMap[getJobLabel(pj)]++
	}
	return jobLabelMap
}

func getJobLabel(pj prowapi.ProwJob) jobLabel {
	jl := jobLabel{jobNamespace: pj.Namespace, jobName: pj.Spec.Job, jobType: string(pj.Spec.Type), state: string(pj.Status.State)}

	if pj.Spec.Refs != nil {
		jl.org = pj.Spec.Refs.Org
		jl.repo = pj.Spec.Refs.Repo
		jl.baseRef = pj.Spec.Refs.BaseRef
	} else if len(pj.Spec.ExtraRefs) > 0 {
		jl.org = pj.Spec.ExtraRefs[0].Org
		jl.repo = pj.Spec.ExtraRefs[0].Repo
		jl.baseRef = pj.Spec.ExtraRefs[0].BaseRef
	}

	return jl
}

type jobIdentifier struct {
	jobLabel
	buildId string
}

func getJobIdentifier(pj prowapi.ProwJob) jobIdentifier {
	return jobIdentifier{
		jobLabel: getJobLabel(pj),
		buildId:  pj.Status.BuildID,
	}
}

// previousStates records the prowJobs we were called with previously
var previousStates map[jobIdentifier]prowapi.ProwJobState

// GatherProwJobMetrics gathers prometheus metrics for prowjobs.
// Not threadsafe, ensure this is called serially.
func GatherProwJobMetrics(current []prowapi.ProwJob) {
	// This may be racing with the prometheus server but we need to remove
	// stale metrics like triggered or pending jobs that are now complete.
	prowJobs.Reset()

	// record the current state of ProwJob CRs on the system
	for jl, count := range getJobLabelMap(current) {
		prowJobs.WithLabelValues(jl.values()...).Set(count)
	}

	// record state transitions since the last time we were called
	currentStates := map[jobIdentifier]prowapi.ProwJobState{}
	for _, pj := range current {
		ji := getJobIdentifier(pj)
		state := prowapi.ProwJobState(ji.state)
		currentStates[ji] = state

		if previousState, seen := previousStates[ji]; !seen || previousState != state {
			prowJobTransitions.WithLabelValues(ji.values()...).Inc()
		}
	}

	previousStates = currentStates
}
<commit_msg>Revert "Merge pull request #17000 from stevekuznetsov/skuznets/revert-label"<commit_after>/*
Copyright 2017 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 kube

import (
	"github.com/prometheus/client_golang/prometheus"

	prowapi "k8s.io/test-infra/prow/apis/prowjobs/v1"
)

var (
	metricLabels = []string{
		// namespace of the job
		"job_namespace",
		// name of the job
		"job_name",
		// type of the prowjob: presubmit, postsubmit, periodic, batch
		"type",
		// state of the prowjob: triggered, pending, success, failure, aborted, error
		"state",
		// the org of the prowjob's repo
		"org",
		// the prowjob's repo
		"repo",
		// the base_ref of the prowjob's repo
		"base_ref",
	}
	prowJobs = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "prowjobs",
		Help: "Number of prowjobs in the system",
	}, metricLabels)
	prowJobTransitions = prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "prowjob_state_transitions",
		Help: "Number of prowjobs transitioning states",
	}, metricLabels)
)

type jobLabel struct {
	jobNamespace string
	jobName      string
	jobType      string
	state        string
	org          string
	repo         string
	baseRef      string
	cluster      string
}

func (jl *jobLabel) values() []string {
	return []string{jl.jobNamespace, jl.jobName, jl.jobType, jl.state, jl.org, jl.repo, jl.baseRef, jl.cluster}
}

func init() {
	prometheus.MustRegister(prowJobs)
	prometheus.MustRegister(prowJobTransitions)
}

func getJobLabelMap(pjs []prowapi.ProwJob) map[jobLabel]float64 {
	jobLabelMap := make(map[jobLabel]float64)

	for _, pj := range pjs {
		jobLabelMap[getJobLabel(pj)]++
	}
	return jobLabelMap
}

func getJobLabel(pj prowapi.ProwJob) jobLabel {
	jl := jobLabel{jobNamespace: pj.Namespace, jobName: pj.Spec.Job, jobType: string(pj.Spec.Type), state: string(pj.Status.State), cluster: pj.Spec.Cluster}

	if pj.Spec.Refs != nil {
		jl.org = pj.Spec.Refs.Org
		jl.repo = pj.Spec.Refs.Repo
		jl.baseRef = pj.Spec.Refs.BaseRef
	} else if len(pj.Spec.ExtraRefs) > 0 {
		jl.org = pj.Spec.ExtraRefs[0].Org
		jl.repo = pj.Spec.ExtraRefs[0].Repo
		jl.baseRef = pj.Spec.ExtraRefs[0].BaseRef
	}

	return jl
}

type jobIdentifier struct {
	jobLabel
	buildId string
}

func getJobIdentifier(pj prowapi.ProwJob) jobIdentifier {
	return jobIdentifier{
		jobLabel: getJobLabel(pj),
		buildId:  pj.Status.BuildID,
	}
}

// previousStates records the prowJobs we were called with previously
var previousStates map[jobIdentifier]prowapi.ProwJobState

// GatherProwJobMetrics gathers prometheus metrics for prowjobs.
// Not threadsafe, ensure this is called serially.
func GatherProwJobMetrics(current []prowapi.ProwJob) {
	// This may be racing with the prometheus server but we need to remove
	// stale metrics like triggered or pending jobs that are now complete.
	prowJobs.Reset()

	// record the current state of ProwJob CRs on the system
	for jl, count := range getJobLabelMap(current) {
		prowJobs.WithLabelValues(jl.values()...).Set(count)
	}

	// record state transitions since the last time we were called
	currentStates := map[jobIdentifier]prowapi.ProwJobState{}
	for _, pj := range current {
		ji := getJobIdentifier(pj)
		state := prowapi.ProwJobState(ji.state)
		currentStates[ji] = state

		if previousState, seen := previousStates[ji]; !seen || previousState != state {
			prowJobTransitions.WithLabelValues(ji.values()...).Inc()
		}
	}

	previousStates = currentStates
}
<|endoftext|>


<commit_before>package exchange

import (
	"context"
	"reflect"
	"sync"
	"testing"
	"time"

	eventstypes "github.com/containerd/containerd/api/events"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/events"
	"github.com/containerd/containerd/namespaces"
	"github.com/containerd/typeurl"
	"github.com/pkg/errors"
)

func TestExchangeBasic(t *testing.T) {
	ctx := namespaces.WithNamespace(context.Background(), t.Name())
	testevents := []events.Event{
		&eventstypes.ContainerCreate{ID: "asdf"},
		&eventstypes.ContainerCreate{ID: "qwer"},
		&eventstypes.ContainerCreate{ID: "zxcv"},
	}
	exchange := NewExchange()

	t.Log("subscribe")
	var cancel1, cancel2 func()

	// Create two subscribers for same set of events and make sure they
	// traverse the exchange.
	ctx1, cancel1 := context.WithCancel(ctx)
	eventq1, errq1 := exchange.Subscribe(ctx1)

	ctx2, cancel2 := context.WithCancel(ctx)
	eventq2, errq2 := exchange.Subscribe(ctx2)

	t.Log("publish")
	var wg sync.WaitGroup
	wg.Add(1)
	errChan := make(chan error)
	go func() {
		defer wg.Done()
		defer close(errChan)
		for _, event := range testevents {
			if err := exchange.Publish(ctx, "/test", event); err != nil {
				errChan <- err
				return
			}
		}

		t.Log("finished publishing")
	}()

	t.Log("waiting")
	wg.Wait()
	if err := <-errChan; err != nil {
		t.Fatal(err)
	}

	for _, subscriber := range []struct {
		eventq <-chan *events.Envelope
		errq   <-chan error
		cancel func()
	}{
		{
			eventq: eventq1,
			errq:   errq1,
			cancel: cancel1,
		},
		{
			eventq: eventq2,
			errq:   errq2,
			cancel: cancel2,
		},
	} {
		var received []events.Event
	subscribercheck:
		for {
			select {
			case env := <-subscriber.eventq:
				ev, err := typeurl.UnmarshalAny(env.Event)
				if err != nil {
					t.Fatal(err)
				}
				received = append(received, ev.(*eventstypes.ContainerCreate))
			case err := <-subscriber.errq:
				if err != nil {
					t.Fatal(err)
				}
				break subscribercheck
			}

			if reflect.DeepEqual(received, testevents) {
				// when we do this, we expect the errs channel to be closed and
				// this will return.
				subscriber.cancel()
			}
		}
	}
}

func TestExchangeValidateTopic(t *testing.T) {
	namespace := t.Name()
	ctx := namespaces.WithNamespace(context.Background(), namespace)
	exchange := NewExchange()

	for _, testcase := range []struct {
		input string
		err   error
	}{
		{
			input: "/test",
		},
		{
			input: "/test/test",
		},
		{
			input: "test",
			err:   errdefs.ErrInvalidArgument,
		},
	} {
		t.Run(testcase.input, func(t *testing.T) {
			event := &eventstypes.ContainerCreate{ID: t.Name()}
			if err := exchange.Publish(ctx, testcase.input, event); errors.Cause(err) != testcase.err {
				if err == nil {
					t.Fatalf("expected error %v, received nil", testcase.err)
				} else {
					t.Fatalf("expected error %v, received %v", testcase.err, err)
				}
			}

			evany, err := typeurl.MarshalAny(event)
			if err != nil {
				t.Fatal(err)
			}

			envelope := events.Envelope{
				Timestamp: time.Now().UTC(),
				Namespace: namespace,
				Topic:     testcase.input,
				Event:     evany,
			}

			// make sure we get same errors with forward.
			if err := exchange.Forward(ctx, &envelope); errors.Cause(err) != testcase.err {
				if err == nil {
					t.Fatalf("expected error %v, received nil", testcase.err)
				} else {
					t.Fatalf("expected error %v, received %v", testcase.err, err)
				}
			}

		})
	}
}
<commit_msg>add testcase for event.exchange and refactor<commit_after>package exchange

import (
	"context"
	"reflect"
	"sync"
	"testing"
	"time"

	eventstypes "github.com/containerd/containerd/api/events"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/events"
	"github.com/containerd/containerd/namespaces"
	"github.com/containerd/typeurl"
	"github.com/pkg/errors"
)

func TestExchangeBasic(t *testing.T) {
	ctx := namespaces.WithNamespace(context.Background(), t.Name())
	testevents := []events.Event{
		&eventstypes.ContainerCreate{ID: "asdf"},
		&eventstypes.ContainerCreate{ID: "qwer"},
		&eventstypes.ContainerCreate{ID: "zxcv"},
	}
	exchange := NewExchange()

	t.Log("subscribe")
	var cancel1, cancel2 func()

	// Create two subscribers for same set of events and make sure they
	// traverse the exchange.
	ctx1, cancel1 := context.WithCancel(ctx)
	eventq1, errq1 := exchange.Subscribe(ctx1)

	ctx2, cancel2 := context.WithCancel(ctx)
	eventq2, errq2 := exchange.Subscribe(ctx2)

	t.Log("publish")
	var wg sync.WaitGroup
	wg.Add(1)
	errChan := make(chan error)
	go func() {
		defer wg.Done()
		defer close(errChan)
		for _, event := range testevents {
			if err := exchange.Publish(ctx, "/test", event); err != nil {
				errChan <- err
				return
			}
		}

		t.Log("finished publishing")
	}()

	t.Log("waiting")
	wg.Wait()
	if err := <-errChan; err != nil {
		t.Fatal(err)
	}

	for _, subscriber := range []struct {
		eventq <-chan *events.Envelope
		errq   <-chan error
		cancel func()
	}{
		{
			eventq: eventq1,
			errq:   errq1,
			cancel: cancel1,
		},
		{
			eventq: eventq2,
			errq:   errq2,
			cancel: cancel2,
		},
	} {
		var received []events.Event
	subscribercheck:
		for {
			select {
			case env := <-subscriber.eventq:
				ev, err := typeurl.UnmarshalAny(env.Event)
				if err != nil {
					t.Fatal(err)
				}
				received = append(received, ev.(*eventstypes.ContainerCreate))
			case err := <-subscriber.errq:
				if err != nil {
					t.Fatal(err)
				}
				break subscribercheck
			}

			if reflect.DeepEqual(received, testevents) {
				// when we do this, we expect the errs channel to be closed and
				// this will return.
				subscriber.cancel()
			}
		}
	}
}

func TestExchangeFilters(t *testing.T) {
	var (
		ctx      = namespaces.WithNamespace(context.Background(), t.Name())
		exchange = NewExchange()

		// config events, All events will be published
		containerCreateEvents = []events.Event{
			&eventstypes.ContainerCreate{ID: "asdf"},
			&eventstypes.ContainerCreate{ID: "qwer"},
			&eventstypes.ContainerCreate{ID: "zxcv"},
		}
		taskExitEvents = []events.Event{
			&eventstypes.TaskExit{ContainerID: "abcdef"},
		}
		testEventSets = []struct {
			topic  string
			events []events.Event
		}{
			{
				topic:  "/containers/create",
				events: containerCreateEvents,
			},
			{
				topic:  "/tasks/exit",
				events: taskExitEvents,
			},
		}
		allTestEvents = func(eventSets []struct {
			topic  string
			events []events.Event
		}) (events []events.Event) {
			for _, v := range eventSets {
				events = append(events, v.events...)
			}
			return
		}(testEventSets)

		// config test cases
		testCases = []struct {
			testName string
			filters  []string

			// The fields as below are for store data. Don't config them.
			expectEvents []events.Event
			eventq       <-chan *events.Envelope
			errq         <-chan error
			cancel       func()
		}{
			{
				testName:     "No Filter",
				expectEvents: allTestEvents,
			},
			{
				testName: "Filter events by one topic",
				filters: []string{
					`topic=="/containers/create"`,
				},
				expectEvents: containerCreateEvents,
			},
			{
				testName: "Filter events by field",
				filters: []string{
					"event.id",
				},
				expectEvents: containerCreateEvents,
			},
			{
				testName: "Filter events by field OR topic",
				filters: []string{
					`topic=="/containers/create"`,
					"event.id",
				},
				expectEvents: containerCreateEvents,
			},
			{
				testName: "Filter events by regex ",
				filters: []string{
					`topic~="/containers/*"`,
				},
				expectEvents: containerCreateEvents,
			},
			{
				testName: "Filter events for by anyone of two topics",
				filters: []string{
					`topic=="/tasks/exit"`,
					`topic=="/containers/create"`,
				},
				expectEvents: append(containerCreateEvents, taskExitEvents...),
			},
			{
				testName: "Filter events for by one topic AND id",
				filters: []string{
					`topic=="/containers/create",event.id=="qwer"`,
				},
				expectEvents: []events.Event{
					&eventstypes.ContainerCreate{ID: "qwer"},
				},
			},
		}
	)

	t.Log("subscribe")
	for i := range testCases {
		var ctx1 context.Context
		ctx1, testCases[i].cancel = context.WithCancel(ctx)
		testCases[i].eventq, testCases[i].errq = exchange.Subscribe(ctx1, testCases[i].filters...)
	}

	t.Log("publish")
	var wg sync.WaitGroup
	wg.Add(1)
	errChan := make(chan error)
	go func() {
		defer wg.Done()
		defer close(errChan)
		for _, es := range testEventSets {
			for _, e := range es.events {
				if err := exchange.Publish(ctx, es.topic, e); err != nil {
					errChan <- err
					return
				}
			}
		}

		t.Log("finished publishing")
	}()

	t.Log("waiting")
	wg.Wait()
	if err := <-errChan; err != nil {
		t.Fatal(err)
	}

	t.Log("receive event")
	for _, subscriber := range testCases {
		t.Logf("test case: %q", subscriber.testName)
		var received []events.Event
	subscribercheck:
		for {
			select {
			case env := <-subscriber.eventq:
				ev, err := typeurl.UnmarshalAny(env.Event)
				if err != nil {
					t.Fatal(err)
				}
				received = append(received, ev)
			case err := <-subscriber.errq:
				if err != nil {
					t.Fatal(err)
				}
				break subscribercheck
			}

			if reflect.DeepEqual(received, subscriber.expectEvents) {
				// when we do this, we expect the errs channel to be closed and
				// this will return.
				subscriber.cancel()
			}
		}
	}
}

func TestExchangeValidateTopic(t *testing.T) {
	namespace := t.Name()
	ctx := namespaces.WithNamespace(context.Background(), namespace)
	exchange := NewExchange()

	for _, testcase := range []struct {
		input string
		err   error
	}{
		{
			input: "/test",
		},
		{
			input: "/test/test",
		},
		{
			input: "test",
			err:   errdefs.ErrInvalidArgument,
		},
	} {
		t.Run(testcase.input, func(t *testing.T) {
			event := &eventstypes.ContainerCreate{ID: t.Name()}
			if err := exchange.Publish(ctx, testcase.input, event); errors.Cause(err) != testcase.err {
				if err == nil {
					t.Fatalf("expected error %v, received nil", testcase.err)
				} else {
					t.Fatalf("expected error %v, received %v", testcase.err, err)
				}
			}

			evany, err := typeurl.MarshalAny(event)
			if err != nil {
				t.Fatal(err)
			}

			envelope := events.Envelope{
				Timestamp: time.Now().UTC(),
				Namespace: namespace,
				Topic:     testcase.input,
				Event:     evany,
			}

			// make sure we get same errors with forward.
			if err := exchange.Forward(ctx, &envelope); errors.Cause(err) != testcase.err {
				if err == nil {
					t.Fatalf("expected error %v, received nil", testcase.err)
				} else {
					t.Fatalf("expected error %v, received %v", testcase.err, err)
				}
			}

		})
	}
}
<|endoftext|>


<commit_before>package printer

import (
	"bytes"
	"fmt"
	"go/parser"
	"go/printer"
	"go/token"
	"text/template"

	"github.com/davyxu/tabtoy/v2/i18n"
	"github.com/davyxu/tabtoy/v2/model"
)

const goTemplate = `// Generated by github.com/davyxu/tabtoy
// Version: {{.ToolVersion}}
// DO NOT EDIT!!
package {{.Package}}

import(
	{{if .HasAnyIndex}}"fmt"{{end}}
	"encoding/json"
	"io/ioutil"
)
{{range $a, $en := .Enums}}
// Defined in table: {{$en.DefinedTable}}
type {{$en.Name}} int32
const (
{{range .GoFields}}
	{{.Comment}}
	{{$en.Name}}_{{.Name}} {{$en.Name}} = {{.Number}}
{{end}}
)
var (
{{$en.Name}}MapperValueByName = map[string]int32{ {{range .GoFields}}
	"{{.Name}}": {{.Number}}, {{end}}
}

{{$en.Name}}MapperNameByValue = map[int32]string{ {{range .GoFields}}
	{{.Number}}: "{{.Name}}" , {{end}}
}
)

func (self {{$en.Name}}) String() string {
	name, _ := {{$en.Name}}MapperNameByValue[int32(self)]
	return name
}
{{end}}

{{range $a, $strus := .Structs}}
// Defined in table: {{$strus.DefinedTable}}
type {{$strus.Name}} struct{
	{{range $b, $fd := $strus.GoFields}}
	{{.Comment}}
	{{$fd.Name}} {{$fd.TypeString}} {{$fd.StructTag}}
	{{end}}
}
{{end}}

// {{$.Name}} 访问接口
type {{$.Name}}Table struct{

	// 表格原始数据
	{{$.Name}}

	// 索引函数表
	indexFuncByName map[string][]func(*{{$.Name}}Table) error

	// 清空函数表
	clearFuncByName map[string][]func(*{{$.Name}}Table) error

	// 加载前回调
	preFuncList []func(*{{$.Name}}Table) error

	// 加载后回调
	postFuncList []func(*{{$.Name}}Table) error

	{{range $a, $strus := .IndexedStructs}} {{range .Indexes}}
	{{$strus.Name}}By{{.Name}} map[{{.KeyType}}]*{{$strus.TypeName}}
	{{end}} {{end}}
}

{{range .VerticalFields}}
{{.Comment}}
func (self *{{$.Name}}Table) Get{{.Name}}( ) {{.ElementTypeString}} {
	return self.{{.Name}}[0]
}
{{end}}

// 从json文件加载
func (self *{{$.Name}}Table) Load(filename string) error {

	data, err := ioutil.ReadFile(filename)

	if err != nil {
		return err
	}

	var newTab {{$.Name}}

	// 读取
	err = json.Unmarshal(data, &newTab)
	if err != nil {
		return err
	}

	// 所有加载前的回调
	for _, v := range self.preFuncList {
		if err = v(self); err != nil {
			return err
		}
	}

	// 清除前通知
	for _, list := range self.clearFuncByName {
		for _, v := range list {
			if err = v(self); err != nil {
				return err
			}
		}
	}

	// 复制数据
	self.{{$.Name}} = newTab

	// 生成索引
	for _, list := range self.indexFuncByName {
		for _, v := range list {
			if err = v(self); err != nil {
				return err
			}
		}
	}

	// 所有完成时的回调
	for _, v := range self.postFuncList {
		if err = v(self); err != nil {
			return err
		}
	}

	return nil
}

// 注册外部索引入口, 索引回调, 清空回调
func (self *{{$.Name}}Table) RegisterIndexEntry(name string, indexCallback func(*{{$.Name}}Table) error, clearCallback func(*{{$.Name}}Table)error) {

	indexList, _ := self.indexFuncByName[name]
	clearList, _ := self.clearFuncByName[name]

	if indexCallback != nil {
		indexList = append(indexList, indexCallback)
	}

	if clearCallback != nil {
		clearList = append(clearList, clearCallback)
	}

	self.indexFuncByName[name] = indexList
	self.clearFuncByName[name] = clearList
}

// 注册加载前回调
func (self *{{$.Name}}Table) RegisterPreEntry(callback func(*{{$.Name}}Table) error) {

	self.preFuncList = append(self.preFuncList, callback)
}

// 注册所有完成时回调
func (self *{{$.Name}}Table) RegisterPostEntry(callback func(*{{$.Name}}Table) error) {

	self.postFuncList = append(self.postFuncList, callback)
}

// 创建一个{{$.Name}}表读取实例
func New{{$.Name}}Table() *{{$.Name}}Table {
	return &{{$.Name}}Table{

		indexFuncByName: map[string][]func(*{{$.Name}}Table) error{

		{{range $a, $strus := .IndexedStructs}}
			"{{$strus.Name}}": {func(tab *{{$.Name}}Table)error {

				// {{$strus.Name}}
				for _, def := range tab.{{$strus.Name}} {
					{{range .Indexes}}
					if _, ok := tab.{{$strus.Name}}By{{.Name}}[def.{{.Name}}]; ok {
						panic(fmt.Sprintf("duplicate index in {{$strus.Name}}By{{.Name}}: %v", def.{{.Name}}))
					}
					{{end}}
					{{range .Indexes}}
					tab.{{$strus.Name}}By{{.Name}}[def.{{.Name}}] = def{{end}}

				}

				return nil
			}},
		{{end}}

		},

		clearFuncByName: map[string][]func(*{{$.Name}}Table)error{

		{{range $a, $strus := .IndexedStructs}}
			"{{$strus.Name}}": {func(tab *{{$.Name}}Table) error{

				// {{$strus.Name}}

				{{range .Indexes}}
				tab.{{$strus.Name}}By{{.Name}} = make(map[{{.KeyType}}]*{{$strus.TypeName}}){{end}}

				return nil
			}},
		{{end}}

		},

		{{range $a, $strus := .IndexedStructs}} {{range .Indexes}}
		{{$strus.Name}}By{{.Name}} : make(map[{{.KeyType}}]*{{$strus.TypeName}}),
		{{end}} {{end}}

	}
}

`

// 每个带有MakeIndex的列
type goFieldModel struct {
	*model.FieldDescriptor
	Number int
}

func (self goFieldModel) Alias() string {
	return self.FieldDescriptor.Meta.GetString("Alias")
}

func (self goFieldModel) Comment() string {

	var out string

	if self.FieldDescriptor.Meta.GetString("Alias") != "" {
		out += "// "
		out += self.FieldDescriptor.Meta.GetString("Alias")
	}

	if self.FieldDescriptor.Comment != "" {
		if out == "" {
			out += "//"
		}
		out += self.FieldDescriptor.Comment
	}

	return out

}

func (self *goFieldModel) KeyType() string {

	// 修复: 当枚举做索引时, 多出包名
	if self.Type == model.FieldType_Enum {
		return self.Complex.Name
	}

	return model.FieldTypeToString(self.Type)
}

func (self *goFieldModel) StructTag() string {

	var buf bytes.Buffer

	buf.WriteString("`")

	var userTagCount int

	self.Meta.VisitUserMeta(func(k string, v interface{}) bool {

		if userTagCount > 0 {
			buf.WriteString(" ")
		}

		buf.WriteString(fmt.Sprintf("%s:\"%s\"", k, v))

		userTagCount++

		return true
	})

	buf.WriteString("`")

	if userTagCount == 0 {
		return ""
	}

	return buf.String()
}

func (self *goFieldModel) ElementTypeString() string {

	switch self.FieldDescriptor.Type {
	case model.FieldType_Float:
		return "float32"
	case model.FieldType_Struct:
		return "*" + self.FieldDescriptor.TypeString()
	default:
		return self.FieldDescriptor.TypeString()
	}
}

func (self *goFieldModel) TypeString() string {

	var prefix string
	if self.IsRepeated {
		prefix = "[]"
	}

	return prefix + self.ElementTypeString()

}

// 对应每个电子表格文件
type goIndexStructModel struct {
	*model.FieldDescriptor

	Indexes []*goFieldModel
}

func (self *goIndexStructModel) TypeName() string {
	return self.Complex.Name
}

type goStructModel struct {
	*model.Descriptor

	GoFields []*goFieldModel
}

func (self *goStructModel) DefinedTable() string {
	return self.File.Name
}

// 整个导出文件
type goFileModel struct {
	*model.FileDescriptor
	ToolVersion    string
	IndexedStructs []*goIndexStructModel
	Structs        []*goStructModel
	Enums          []*goStructModel
	IndexCount     int

	// 配置的字段
	VerticalFields []*goFieldModel
}

func (self *goFileModel) HasAnyIndex() bool {
	return self.IndexCount > 0
}

func (self *goFileModel) HasAnyStruct() bool {
	return len(self.Structs) > 0
}

func (self *goFileModel) Package() string {
	return self.FileDescriptor.Pragma.GetString("Package")
}

type goPrinter struct {
}

func collectIndexInfo(g *Globals, fm *goFileModel) {

	// 遍历需要导出的合并字段
	for _, fd := range g.CombineStruct.Fields {

		// fd --> 对应每个文件的Row定义, 也就是XXDefine, 在CombineStruct上, 只是一个字段

		// 对CombineStruct的XXDefine对应的字段
		if g.CombineStruct.Usage != model.DescriptorUsage_CombineStruct {
			continue
		}

		//非结构体不输出(其实根本不会击中)
		if fd.Complex == nil {
			continue
		}

		// 这个字段被限制输出
		if !fd.Complex.File.MatchTag(".go") {
			continue
		}

		// 这个结构有索引才创建
		if len(fd.Complex.Indexes) == 0 {
			continue
		}

		rm := goIndexStructModel{FieldDescriptor: fd}

		// 索引字段
		for _, key := range fd.Complex.Indexes {

			rm.Indexes = append(rm.Indexes, &goFieldModel{
				FieldDescriptor: key,
			})
			fm.IndexCount++
		}

		fm.IndexedStructs = append(fm.IndexedStructs, &rm)

	}
}

func collectAllStructInfo(g *Globals, fm *goFileModel) {

	// 遍历所有类型
	for _, d := range g.FileDescriptor.Descriptors {

		// 这给被限制输出
		if !d.File.MatchTag(".go") {
			log.Infof("%s: %s", i18n.String(i18n.Printer_IgnoredByOutputTag), d.Name)
			continue
		}

		structM := &goStructModel{Descriptor: d}

		// 遍历字段
		for index, fd := range d.Fields {

			// 对CombineStruct的XXDefine对应的字段
			if d.Usage == model.DescriptorUsage_CombineStruct {

				// 这个字段被限制输出
				if fd.Complex != nil && !fd.Complex.File.MatchTag(".go") {
					continue
				}

				if fd.Complex != nil && fd.Complex.File.Pragma.GetBool("Vertical") {
					fm.VerticalFields = append(fm.VerticalFields, &goFieldModel{FieldDescriptor: fd})
				}
			}

			field := &goFieldModel{FieldDescriptor: fd}

			switch d.Kind {
			case model.DescriptorKind_Struct:
				field.Number = index + 1
			case model.DescriptorKind_Enum:
				field.Number = int(fd.EnumValue)
			}

			structM.GoFields = append(structM.GoFields, field)

		}

		switch d.Kind {
		case model.DescriptorKind_Struct:
			fm.Structs = append(fm.Structs, structM)
		case model.DescriptorKind_Enum:
			fm.Enums = append(fm.Enums, structM)
		}

	}
}

func (self *goPrinter) Run(g *Globals) *Stream {

	tpl, err := template.New("golang").Parse(goTemplate)
	if err != nil {
		log.Errorln(err)
		return nil
	}

	var fm goFileModel
	fm.ToolVersion = g.Version
	fm.FileDescriptor = g.FileDescriptor

	collectIndexInfo(g, &fm)
	collectAllStructInfo(g, &fm)

	bf := NewStream()

	err = tpl.Execute(bf.Buffer(), &fm)
	if err != nil {
		log.Errorln(err)
		return nil
	}

	if err := formatCode(bf.Buffer()); err != nil {
		log.Errorln("format golang code err", err)
		log.Errorln(bf.Buffer().String())
		return nil
	}

	return bf
}

func formatCode(bf *bytes.Buffer) error {

	fset := token.NewFileSet()

	ast, err := parser.ParseFile(fset, "", bf, parser.ParseComments)
	if err != nil {
		return err
	}

	bf.Reset()

	err = (&printer.Config{Mode: printer.TabIndent | printer.UseSpaces, Tabwidth: 8}).Fprint(bf, fset, ast)
	if err != nil {
		return err
	}

	return nil
}

func init() {

	RegisterPrinter("go", &goPrinter{})

}
<commit_msg>add golang代码生成增加从数据加载<commit_after>package printer

import (
	"bytes"
	"fmt"
	"go/parser"
	"go/printer"
	"go/token"
	"text/template"

	"github.com/davyxu/tabtoy/v2/i18n"
	"github.com/davyxu/tabtoy/v2/model"
)

const goTemplate = `// Generated by github.com/davyxu/tabtoy
// Version: {{.ToolVersion}}
// DO NOT EDIT!!
package {{.Package}}

import(
	{{if .HasAnyIndex}}"fmt"{{end}}
	"encoding/json"
	"io/ioutil"
)
{{range $a, $en := .Enums}}
// Defined in table: {{$en.DefinedTable}}
type {{$en.Name}} int32
const (
{{range .GoFields}}
	{{.Comment}}
	{{$en.Name}}_{{.Name}} {{$en.Name}} = {{.Number}}
{{end}}
)
var (
{{$en.Name}}MapperValueByName = map[string]int32{ {{range .GoFields}}
	"{{.Name}}": {{.Number}}, {{end}}
}

{{$en.Name}}MapperNameByValue = map[int32]string{ {{range .GoFields}}
	{{.Number}}: "{{.Name}}" , {{end}}
}
)

func (self {{$en.Name}}) String() string {
	name, _ := {{$en.Name}}MapperNameByValue[int32(self)]
	return name
}
{{end}}

{{range $a, $strus := .Structs}}
// Defined in table: {{$strus.DefinedTable}}
type {{$strus.Name}} struct{
	{{range $b, $fd := $strus.GoFields}}
	{{.Comment}}
	{{$fd.Name}} {{$fd.TypeString}} {{$fd.StructTag}}
	{{end}}
}
{{end}}

// {{$.Name}} 访问接口
type {{$.Name}}Table struct{

	// 表格原始数据
	{{$.Name}}

	// 索引函数表
	indexFuncByName map[string][]func(*{{$.Name}}Table) error

	// 清空函数表
	clearFuncByName map[string][]func(*{{$.Name}}Table) error

	// 加载前回调
	preFuncList []func(*{{$.Name}}Table) error

	// 加载后回调
	postFuncList []func(*{{$.Name}}Table) error

	{{range $a, $strus := .IndexedStructs}} {{range .Indexes}}
	{{$strus.Name}}By{{.Name}} map[{{.KeyType}}]*{{$strus.TypeName}}
	{{end}} {{end}}
}

{{range .VerticalFields}}
{{.Comment}}
func (self *{{$.Name}}Table) Get{{.Name}}( ) {{.ElementTypeString}} {
	return self.{{.Name}}[0]
}
{{end}}

// 从json文件加载
func (self *{{$.Name}}Table) Load(filename string) error {
	data, err := ioutil.ReadFile(filename)

	if err != nil {
		return err
	}

	return self.LoadData(data)
}

// 从二进制加载
func (self *{{$.Name}}Table) LoadData(data []byte) error {

	var newTab {{$.Name}}

	// 读取
	err := json.Unmarshal(data, &newTab)
	if err != nil {
		return err
	}

	// 所有加载前的回调
	for _, v := range self.preFuncList {
		if err = v(self); err != nil {
			return err
		}
	}

	// 清除前通知
	for _, list := range self.clearFuncByName {
		for _, v := range list {
			if err = v(self); err != nil {
				return err
			}
		}
	}

	// 复制数据
	self.{{$.Name}} = newTab

	// 生成索引
	for _, list := range self.indexFuncByName {
		for _, v := range list {
			if err = v(self); err != nil {
				return err
			}
		}
	}

	// 所有完成时的回调
	for _, v := range self.postFuncList {
		if err = v(self); err != nil {
			return err
		}
	}

	return nil
}

// 注册外部索引入口, 索引回调, 清空回调
func (self *{{$.Name}}Table) RegisterIndexEntry(name string, indexCallback func(*{{$.Name}}Table) error, clearCallback func(*{{$.Name}}Table)error) {

	indexList, _ := self.indexFuncByName[name]
	clearList, _ := self.clearFuncByName[name]

	if indexCallback != nil {
		indexList = append(indexList, indexCallback)
	}

	if clearCallback != nil {
		clearList = append(clearList, clearCallback)
	}

	self.indexFuncByName[name] = indexList
	self.clearFuncByName[name] = clearList
}

// 注册加载前回调
func (self *{{$.Name}}Table) RegisterPreEntry(callback func(*{{$.Name}}Table) error) {

	self.preFuncList = append(self.preFuncList, callback)
}

// 注册所有完成时回调
func (self *{{$.Name}}Table) RegisterPostEntry(callback func(*{{$.Name}}Table) error) {

	self.postFuncList = append(self.postFuncList, callback)
}

// 创建一个{{$.Name}}表读取实例
func New{{$.Name}}Table() *{{$.Name}}Table {
	return &{{$.Name}}Table{

		indexFuncByName: map[string][]func(*{{$.Name}}Table) error{

		{{range $a, $strus := .IndexedStructs}}
			"{{$strus.Name}}": {func(tab *{{$.Name}}Table)error {

				// {{$strus.Name}}
				for _, def := range tab.{{$strus.Name}} {
					{{range .Indexes}}
					if _, ok := tab.{{$strus.Name}}By{{.Name}}[def.{{.Name}}]; ok {
						panic(fmt.Sprintf("duplicate index in {{$strus.Name}}By{{.Name}}: %v", def.{{.Name}}))
					}
					{{end}}
					{{range .Indexes}}
					tab.{{$strus.Name}}By{{.Name}}[def.{{.Name}}] = def{{end}}

				}

				return nil
			}},
		{{end}}

		},

		clearFuncByName: map[string][]func(*{{$.Name}}Table)error{

		{{range $a, $strus := .IndexedStructs}}
			"{{$strus.Name}}": {func(tab *{{$.Name}}Table) error{

				// {{$strus.Name}}

				{{range .Indexes}}
				tab.{{$strus.Name}}By{{.Name}} = make(map[{{.KeyType}}]*{{$strus.TypeName}}){{end}}

				return nil
			}},
		{{end}}

		},

		{{range $a, $strus := .IndexedStructs}} {{range .Indexes}}
		{{$strus.Name}}By{{.Name}} : make(map[{{.KeyType}}]*{{$strus.TypeName}}),
		{{end}} {{end}}

	}
}

`

// 每个带有MakeIndex的列
type goFieldModel struct {
	*model.FieldDescriptor
	Number int
}

func (self goFieldModel) Alias() string {
	return self.FieldDescriptor.Meta.GetString("Alias")
}

func (self goFieldModel) Comment() string {

	var out string

	if self.FieldDescriptor.Meta.GetString("Alias") != "" {
		out += "// "
		out += self.FieldDescriptor.Meta.GetString("Alias")
	}

	if self.FieldDescriptor.Comment != "" {
		if out == "" {
			out += "//"
		}
		out += self.FieldDescriptor.Comment
	}

	return out

}

func (self *goFieldModel) KeyType() string {

	// 修复: 当枚举做索引时, 多出包名
	if self.Type == model.FieldType_Enum {
		return self.Complex.Name
	}

	return model.FieldTypeToString(self.Type)
}

func (self *goFieldModel) StructTag() string {

	var buf bytes.Buffer

	buf.WriteString("`")

	var userTagCount int

	self.Meta.VisitUserMeta(func(k string, v interface{}) bool {

		if userTagCount > 0 {
			buf.WriteString(" ")
		}

		buf.WriteString(fmt.Sprintf("%s:\"%s\"", k, v))

		userTagCount++

		return true
	})

	buf.WriteString("`")

	if userTagCount == 0 {
		return ""
	}

	return buf.String()
}

func (self *goFieldModel) ElementTypeString() string {

	switch self.FieldDescriptor.Type {
	case model.FieldType_Float:
		return "float32"
	case model.FieldType_Struct:
		return "*" + self.FieldDescriptor.TypeString()
	default:
		return self.FieldDescriptor.TypeString()
	}
}

func (self *goFieldModel) TypeString() string {

	var prefix string
	if self.IsRepeated {
		prefix = "[]"
	}

	return prefix + self.ElementTypeString()

}

// 对应每个电子表格文件
type goIndexStructModel struct {
	*model.FieldDescriptor

	Indexes []*goFieldModel
}

func (self *goIndexStructModel) TypeName() string {
	return self.Complex.Name
}

type goStructModel struct {
	*model.Descriptor

	GoFields []*goFieldModel
}

func (self *goStructModel) DefinedTable() string {
	return self.File.Name
}

// 整个导出文件
type goFileModel struct {
	*model.FileDescriptor
	ToolVersion    string
	IndexedStructs []*goIndexStructModel
	Structs        []*goStructModel
	Enums          []*goStructModel
	IndexCount     int

	// 配置的字段
	VerticalFields []*goFieldModel
}

func (self *goFileModel) HasAnyIndex() bool {
	return self.IndexCount > 0
}

func (self *goFileModel) HasAnyStruct() bool {
	return len(self.Structs) > 0
}

func (self *goFileModel) Package() string {
	return self.FileDescriptor.Pragma.GetString("Package")
}

type goPrinter struct {
}

func collectIndexInfo(g *Globals, fm *goFileModel) {

	// 遍历需要导出的合并字段
	for _, fd := range g.CombineStruct.Fields {

		// fd --> 对应每个文件的Row定义, 也就是XXDefine, 在CombineStruct上, 只是一个字段

		// 对CombineStruct的XXDefine对应的字段
		if g.CombineStruct.Usage != model.DescriptorUsage_CombineStruct {
			continue
		}

		//非结构体不输出(其实根本不会击中)
		if fd.Complex == nil {
			continue
		}

		// 这个字段被限制输出
		if !fd.Complex.File.MatchTag(".go") {
			continue
		}

		// 这个结构有索引才创建
		if len(fd.Complex.Indexes) == 0 {
			continue
		}

		rm := goIndexStructModel{FieldDescriptor: fd}

		// 索引字段
		for _, key := range fd.Complex.Indexes {

			rm.Indexes = append(rm.Indexes, &goFieldModel{
				FieldDescriptor: key,
			})
			fm.IndexCount++
		}

		fm.IndexedStructs = append(fm.IndexedStructs, &rm)

	}
}

func collectAllStructInfo(g *Globals, fm *goFileModel) {

	// 遍历所有类型
	for _, d := range g.FileDescriptor.Descriptors {

		// 这给被限制输出
		if !d.File.MatchTag(".go") {
			log.Infof("%s: %s", i18n.String(i18n.Printer_IgnoredByOutputTag), d.Name)
			continue
		}

		structM := &goStructModel{Descriptor: d}

		// 遍历字段
		for index, fd := range d.Fields {

			// 对CombineStruct的XXDefine对应的字段
			if d.Usage == model.DescriptorUsage_CombineStruct {

				// 这个字段被限制输出
				if fd.Complex != nil && !fd.Complex.File.MatchTag(".go") {
					continue
				}

				if fd.Complex != nil && fd.Complex.File.Pragma.GetBool("Vertical") {
					fm.VerticalFields = append(fm.VerticalFields, &goFieldModel{FieldDescriptor: fd})
				}
			}

			field := &goFieldModel{FieldDescriptor: fd}

			switch d.Kind {
			case model.DescriptorKind_Struct:
				field.Number = index + 1
			case model.DescriptorKind_Enum:
				field.Number = int(fd.EnumValue)
			}

			structM.GoFields = append(structM.GoFields, field)

		}

		switch d.Kind {
		case model.DescriptorKind_Struct:
			fm.Structs = append(fm.Structs, structM)
		case model.DescriptorKind_Enum:
			fm.Enums = append(fm.Enums, structM)
		}

	}
}

func (self *goPrinter) Run(g *Globals) *Stream {

	tpl, err := template.New("golang").Parse(goTemplate)
	if err != nil {
		log.Errorln(err)
		return nil
	}

	var fm goFileModel
	fm.ToolVersion = g.Version
	fm.FileDescriptor = g.FileDescriptor

	collectIndexInfo(g, &fm)
	collectAllStructInfo(g, &fm)

	bf := NewStream()

	err = tpl.Execute(bf.Buffer(), &fm)
	if err != nil {
		log.Errorln(err)
		return nil
	}

	if err := formatCode(bf.Buffer()); err != nil {
		log.Errorln("format golang code err", err)
		log.Errorln(bf.Buffer().String())
		return nil
	}

	return bf
}

func formatCode(bf *bytes.Buffer) error {

	fset := token.NewFileSet()

	ast, err := parser.ParseFile(fset, "", bf, parser.ParseComments)
	if err != nil {
		return err
	}

	bf.Reset()

	err = (&printer.Config{Mode: printer.TabIndent | printer.UseSpaces, Tabwidth: 8}).Fprint(bf, fset, ast)
	if err != nil {
		return err
	}

	return nil
}

func init() {

	RegisterPrinter("go", &goPrinter{})

}
<|endoftext|>


<commit_before>package docker

import (
	"bufio"
	"context"
	"encoding/json"
	"github.com/ViBiOh/dashboard/auth"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"regexp"
)

const ignoredByteLogSize = 8
const tailSize = `100`

var logWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/logs`)
var statsWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/stats`)
var eventsWebsocketRequest = regexp.MustCompile(`events`)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return hostCheck.MatchString(r.Host)
	},
}

func upgradeAndAuth(w http.ResponseWriter, r *http.Request) (*websocket.Conn, *auth.User, error) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		defer ws.Close()
		return nil, nil, err
	}

	_, basicAuth, err := ws.ReadMessage()
	if err != nil {
		defer ws.Close()
		return nil, nil, err
	}

	user, err := auth.IsAuthenticatedByAuth(string(basicAuth))
	if err != nil {
		ws.WriteMessage(websocket.TextMessage, []byte(err.Error()))

		defer ws.Close()
		return nil, nil, err
	}

	return ws, user, nil
}

func logsContainerWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, _, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	logs, err := docker.ContainerLogs(ctx, string(containerID), types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true, Follow: true, Tail: tailSize})
	if err != nil {
		log.Print(err)
		return
	}
	defer logs.Close()

	go func() {
		scanner := bufio.NewScanner(logs)
		for scanner.Scan() {
			select {
			case <-ctx.Done():
				return

			default:
				logLine := scanner.Bytes()
				if len(logLine) > ignoredByteLogSize {
					if err = ws.WriteMessage(websocket.TextMessage, logLine[ignoredByteLogSize:]); err != nil {
						log.Printf(`Error while writing to logs socket: %v`, err)
						cancel()
						return
					}
				}
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			return

		default:
			if messageType, content, err := ws.NextReader(); err != nil {
				log.Printf(`Error while reading from logs socket: %v | %v | %v`, err, messageType, content)
				cancel()
				return
			}
		}
	}
}

func eventsWebsocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		return
	}
	defer ws.Close()

	filtersArgs := filters.NewArgs()
	if labelFilters(&filtersArgs, user, nil) != nil {
		log.Printf(`Error while defining label filters: %v`, err)
		return
	}
	if eventFilters(&filtersArgs) != nil {
		log.Printf(`Error while defining event filters: %v`, err)
		return
	}

	context := context.Background()
	defer context.Done()
	messages, errors := docker.Events(context, types.EventsOptions{Filters: filtersArgs})

	done := make(chan struct{})

	go func() {
		for {
			select {
			case <-done:
				return

			case message := <-messages:
				messageJSON, err := json.Marshal(message)
				if err != nil {
					log.Printf(`Error while marshalling event: %v`, err)
					close(done)
					return
				}

				if err = ws.WriteMessage(websocket.TextMessage, messageJSON); err != nil {
					log.Printf(`Error while writing to events socket: %v`, err)
					close(done)
					return
				}
				break

			case err := <-errors:
				log.Printf(`Error while reading events: %v`, err)
				close(done)
				return
			}
		}
	}()

	for {
		select {
		case <-done:
			return
		default:
			if _, _, err := ws.NextReader(); err != nil {
				log.Printf(`Error while reading from events socket: %v`, err)
				close(done)
				return
			}
		}
	}
}

func statsWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, _, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	context := context.Background()
	defer context.Done()

	stats, err := docker.ContainerStats(context, string(containerID), true)
	if err != nil {
		log.Print(err)
		return
	}
	defer stats.Body.Close()

	done := make(chan struct{})

	go func() {
		scanner := bufio.NewScanner(stats.Body)
		for scanner.Scan() {
			select {
			case <-done:
				return

			default:
				if err = ws.WriteMessage(websocket.TextMessage, scanner.Bytes()); err != nil {
					log.Printf(`Error while writing to stats socket: %v`, err)
					close(done)
					return
				}
			}
		}
	}()

	for {
		select {
		case <-done:
			return

		default:
			if _, _, err := ws.NextReader(); err != nil {
				log.Printf(`Error while reading from stats socket: %v`, err)
				close(done)
				return
			}
		}
	}
}

// WebsocketHandler for Docker Websocket request. Should be use with net/http
type WebsocketHandler struct {
}

func (handler WebsocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	urlPath := []byte(r.URL.Path)

	if logWebsocketRequest.Match(urlPath) {
		logsContainerWebsocketHandler(w, r, logWebsocketRequest.FindSubmatch(urlPath)[1])
	} else if eventsWebsocketRequest.Match((urlPath)) {
		eventsWebsocketHandler(w, r)
	} else if statsWebsocketRequest.Match(urlPath) {
		statsWebsocketHandler(w, r, statsWebsocketRequest.FindSubmatch(urlPath)[1])
	}
}
<commit_msg>Update websocket.go<commit_after>package docker

import (
	"bufio"
	"context"
	"encoding/json"
	"github.com/ViBiOh/dashboard/auth"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"regexp"
)

const ignoredByteLogSize = 8
const tailSize = `100`

var logWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/logs`)
var statsWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/stats`)
var eventsWebsocketRequest = regexp.MustCompile(`events`)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return hostCheck.MatchString(r.Host)
	},
}

func upgradeAndAuth(w http.ResponseWriter, r *http.Request) (*websocket.Conn, *auth.User, error) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		defer ws.Close()
		return nil, nil, err
	}

	_, basicAuth, err := ws.ReadMessage()
	if err != nil {
		defer ws.Close()
		return nil, nil, err
	}

	user, err := auth.IsAuthenticatedByAuth(string(basicAuth))
	if err != nil {
		ws.WriteMessage(websocket.TextMessage, []byte(err.Error()))

		defer ws.Close()
		return nil, nil, err
	}

	return ws, user, nil
}

func logsContainerWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, _, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	logs, err := docker.ContainerLogs(ctx, string(containerID), types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true, Follow: true, Tail: tailSize})
	if err != nil {
		log.Print(err)
		return
	}
	defer logs.Close()

	go func() {
		defer cancel()

		scanner := bufio.NewScanner(logs)
		for scanner.Scan() {
			select {
			case <-ctx.Done():
				return

			default:
				logLine := scanner.Bytes()
				if len(logLine) > ignoredByteLogSize {
					if err = ws.WriteMessage(websocket.TextMessage, logLine[ignoredByteLogSize:]); err != nil {
						log.Printf(`Error while writing to logs socket: %v`, err)
						return
					}
				}
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			return

		default:
			if messageType, content, err := ws.NextReader(); err != nil {
				log.Printf(`Error while reading from logs socket: %v | %v | %v`, err, messageType, content)
				return
			}
		}
	}
}

func eventsWebsocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		return
	}
	defer ws.Close()

	filtersArgs := filters.NewArgs()
	if labelFilters(&filtersArgs, user, nil) != nil {
		log.Printf(`Error while defining label filters: %v`, err)
		return
	}
	if eventFilters(&filtersArgs) != nil {
		log.Printf(`Error while defining event filters: %v`, err)
		return
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	messages, errors := docker.Events(ctx, types.EventsOptions{Filters: filtersArgs})

	go func() {
		defer cancel()

		for {
			select {
			case <-ctx.Done():
				return

			case message := <-messages:
				messageJSON, err := json.Marshal(message)
				if err != nil {
					log.Printf(`Error while marshalling event: %v`, err)
					return
				}

				if err = ws.WriteMessage(websocket.TextMessage, messageJSON); err != nil {
					log.Printf(`Error while writing to events socket: %v`, err)
					return
				}

			case err := <-errors:
				log.Printf(`Error while reading events: %v`, err)
				return
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			return
		default:
			if _, _, err := ws.NextReader(); err != nil {
				log.Printf(`Error while reading from events socket: %v`, err)
				return
			}
		}
	}
}

func statsWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, _, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}
	defer ws.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	stats, err := docker.ContainerStats(ctx, string(containerID), true)
	if err != nil {
		log.Print(err)
		return
	}
	defer stats.Body.Close()

	go func() {
		defer cancel()

		scanner := bufio.NewScanner(stats.Body)
		for scanner.Scan() {
			select {
			case <-ctx.Done():
				return

			default:
				if err = ws.WriteMessage(websocket.TextMessage, scanner.Bytes()); err != nil {
					log.Printf(`Error while writing to stats socket: %v`, err)
					return
				}
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			return

		default:
			if _, _, err := ws.NextReader(); err != nil {
				log.Printf(`Error while reading from stats socket: %v`, err)
				return
			}
		}
	}
}

// WebsocketHandler for Docker Websocket request. Should be use with net/http
type WebsocketHandler struct {
}

func (handler WebsocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	urlPath := []byte(r.URL.Path)

	if logWebsocketRequest.Match(urlPath) {
		logsContainerWebsocketHandler(w, r, logWebsocketRequest.FindSubmatch(urlPath)[1])
	} else if eventsWebsocketRequest.Match((urlPath)) {
		eventsWebsocketHandler(w, r)
	} else if statsWebsocketRequest.Match(urlPath) {
		statsWebsocketHandler(w, r, statsWebsocketRequest.FindSubmatch(urlPath)[1])
	}
}
<|endoftext|>