<commit_before>package main

import (
	"fmt"
	"github.com/hsoj/asn1-ber"
	"io"
	"koding/tools/db"
	"koding/tools/lifecycle"
	"koding/tools/log"
	"koding/virt"
	"labix.org/v2/mgo/bson"
	"net"
	"strconv"
	"strings"
)

func main() {
	lifecycle.Startup("ldapserver", true)

	ln, err := net.Listen("tcp", ":389")
	if err != nil {
		fmt.Println(err)
		return
	}
	for {
		conn, err := ln.Accept()
		if err != nil {
			fmt.Println(err)
			continue
		}
		go handleConnection(conn)
	}
}

func handleConnection(conn net.Conn) {
	defer log.RecoverAndLog()

	bound := false
	var vm *virt.VM

	for {
		packet, err := ber.ReadPacket(conn)
		if err != nil {
			if err == io.EOF {
				break
			}
			panic(err)
		}

		messageID := packet.Children[0].Value.(uint64)
		request := packet.Children[1]

		switch request.Tag {
		case ApplicationBindRequest:
			name := request.Children[1].Value.(string)
			password := request.Children[2].Data.String()
			bound = authenticate(name, password)

			var resultCode uint64 = LDAPResultInvalidCredentials
			if bound {
				resultCode = LDAPResultSuccess
			}
			conn.Write(createSimpleResponse(messageID, ApplicationBindResponse, resultCode).Bytes())

		case ApplicationUnbindRequest:
			bound = false

		case ApplicationSearchRequest:
			if !bound {
				conn.Write(createSimpleResponse(messageID, ApplicationSearchResultDone, LDAPResultInsufficientAccessRights).Bytes())
				continue
			}

			if !lookupUser(request.Children[6], messageID, vm, conn) {
				// ber.PrintPacket(packet) // for debugging
			}

			conn.Write(createSimpleResponse(messageID, ApplicationSearchResultDone, LDAPResultSuccess).Bytes())

		default:
			panic("Unsupported LDAP command")
		}
	}
}

func authenticate(name, password string) bool {
	if name == "vmhost" && password == "abc" {
		return true
	}

	if strings.HasPrefix(name, "vm-") && bson.IsObjectIdHex(name[3:]) {
		var vm virt.VM
		err := db.VMs.FindId(bson.ObjectIdHex(name[3:])).One(&vm)
		return err == nil && password == vm.LdapPassword
	}

	user, err := findUserByName(name)
	return err == nil && user.HasPassword(password)
}

func lookupUser(filter *ber.Packet, messageID uint64, vm *virt.VM, conn net.Conn) bool {
	switch findAttributeInFilter(filter, "objectClass") {
	case "posixGroup":
		if memberUid := findAttributeInFilter(filter, "memberUid"); memberUid != "" {
			if vm == nil {
				return true
			}

			user, err := findUserByName(memberUid)
			if err != nil {
				return true
			}
			userEntry := vm.GetUserEntry(user)

			if userEntry != nil && userEntry.Sudo {
				conn.Write(createSearchResultEntry(messageID, map[string]string{
					"objectClass": "posixGroup",
					"cn":          "sudo",
					"gidNumber":   "27",
				}).Bytes())
			}
			return true
		}

		if gidStr := findAttributeInFilter(filter, "gidNumber"); gidStr != "" {
			gid, _ := strconv.Atoi(gidStr)
			user, err := findUserByUid(gid)
			if err != nil || (vm != nil && vm.GetUserEntry(user) == nil) {
				return true
			}

			conn.Write(createSearchResultEntry(messageID, map[string]string{
				"objectClass": "posixGroup",
				"cn":          user.Name,
				"gidNumber":   gidStr,
			}).Bytes())
			return true
		}

	default: // including "posixAccount"
		user, err := findUserInFilter(filter)
		if err == nil && user == nil {
			return false
		}
		if err != nil || (vm != nil && vm.GetUserEntry(user) == nil) {
			return true
		}

		conn.Write(createSearchResultEntry(messageID, map[string]string{
			"objectClass":   "posixAccount",
			"cn":            user.Name,
			"uid":           user.Name,
			"userPassword":  "{SSHA}MmhmXrqch9NbAcRA6Z79OTSj6MqNXQxF",
			"uidNumber":     strconv.Itoa(user.Uid),
			"gidNumber":     strconv.Itoa(user.Uid),
			"homeDirectory": "/home/" + user.Name,
			"loginShell":    "/bin/bash",
		}).Bytes())
		return true

	}

	return false
}

func findUserInFilter(filter *ber.Packet) (*virt.User, error) {
	if name := findAttributeInFilter(filter, "uid"); name != "" {
		return findUserByName(name)
	}
	if uidStr := findAttributeInFilter(filter, "uidNumber"); uidStr != "" {
		uid, _ := strconv.Atoi(uidStr)
		return findUserByUid(uid)
	}
	return nil, nil
}

func findUser(query interface{}) (*virt.User, error) {
	var user virt.User
	err := db.Users.Find(query).One(&user)
	if err == nil && user.Uid < virt.UserIdOffset {
		panic("User with too low uid.")
	}
	return &user, err
}

func findUserByUid(id int) (*virt.User, error) {
	return findUser(bson.M{"uid": id})
}

func findUserByName(name string) (*virt.User, error) {
	return findUser(bson.M{"username": name})
}

func findAttributeInFilter(filter *ber.Packet, name string) string {
	switch filter.Tag {
	case 0, 1: // and, or
		for _, child := range filter.Children {
			if result := findAttributeInFilter(child, name); result != "" {
				return result
			}
		}
	case 3: // equalityMatch
		if filter.Children[0].Value == name {
			return filter.Children[1].Value.(string)
		}
	}
	return ""
}

func createLDAPMessage(messageID uint64) *ber.Packet {
	packet := ber.Encode(ber.ClassUniversal, ber.TypeConstructed, ber.TagSequence, nil, "LDAPMessage")
	packet.AppendChild(ber.NewInteger(ber.ClassUniversal, ber.TypePrimative, ber.TagInteger, messageID, "MessageID"))
	return packet
}

func createLDAPResponse(tag uint8, resultCode uint64) *ber.Packet {
	packet := ber.Encode(ber.ClassApplication, ber.TypeConstructed, tag, nil, "LDAPResponse")
	packet.AppendChild(ber.NewInteger(ber.ClassUniversal, ber.TypePrimative, ber.TagInteger, resultCode, "resultCode"))
	packet.AppendChild(ber.NewString(ber.ClassUniversal, ber.TypePrimative, ber.TagOctetString, "", "matchedDN"))
	packet.AppendChild(ber.NewString(ber.ClassUniversal, ber.TypePrimative, ber.TagOctetString, "", "errorMessage"))
	return packet
}

func createSimpleResponse(messageID uint64, tag uint8, resultCode uint64) *ber.Packet {
	response := createLDAPMessage(messageID)
	response.AppendChild(createLDAPResponse(tag, resultCode))
	return response
}

func createSearchResultEntry(messageID uint64, attributes map[string]string) *ber.Packet {
	response := createLDAPMessage(messageID)
	entry := ber.Encode(ber.ClassApplication, ber.TypeConstructed, ApplicationSearchResultEntry, nil, "SearchResultEntry")
	entry.AppendChild(ber.NewString(ber.ClassUniversal, ber.TypePrimative, ber.TagOctetString, attributes["cn"], "objectName"))
	attributeSequence := ber.Encode(ber.ClassUniversal, ber.TypeConstructed, ber.TagSequence, nil, "attributes")
	for key, value := range attributes {
		attributeSequence.AppendChild(createAttribute(key, value))
	}
	entry.AppendChild(attributeSequence)
	response.AppendChild(entry)
	return response
}

func createAttribute(name, value string) *ber.Packet {
	attribute := ber.Encode(ber.ClassUniversal, ber.TypeConstructed, ber.TagSequence, nil, "attribute")
	attribute.AppendChild(ber.NewString(ber.ClassUniversal, ber.TypePrimative, ber.TagOctetString, name, "type"))
	values := ber.Encode(ber.ClassUniversal, ber.TypeConstructed, ber.TagSet, nil, "values")
	values.AppendChild(ber.NewString(ber.ClassUniversal, ber.TypePrimative, ber.TagOctetString, value, "value"))
	attribute.AppendChild(values)
	return attribute
}

const (
	ApplicationBindRequest           = 0
	ApplicationBindResponse          = 1
	ApplicationUnbindRequest         = 2
	ApplicationSearchRequest         = 3
	ApplicationSearchResultEntry     = 4
	ApplicationSearchResultDone      = 5
	ApplicationModifyRequest         = 6
	ApplicationModifyResponse        = 7
	ApplicationAddRequest            = 8
	ApplicationAddResponse           = 9
	ApplicationDelRequest            = 10
	ApplicationDelResponse           = 11
	ApplicationModifyDNRequest       = 12
	ApplicationModifyDNResponse      = 13
	ApplicationCompareRequest        = 14
	ApplicationCompareResponse       = 15
	ApplicationAbandonRequest        = 16
	ApplicationSearchResultReference = 19
	ApplicationExtendedRequest       = 23
	ApplicationExtendedResponse      = 24
)

const (
	LDAPResultSuccess                      = 0
	LDAPResultOperationsError              = 1
	LDAPResultProtocolError                = 2
	LDAPResultTimeLimitExceeded            = 3
	LDAPResultSizeLimitExceeded            = 4
	LDAPResultCompareFalse                 = 5
	LDAPResultCompareTrue                  = 6
	LDAPResultAuthMethodNotSupported       = 7
	LDAPResultStrongAuthRequired           = 8
	LDAPResultReferral                     = 10
	LDAPResultAdminLimitExceeded           = 11
	LDAPResultUnavailableCriticalExtension = 12
	LDAPResultConfidentialityRequired      = 13
	LDAPResultSaslBindInProgress           = 14
	LDAPResultNoSuchAttribute              = 16
	LDAPResultUndefinedAttributeType       = 17
	LDAPResultInappropriateMatching        = 18
	LDAPResultConstraintViolation          = 19
	LDAPResultAttributeOrValueExists       = 20
	LDAPResultInvalidAttributeSyntax       = 21
	LDAPResultNoSuchObject                 = 32
	LDAPResultAliasProblem                 = 33
	LDAPResultInvalidDNSyntax              = 34
	LDAPResultAliasDereferencingProblem    = 36
	LDAPResultInappropriateAuthentication  = 48
	LDAPResultInvalidCredentials           = 49
	LDAPResultInsufficientAccessRights     = 50
	LDAPResultBusy                         = 51
	LDAPResultUnavailable                  = 52
	LDAPResultUnwillingToPerform           = 53
	LDAPResultLoopDetect                   = 54
	LDAPResultNamingViolation              = 64
	LDAPResultObjectClassViolation         = 65
	LDAPResultNotAllowedOnNonLeaf          = 66
	LDAPResultNotAllowedOnRDN              = 67
	LDAPResultEntryAlreadyExists           = 68
	LDAPResultObjectClassModsProhibited    = 69
	LDAPResultAffectsMultipleDSAs          = 71
	LDAPResultOther                        = 80

	ErrorNetwork         = 200
	ErrorFilterCompile   = 201
	ErrorFilterDecompile = 202
	ErrorDebugging       = 203
)
<commit_msg>virt: Fixed sudo.<commit_after>package main

import (
	"fmt"
	"github.com/hsoj/asn1-ber"
	"io"
	"koding/tools/db"
	"koding/tools/lifecycle"
	"koding/tools/log"
	"koding/virt"
	"labix.org/v2/mgo/bson"
	"net"
	"strconv"
	"strings"
)

func main() {
	lifecycle.Startup("ldapserver", true)

	ln, err := net.Listen("tcp", ":389")
	if err != nil {
		fmt.Println(err)
		return
	}
	for {
		conn, err := ln.Accept()
		if err != nil {
			fmt.Println(err)
			continue
		}
		go handleConnection(conn)
	}
}

func handleConnection(conn net.Conn) {
	defer log.RecoverAndLog()

	bound := false
	var vm *virt.VM

	for {
		packet, err := ber.ReadPacket(conn)
		if err != nil {
			if err == io.EOF {
				break
			}
			panic(err)
		}

		messageID := packet.Children[0].Value.(uint64)
		request := packet.Children[1]

		switch request.Tag {
		case ApplicationBindRequest:
			name := request.Children[1].Value.(string)
			password := request.Children[2].Data.String()
			bound, vm = authenticate(name, password)

			var resultCode uint64 = LDAPResultInvalidCredentials
			if bound {
				resultCode = LDAPResultSuccess
			}
			conn.Write(createSimpleResponse(messageID, ApplicationBindResponse, resultCode).Bytes())

		case ApplicationUnbindRequest:
			bound = false

		case ApplicationSearchRequest:
			if !bound {
				conn.Write(createSimpleResponse(messageID, ApplicationSearchResultDone, LDAPResultInsufficientAccessRights).Bytes())
				continue
			}

			if !lookupUser(request.Children[6], messageID, vm, conn) {
				// ber.PrintPacket(packet) // for debugging
			}

			conn.Write(createSimpleResponse(messageID, ApplicationSearchResultDone, LDAPResultSuccess).Bytes())

		default:
			panic("Unsupported LDAP command")
		}
	}
}

func authenticate(name, password string) (bool, *virt.VM) {
	if name == "vmhost" && password == "abc" {
		return true, nil
	}

	if strings.HasPrefix(name, "vm-") && bson.IsObjectIdHex(name[3:]) {
		var vm virt.VM
		err := db.VMs.FindId(bson.ObjectIdHex(name[3:])).One(&vm)
		return err == nil && password == vm.LdapPassword, &vm
	}

	user, err := findUserByName(name)
	return err == nil && user.HasPassword(password), nil
}

func lookupUser(filter *ber.Packet, messageID uint64, vm *virt.VM, conn net.Conn) bool {
	switch findAttributeInFilter(filter, "objectClass") {
	case "posixGroup":
		if memberUid := findAttributeInFilter(filter, "memberUid"); memberUid != "" {
			if vm == nil {
				return true
			}

			user, err := findUserByName(memberUid)
			if err != nil {
				return true
			}
			userEntry := vm.GetUserEntry(user)

			if userEntry != nil && userEntry.Sudo {
				conn.Write(createSearchResultEntry(messageID, map[string]string{
					"objectClass": "posixGroup",
					"cn":          "sudo",
					"gidNumber":   "27",
				}).Bytes())
			}
			return true
		}

		if gidStr := findAttributeInFilter(filter, "gidNumber"); gidStr != "" {
			gid, _ := strconv.Atoi(gidStr)
			user, err := findUserByUid(gid)
			if err != nil || (vm != nil && vm.GetUserEntry(user) == nil) {
				return true
			}

			conn.Write(createSearchResultEntry(messageID, map[string]string{
				"objectClass": "posixGroup",
				"cn":          user.Name,
				"gidNumber":   gidStr,
			}).Bytes())
			return true
		}

	default: // including "posixAccount"
		user, err := findUserInFilter(filter)
		if err == nil && user == nil {
			return false
		}
		if err != nil || (vm != nil && vm.GetUserEntry(user) == nil) {
			return true
		}

		conn.Write(createSearchResultEntry(messageID, map[string]string{
			"objectClass":   "posixAccount",
			"cn":            user.Name,
			"uid":           user.Name,
			"userPassword":  "{SSHA}MmhmXrqch9NbAcRA6Z79OTSj6MqNXQxF",
			"uidNumber":     strconv.Itoa(user.Uid),
			"gidNumber":     strconv.Itoa(user.Uid),
			"homeDirectory": "/home/" + user.Name,
			"loginShell":    "/bin/bash",
		}).Bytes())
		return true

	}

	return false
}

func findUserInFilter(filter *ber.Packet) (*virt.User, error) {
	if name := findAttributeInFilter(filter, "uid"); name != "" {
		return findUserByName(name)
	}
	if uidStr := findAttributeInFilter(filter, "uidNumber"); uidStr != "" {
		uid, _ := strconv.Atoi(uidStr)
		return findUserByUid(uid)
	}
	return nil, nil
}

func findUser(query interface{}) (*virt.User, error) {
	var user virt.User
	err := db.Users.Find(query).One(&user)
	if err == nil && user.Uid < virt.UserIdOffset {
		panic("User with too low uid.")
	}
	return &user, err
}

func findUserByUid(id int) (*virt.User, error) {
	return findUser(bson.M{"uid": id})
}

func findUserByName(name string) (*virt.User, error) {
	return findUser(bson.M{"username": name})
}

func findAttributeInFilter(filter *ber.Packet, name string) string {
	switch filter.Tag {
	case 0, 1: // and, or
		for _, child := range filter.Children {
			if result := findAttributeInFilter(child, name); result != "" {
				return result
			}
		}
	case 3: // equalityMatch
		if filter.Children[0].Value == name {
			return filter.Children[1].Value.(string)
		}
	}
	return ""
}

func createLDAPMessage(messageID uint64) *ber.Packet {
	packet := ber.Encode(ber.ClassUniversal, ber.TypeConstructed, ber.TagSequence, nil, "LDAPMessage")
	packet.AppendChild(ber.NewInteger(ber.ClassUniversal, ber.TypePrimative, ber.TagInteger, messageID, "MessageID"))
	return packet
}

func createLDAPResponse(tag uint8, resultCode uint64) *ber.Packet {
	packet := ber.Encode(ber.ClassApplication, ber.TypeConstructed, tag, nil, "LDAPResponse")
	packet.AppendChild(ber.NewInteger(ber.ClassUniversal, ber.TypePrimative, ber.TagInteger, resultCode, "resultCode"))
	packet.AppendChild(ber.NewString(ber.ClassUniversal, ber.TypePrimative, ber.TagOctetString, "", "matchedDN"))
	packet.AppendChild(ber.NewString(ber.ClassUniversal, ber.TypePrimative, ber.TagOctetString, "", "errorMessage"))
	return packet
}

func createSimpleResponse(messageID uint64, tag uint8, resultCode uint64) *ber.Packet {
	response := createLDAPMessage(messageID)
	response.AppendChild(createLDAPResponse(tag, resultCode))
	return response
}

func createSearchResultEntry(messageID uint64, attributes map[string]string) *ber.Packet {
	response := createLDAPMessage(messageID)
	entry := ber.Encode(ber.ClassApplication, ber.TypeConstructed, ApplicationSearchResultEntry, nil, "SearchResultEntry")
	entry.AppendChild(ber.NewString(ber.ClassUniversal, ber.TypePrimative, ber.TagOctetString, attributes["cn"], "objectName"))
	attributeSequence := ber.Encode(ber.ClassUniversal, ber.TypeConstructed, ber.TagSequence, nil, "attributes")
	for key, value := range attributes {
		attributeSequence.AppendChild(createAttribute(key, value))
	}
	entry.AppendChild(attributeSequence)
	response.AppendChild(entry)
	return response
}

func createAttribute(name, value string) *ber.Packet {
	attribute := ber.Encode(ber.ClassUniversal, ber.TypeConstructed, ber.TagSequence, nil, "attribute")
	attribute.AppendChild(ber.NewString(ber.ClassUniversal, ber.TypePrimative, ber.TagOctetString, name, "type"))
	values := ber.Encode(ber.ClassUniversal, ber.TypeConstructed, ber.TagSet, nil, "values")
	values.AppendChild(ber.NewString(ber.ClassUniversal, ber.TypePrimative, ber.TagOctetString, value, "value"))
	attribute.AppendChild(values)
	return attribute
}

const (
	ApplicationBindRequest           = 0
	ApplicationBindResponse          = 1
	ApplicationUnbindRequest         = 2
	ApplicationSearchRequest         = 3
	ApplicationSearchResultEntry     = 4
	ApplicationSearchResultDone      = 5
	ApplicationModifyRequest         = 6
	ApplicationModifyResponse        = 7
	ApplicationAddRequest            = 8
	ApplicationAddResponse           = 9
	ApplicationDelRequest            = 10
	ApplicationDelResponse           = 11
	ApplicationModifyDNRequest       = 12
	ApplicationModifyDNResponse      = 13
	ApplicationCompareRequest        = 14
	ApplicationCompareResponse       = 15
	ApplicationAbandonRequest        = 16
	ApplicationSearchResultReference = 19
	ApplicationExtendedRequest       = 23
	ApplicationExtendedResponse      = 24
)

const (
	LDAPResultSuccess                      = 0
	LDAPResultOperationsError              = 1
	LDAPResultProtocolError                = 2
	LDAPResultTimeLimitExceeded            = 3
	LDAPResultSizeLimitExceeded            = 4
	LDAPResultCompareFalse                 = 5
	LDAPResultCompareTrue                  = 6
	LDAPResultAuthMethodNotSupported       = 7
	LDAPResultStrongAuthRequired           = 8
	LDAPResultReferral                     = 10
	LDAPResultAdminLimitExceeded           = 11
	LDAPResultUnavailableCriticalExtension = 12
	LDAPResultConfidentialityRequired      = 13
	LDAPResultSaslBindInProgress           = 14
	LDAPResultNoSuchAttribute              = 16
	LDAPResultUndefinedAttributeType       = 17
	LDAPResultInappropriateMatching        = 18
	LDAPResultConstraintViolation          = 19
	LDAPResultAttributeOrValueExists       = 20
	LDAPResultInvalidAttributeSyntax       = 21
	LDAPResultNoSuchObject                 = 32
	LDAPResultAliasProblem                 = 33
	LDAPResultInvalidDNSyntax              = 34
	LDAPResultAliasDereferencingProblem    = 36
	LDAPResultInappropriateAuthentication  = 48
	LDAPResultInvalidCredentials           = 49
	LDAPResultInsufficientAccessRights     = 50
	LDAPResultBusy                         = 51
	LDAPResultUnavailable                  = 52
	LDAPResultUnwillingToPerform           = 53
	LDAPResultLoopDetect                   = 54
	LDAPResultNamingViolation              = 64
	LDAPResultObjectClassViolation         = 65
	LDAPResultNotAllowedOnNonLeaf          = 66
	LDAPResultNotAllowedOnRDN              = 67
	LDAPResultEntryAlreadyExists           = 68
	LDAPResultObjectClassModsProhibited    = 69
	LDAPResultAffectsMultipleDSAs          = 71
	LDAPResultOther                        = 80

	ErrorNetwork         = 200
	ErrorFilterCompile   = 201
	ErrorFilterDecompile = 202
	ErrorDebugging       = 203
)
<|endoftext|>


<commit_before>package misspell

import (
	"testing"
)

func TestFalsePositives(t *testing.T) {
	cases := []string{
		"replicaset",
		"parallelise",
		"parallelize",
		"perceptron",  // http://foldoc.org/perceptron
		"perceptrons", // ^^
		"convertors",  // alt spelling
		"adventurers",
		" s.svc.GetObject ",
		"infinitie.net",
		"foo summaries\n",
		"thru",
		"publically",
		"6YUO5",   // base64
		"cleaner", // triggered by "cleane->cleanser" and partial word FP
		"  http.Redirect(w, req, req.URL.Path, http.StatusFound) ",
		"url is http://zeebra.com ",
		"path is /zeebra?zeebra=zeebra ",
		"Malcom_McLean",
		"implementor", // alt spelling, see https://github.com/client9/misspell/issues/46
		"searchtypes",
		" witness",
		"returndata",
		"UNDERSTOOD",
		"textinterface",
		" committed ",
		"committed",
		"Bengali",
		"Portuguese",
		"scientists",
		"causally",
		"embarrassing",
		"setuptools", // python package
		"committing",
		"guises",
		"disguise",
		"begging",
		"cmo",
		"cmos",
		"borked",
		"hadn't",
		"Iceweasel",
		"summarised",
		"autorenew",
		"travelling",
		"republished",
		"fallthru",
		"pruning",
		"deb.VersionDontCare",
		"authtag",
		"intrepid",
		"usefully",
		"there",
		"definite",
		"earliest",
		"Japanese",
		"international",
		"excellent",
		"gracefully",
		"carefully",
		"class",
		"include",
		"process",
		"address",
		"attempt",
		"large",
		"although",
		"specific",
		"taste",
		"against",
		"successfully",
		"unsuccessfully",
		"occurred",
		"agree",
		"controlled",
		"publisher",
		"strategy",
		"geoposition",
		"paginated",
		"happened",
		"relative",
		"computing",
		"language",
		"manual",
		"token",
		"into",
		"nothing",
		"datatool",
		"propose",
		"learnt",
		"tolerant",
		"whitehat",
		"monotonic",
		"comprised",
		"indemnity",
		"flattened",
		"interrupted",
		"inotify",
		"occasional",
		"forging",
		"ampersand",
		"decomposition",
		"commit",
		"programmer", // "grammer"
		//		"requestsinserted",
		"seeked",      // technical word
		"bodyreader",  // variable name
		"cantPrepare", // variable name
		"dontPrepare", // variable name
	}
	r := New()
	r.Debug = true
	for casenum, tt := range cases {
		got, _ := r.Replace(tt)
		if got != tt {
			t.Errorf("%d: %q got converted to %q", casenum, tt, got)
		}
	}
}
<commit_msg>falsepositives_test.go<commit_after>package misspell

import (
	"testing"
)

func TestFalsePositives(t *testing.T) {
	cases := []string{
		"requesters",
		"requestors",
		"replicaset",
		"parallelise",
		"parallelize",
		"perceptron",  // http://foldoc.org/perceptron
		"perceptrons", // ^^
		"convertors",  // alt spelling
		"adventurers",
		" s.svc.GetObject ",
		"infinitie.net",
		"foo summaries\n",
		"thru",
		"publically",
		"6YUO5",   // base64
		"cleaner", // triggered by "cleane->cleanser" and partial word FP
		"  http.Redirect(w, req, req.URL.Path, http.StatusFound) ",
		"url is http://zeebra.com ",
		"path is /zeebra?zeebra=zeebra ",
		"Malcom_McLean",
		"implementor", // alt spelling, see https://github.com/client9/misspell/issues/46
		"searchtypes",
		" witness",
		"returndata",
		"UNDERSTOOD",
		"textinterface",
		" committed ",
		"committed",
		"Bengali",
		"Portuguese",
		"scientists",
		"causally",
		"embarrassing",
		"setuptools", // python package
		"committing",
		"guises",
		"disguise",
		"begging",
		"cmo",
		"cmos",
		"borked",
		"hadn't",
		"Iceweasel",
		"summarised",
		"autorenew",
		"travelling",
		"republished",
		"fallthru",
		"pruning",
		"deb.VersionDontCare",
		"authtag",
		"intrepid",
		"usefully",
		"there",
		"definite",
		"earliest",
		"Japanese",
		"international",
		"excellent",
		"gracefully",
		"carefully",
		"class",
		"include",
		"process",
		"address",
		"attempt",
		"large",
		"although",
		"specific",
		"taste",
		"against",
		"successfully",
		"unsuccessfully",
		"occurred",
		"agree",
		"controlled",
		"publisher",
		"strategy",
		"geoposition",
		"paginated",
		"happened",
		"relative",
		"computing",
		"language",
		"manual",
		"token",
		"into",
		"nothing",
		"datatool",
		"propose",
		"learnt",
		"tolerant",
		"whitehat",
		"monotonic",
		"comprised",
		"indemnity",
		"flattened",
		"interrupted",
		"inotify",
		"occasional",
		"forging",
		"ampersand",
		"decomposition",
		"commit",
		"programmer", // "grammer"
		//		"requestsinserted",
		"seeked",      // technical word
		"bodyreader",  // variable name
		"cantPrepare", // variable name
		"dontPrepare", // variable name
	}
	r := New()
	r.Debug = true
	for casenum, tt := range cases {
		got, _ := r.Replace(tt)
		if got != tt {
			t.Errorf("%d: %q got converted to %q", casenum, tt, got)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 Herbert G. Fischer. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package imagick

/*
#include <wand/MagickWand.h>
*/
import "C"
import "unsafe"

// Returns the ImageMagick API copyright as a string constant.
func GetCopyright() string {
	cstr := C.MagickGetCopyright()
	defer C.free(unsafe.Pointer(cstr))
	return C.GoString(cstr)
}

// Returns the ImageMagick home URL.
func GetHomeURL() string {
	cstr := C.MagickGetHomeURL()
	defer C.free(unsafe.Pointer(cstr))
	return C.GoString(cstr)
}

// Returns the ImageMagick package name as a string constant.
func GetPackageName() string {
	cstr := C.MagickGetPackageName()
	defer C.free(unsafe.Pointer(cstr))
	return C.GoString(cstr)
}

// Returns the ImageMagick release date as a string constant.
func GetReleaseDate() string {
	cstr := C.MagickGetReleaseDate()
	defer C.free(unsafe.Pointer(cstr))
	return C.GoString(cstr)
}

// Returns the ImageMagick quantum depth as a string constant.
func GetQuantumDepth() (string, uint) {
	cst := C.size_t(0)
	csq := C.MagickGetQuantumDepth(&cst)
	return C.GoString(csq), uint(cst)
}

// Returns the ImageMagick quantum range as a string constant.
func GetQuantumRange() (string, uint) {
	cst := C.size_t(0)
	csq := C.MagickGetQuantumRange(&cst)
	return C.GoString(csq), uint(cst)
}

// Returns the specified resource in megabytes.
func GetResource(rtype ResourceType) int64 {
	return int64(C.MagickGetResource(C.ResourceType(rtype)))

}

// Returns the specified resource limit in megabytes.
func GetResourceLimit(rtype ResourceType) int64 {
	return int64(C.MagickGetResourceLimit(C.ResourceType(rtype)))
}

// Returns the ImageMagick API version as a string constant and as a number.
func GetVersion() (version string, nversion uint) {
	cnver := C.size_t(0)
	csver := C.MagickGetVersion(&cnver)
	defer C.free(unsafe.Pointer(csver))
	version = C.GoString(csver)
	nversion = uint(cnver)
	return
}
<commit_msg>Do not free C constants<commit_after>// Copyright 2013 Herbert G. Fischer. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package imagick

/*
#include <wand/MagickWand.h>
*/
import "C"
import "unsafe"

// Returns the ImageMagick API copyright as a string constant.
func GetCopyright() string {
	cstr := C.MagickGetCopyright()
	return C.GoString(cstr)
}

// Returns the ImageMagick home URL.
func GetHomeURL() string {
	cstr := C.MagickGetHomeURL()
	return C.GoString(cstr)
}

// Returns the ImageMagick package name as a string constant.
func GetPackageName() string {
	cstr := C.MagickGetPackageName()
	return C.GoString(cstr)
}

// Returns the ImageMagick release date as a string constant.
func GetReleaseDate() string {
	cstr := C.MagickGetReleaseDate()
	return C.GoString(cstr)
}

// Returns the ImageMagick quantum depth as a string constant.
func GetQuantumDepth() (string, uint) {
	cst := C.size_t(0)
	csq := C.MagickGetQuantumDepth(&cst)
	return C.GoString(csq), uint(cst)
}

// Returns the ImageMagick quantum range as a string constant.
func GetQuantumRange() (string, uint) {
	cst := C.size_t(0)
	csq := C.MagickGetQuantumRange(&cst)
	return C.GoString(csq), uint(cst)
}

// Returns the specified resource in megabytes.
func GetResource(rtype ResourceType) int64 {
	return int64(C.MagickGetResource(C.ResourceType(rtype)))

}

// Returns the specified resource limit in megabytes.
func GetResourceLimit(rtype ResourceType) int64 {
	return int64(C.MagickGetResourceLimit(C.ResourceType(rtype)))
}

// Returns the ImageMagick API version as a string constant and as a number.
func GetVersion() (version string, nversion uint) {
	cnver := C.size_t(0)
	csver := C.MagickGetVersion(&cnver)
	version = C.GoString(csver)
	nversion = uint(cnver)
	return
}
<|endoftext|>


<commit_before>package pointer

import (
	"bufio"
	"errors"
	"fmt"
	"github.com/github/git-media/git"
	"github.com/github/git-media/gitmedia"
	"io"
	"os"
	"path/filepath"
	"strings"
)

var (
	linkTemplate = `oid %s
name %s
`
)

// Link provides a link between a git sha1 and a git media oid. Link files
// are stored under .git/media/objects in a structure similar to that of
// .git/objects. The link file contains the git media oid and the name of
// file if it was available when created.
type Link struct {
	Oid  string
	Name string
}

// CreateLink will create a link file from a Pointer. The filename passed
// in is what will be stored in the link contents.
func (p *Pointer) CreateLink(filename string) error {
	hash, err := git.HashObject([]byte(p.Encoded()))
	if err != nil {
		return err
	}

	linkFile, err := gitmedia.LocalLinkPath(hash)
	if err != nil {
		return err
	}
	linkLock := linkFile + ".lock"

	file, err := os.OpenFile(linkLock, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0644)
	if err != nil {
		return err
	}

	_, err = file.WriteString(fmt.Sprintf(linkTemplate, p.Oid, filename))
	if err != nil {
		os.Remove(linkLock)
		return err
	}

	file.Close()

	return os.Rename(linkLock, linkFile)
}

// FindLink takes a git sha1 and attempts to find the link file associated
// with it. While this can be used to determine whether the git object is
// a git media object, it should not be considered authoritative. While a link
// file's presence is an indicator that it is a git media file, its absence does
// not mean that it is not a git media file.
func FindLink(sha1 string) (*Link, error) {
	linkPath := filepath.Join(gitmedia.LocalLinkDir, sha1[0:2], sha1[2:len(sha1)])

	linkFile, err := os.Open(linkPath)
	if err != nil {
		return nil, err
	}

	link, err := DecodeLink(linkFile)
	if err != nil {
		return nil, err
	}

	return link, nil
}

func DecodeLink(reader io.Reader) (*Link, error) {
	link := &Link{}

	m := make(map[string]string)
	scanner := bufio.NewScanner(reader)
	for scanner.Scan() {
		text := scanner.Text()
		if len(text) == 0 {
			continue
		}

		parts := strings.SplitN(text, " ", 2)
		key := parts[0]
		m[key] = parts[1]
	}

	oid, ok := m["oid"]
	if !ok {
		return nil, errors.New("No Oid in link file")
	}

	link.Oid = oid
	link.Name = m["name"]
	return link, nil
}
<commit_msg>Remove the original file before mv'ing to it, Windows errors if the file exists<commit_after>package pointer

import (
	"bufio"
	"errors"
	"fmt"
	"github.com/github/git-media/git"
	"github.com/github/git-media/gitmedia"
	"io"
	"os"
	"path/filepath"
	"strings"
)

var (
	linkTemplate = `oid %s
name %s
`
)

// Link provides a link between a git sha1 and a git media oid. Link files
// are stored under .git/media/objects in a structure similar to that of
// .git/objects. The link file contains the git media oid and the name of
// file if it was available when created.
type Link struct {
	Oid  string
	Name string
}

// CreateLink will create a link file from a Pointer. The filename passed
// in is what will be stored in the link contents.
func (p *Pointer) CreateLink(filename string) error {
	hash, err := git.HashObject([]byte(p.Encoded()))
	if err != nil {
		return err
	}

	linkFile, err := gitmedia.LocalLinkPath(hash)
	if err != nil {
		return err
	}
	linkLock := linkFile + ".lock"

	file, err := os.OpenFile(linkLock, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0644)
	if err != nil {
		return err
	}

	_, err = file.WriteString(fmt.Sprintf(linkTemplate, p.Oid, filename))
	if err != nil {
		os.Remove(linkLock)
		return err
	}

	file.Close()

	os.Remove(linkFile) // Remove the link file if it already existed

	return os.Rename(linkLock, linkFile)
}

// FindLink takes a git sha1 and attempts to find the link file associated
// with it. While this can be used to determine whether the git object is
// a git media object, it should not be considered authoritative. While a link
// file's presence is an indicator that it is a git media file, its absence does
// not mean that it is not a git media file.
func FindLink(sha1 string) (*Link, error) {
	linkPath := filepath.Join(gitmedia.LocalLinkDir, sha1[0:2], sha1[2:len(sha1)])

	linkFile, err := os.Open(linkPath)
	if err != nil {
		return nil, err
	}

	link, err := DecodeLink(linkFile)
	if err != nil {
		return nil, err
	}

	return link, nil
}

func DecodeLink(reader io.Reader) (*Link, error) {
	link := &Link{}

	m := make(map[string]string)
	scanner := bufio.NewScanner(reader)
	for scanner.Scan() {
		text := scanner.Text()
		if len(text) == 0 {
			continue
		}

		parts := strings.SplitN(text, " ", 2)
		key := parts[0]
		m[key] = parts[1]
	}

	oid, ok := m["oid"]
	if !ok {
		return nil, errors.New("No Oid in link file")
	}

	link.Oid = oid
	link.Name = m["name"]
	return link, nil
}
<|endoftext|>


<commit_before>package tsi1

import (
	"context"
	"fmt"
	"io"
	"math"
	"os"
	"path/filepath"
	"sort"
	"text/tabwriter"

	"github.com/influxdata/influxdb"
	"github.com/influxdata/influxdb/tsdb"
	"go.uber.org/zap"
)

const (
	// Number of series IDs to stored in slice before we convert to a roaring
	// bitmap. Roaring bitmaps have a non-trivial initial cost to construct.
	useBitmapN = 25
)

// ReportCommand represents the program execution for "influxd inspect report-tsi".
type ReportCommand struct {
	// Standard input/output, overridden for testing.
	Stderr io.Writer
	Stdout io.Writer
	Logger *zap.Logger

	// Filters
	DataPath        string
	OrgID, BucketID *influxdb.ID

	byOrg               map[influxdb.ID]*cardinality
	byBucket            map[influxdb.ID]*cardinality
	byBucketMeasurement map[influxdb.ID]map[string]*cardinality
	orgToBucket         map[influxdb.ID][]influxdb.ID

	SeriesDirPath string // optional. Defaults to dbPath/_series
	sfile         *tsdb.SeriesFile
	indexFile     *Index

	TopN          int
	ByMeasurement bool
	byTagKey      bool
}

// NewReportCommand returns a new instance of ReportCommand with default setting applied.
func NewReportCommand() *ReportCommand {
	return &ReportCommand{
		Logger:              zap.NewNop(),
		byOrg:               make(map[influxdb.ID]*cardinality),
		byBucket:            make(map[influxdb.ID]*cardinality),
		byBucketMeasurement: make(map[influxdb.ID]map[string]*cardinality),
		orgToBucket:         make(map[influxdb.ID][]influxdb.ID),
		TopN:                0,
		byTagKey:            false,
	}
}

// ReportTSISummary is returned by a report-tsi Run() command and is used to access cardinality information
type Summary struct {
	OrgCardinality               map[influxdb.ID]int64
	BucketCardinality            map[influxdb.ID]int64
	BucketMeasurementCardinality map[influxdb.ID]map[string]int64
}

func newSummary() *Summary {
	return &Summary{
		OrgCardinality:               make(map[influxdb.ID]int64),
		BucketCardinality:            make(map[influxdb.ID]int64),
		BucketMeasurementCardinality: make(map[influxdb.ID]map[string]int64),
	}
}

// Run runs the report-tsi tool which can be used to find the cardinality
// any org or bucket. Run returns a *ReportTSISummary, which contains maps for finding
// the cardinality of a bucket or org based on its influxdb.ID
func (report *ReportCommand) Run(print bool) (*Summary, error) {
	report.Stdout = os.Stdout

	if report.SeriesDirPath == "" {
		report.SeriesDirPath = filepath.Join(report.DataPath, "_series")
	}

	sFile := tsdb.NewSeriesFile(report.SeriesDirPath)

	// TODO: do we actually want the seriesfile logging?
	// sFile.WithLogger(report.Logger)

	if err := sFile.Open(context.Background()); err != nil {
		report.Logger.Error("failed to open series")
		return nil, err
	}
	defer sFile.Close()
	report.sfile = sFile

	indexPath := filepath.Join(report.DataPath, "index")
	report.indexFile = NewIndex(sFile, NewConfig(), WithPath(indexPath))
	if err := report.indexFile.Open(context.Background()); err != nil {
		return nil, err
	}
	defer report.indexFile.Close()

	summary, err := report.calculateOrgBucketCardinality()
	if err != nil {
		return nil, err
	}

	if print {
		report.printCardinalitySummary(summary)
	}

	return summary, nil
}

type cardinality struct {
	name  []byte
	short []uint32
	set   *tsdb.SeriesIDSet
}

func (c *cardinality) add(x uint64) {
	if c.set != nil {
		c.set.AddNoLock(tsdb.NewSeriesID(x))
		return
	}

	c.short = append(c.short, uint32(x)) // Series IDs never get beyond 2^32

	// Cheaper to store in bitmap.
	if len(c.short) > useBitmapN {
		c.set = tsdb.NewSeriesIDSet()
		for i := 0; i < len(c.short); i++ {
			c.set.AddNoLock(tsdb.NewSeriesID(uint64(c.short[i])))
		}
		c.short = nil
		return
	}
}

func (c *cardinality) cardinality() int64 {
	if c == nil || (c.short == nil && c.set == nil) {
		return 0
	}

	if c.short != nil {
		return int64(len(c.short))
	}
	return int64(c.set.Cardinality())
}

func (report *ReportCommand) calculateCardinalities() error {
	idx := report.indexFile
	itr, err := idx.MeasurementIterator()
	if err != nil {
		return err
	} else if itr == nil {
		return nil
	}
	defer itr.Close()

	var totalCard = &cardinality{name: []byte("total")}
OUTER:
	for {
		name, err := itr.Next()
		if err != nil {
			return err
		} else if name == nil {
			break OUTER
		}

		// decode org and bucket from measurement name
		var a [16]byte
		copy(a[:], name[:16])
		org, bucket := tsdb.DecodeName(a)

		if report.OrgID != nil && *report.OrgID != org {
			continue
		} else if report.BucketID != nil && *report.BucketID != bucket {
			continue
		}

		sitr, err := idx.MeasurementSeriesIDIterator(name)
		if err != nil {
			return err
		} else if sitr == nil {
			continue
		}

		var orgCard, bucketCard *cardinality

		// initialize map of bucket to measurements
		if _, ok := report.byBucketMeasurement[bucket]; !ok {
			report.byBucketMeasurement[bucket] = make(map[string]*cardinality)
		}

		// initialize total cardinality tracking struct for this bucket
		if c, ok := report.byBucket[bucket]; !ok {
			bucketCard = &cardinality{name: []byte(bucket.String())}
			report.byBucket[bucket] = bucketCard
		} else {
			bucketCard = c
		}

		// initialize total cardinality tracking struct for this org
		if c, ok := report.byOrg[org]; !ok {
			orgCard = &cardinality{name: []byte(bucket.String())}
			report.byOrg[org] = orgCard
		} else {
			orgCard = c
		}

		// track org to bucket mappings so we can associate org an bucket in the printed summary
		if _, ok := report.orgToBucket[org]; !ok {
			report.orgToBucket[org] = []influxdb.ID{}
		}

		report.orgToBucket[org] = append(report.orgToBucket[org], bucket)

		var e tsdb.SeriesIDElem
		for e, err = sitr.Next(); err == nil && e.SeriesID.ID != 0; e, err = sitr.Next() {
			id := e.SeriesID.ID

			if id > math.MaxUint32 {
				panic(fmt.Sprintf("series ID is too large: %d (max %d). Corrupted series file?", e.SeriesID, uint32(math.MaxUint32)))
			}

			totalCard.add(id)

			// add cardinalities to org and bucket maps
			orgCard.add(id)
			bucketCard.add(id)

			// retrieve tags associated with series id so we can get
			// associated measurement
			_, tags := report.sfile.Series(e.SeriesID)
			if len(tags) == 0 {
				panic(fmt.Sprintf("empty series key"))
			}

			// measurement name should be first tag.
			mName := string(tags[0].Value)

			// update measurement-level cardinality if tracking by measurement
			if report.ByMeasurement {
				var mCard *cardinality
				if cardForM, ok := report.byBucketMeasurement[bucket][mName]; !ok {
					mCard = &cardinality{name: []byte(mName)}
					report.byBucketMeasurement[bucket][mName] = mCard
				} else {
					mCard = cardForM
				}

				mCard.add(id)
			}
		}

		sitr.Close()

		if err != nil {
			return err
		}

	}
	return nil
}

func (report *ReportCommand) calculateOrgBucketCardinality() (*Summary, error) {
	if err := report.calculateCardinalities(); err != nil {
		return nil, err
	}
	// Generate a new summary
	summary := newSummary()
	for bucketID, bucketCard := range report.byBucket {
		summary.BucketCardinality[bucketID] = bucketCard.cardinality()
		summary.BucketMeasurementCardinality[bucketID] = make(map[string]int64)
	}

	for orgID, orgCard := range report.byOrg {
		summary.OrgCardinality[orgID] = orgCard.cardinality()
	}

	for bucketID, bucketMeasurement := range report.byBucketMeasurement {
		for mName, mCard := range bucketMeasurement {
			summary.BucketMeasurementCardinality[bucketID][mName] = mCard.cardinality()
		}
	}

	return summary, nil
}

func (report *ReportCommand) printCardinalitySummary(summary *Summary) {
	tw := tabwriter.NewWriter(report.Stdout, 4, 4, 1, '\t', 0)
	fmt.Fprint(tw, "\n")

	// sort total org and bucket cardinalities and limit to top n values
	sortedOrgs := sortKeys(summary.OrgCardinality, report.TopN)
	sortedBuckets := sortKeys(summary.BucketCardinality, report.TopN)

	for orgIndex := range sortedOrgs {
		// if we specify a bucket, we do not print the org cardinality
		if report.BucketID == nil {
			fmt.Fprintf(tw, "Org %s total: %d\n", sortedOrgs[orgIndex].id, sortedOrgs[orgIndex].card)
		}

		for bucketIndex := range sortedBuckets {
			fmt.Fprintf(tw, "\tBucket %s total: %d\n", sortedBuckets[bucketIndex].id, sortedBuckets[bucketIndex].card)

			if report.ByMeasurement {
				bucketID, _ := influxdb.IDFromString(sortedBuckets[bucketIndex].id)
				sortedMeasurements := sortMeasurements(summary.BucketMeasurementCardinality[*bucketID], report.TopN)

				for measIndex := range sortedMeasurements {
					fmt.Fprintf(tw, "\t\t_m=%s\t%d\n", sortedMeasurements[measIndex].id, sortedMeasurements[measIndex].card)
				}
			}
		}
	}
	fmt.Fprint(tw, "\n")

	tw.Flush()
}

// sortKeys is a quick helper to return the sorted set of a map's keys
// sortKeys will only return report.topN keys if the flag is set
type result struct {
	id   string
	card int64
}

type resultList []result

func (a resultList) Len() int           { return len(a) }
func (a resultList) Less(i, j int) bool { return a[i].card < a[j].card }
func (a resultList) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }

func sortKeys(vals map[influxdb.ID]int64, topN int) resultList {
	sorted := make(resultList, 0)
	for k, v := range vals {
		sorted = append(sorted, result{k.String(), v})
	}
	sort.Sort(sort.Reverse(sorted))

	if topN == 0 {
		return sorted
	}
	if topN > len(sorted) {
		topN = len(sorted)
	}
	return sorted[:topN]
}

func sortMeasurements(vals map[string]int64, topN int) resultList {
	sorted := make(resultList, 0)
	for k, v := range vals {
		sorted = append(sorted, result{k, v})
	}
	sort.Sort(sort.Reverse(sorted))

	if topN == 0 {
		return sorted
	}
	if topN > len(sorted) {
		topN = len(sorted)
	}
	return sorted[:topN]
}
<commit_msg>fix(tsi1): map org to bucket in report-tsi tool so output is more useful<commit_after>package tsi1

import (
	"context"
	"fmt"
	"io"
	"math"
	"os"
	"path/filepath"
	"sort"
	"text/tabwriter"
	"time"

	"github.com/influxdata/influxdb"
	"github.com/influxdata/influxdb/tsdb"
	"go.uber.org/zap"
)

const (
	// Number of series IDs to stored in slice before we convert to a roaring
	// bitmap. Roaring bitmaps have a non-trivial initial cost to construct.
	useBitmapN = 25
)

// ReportCommand represents the program execution for "influxd inspect report-tsi".
type ReportCommand struct {
	// Standard input/output, overridden for testing.
	Stderr io.Writer
	Stdout io.Writer
	Logger *zap.Logger

	// Filters
	DataPath        string
	OrgID, BucketID *influxdb.ID

	byOrgBucket         map[influxdb.ID]map[influxdb.ID]*cardinality
	byBucketMeasurement map[influxdb.ID]map[string]*cardinality
	orgToBucket         map[influxdb.ID][]influxdb.ID

	SeriesDirPath string // optional. Defaults to dbPath/_series
	sfile         *tsdb.SeriesFile
	indexFile     *Index

	TopN          int
	ByMeasurement bool
	byTagKey      bool

	start time.Time
}

// NewReportCommand returns a new instance of ReportCommand with default setting applied.
func NewReportCommand() *ReportCommand {
	return &ReportCommand{
		Logger:              zap.NewNop(),
		byOrgBucket:         make(map[influxdb.ID]map[influxdb.ID]*cardinality),
		byBucketMeasurement: make(map[influxdb.ID]map[string]*cardinality),
		orgToBucket:         make(map[influxdb.ID][]influxdb.ID),
		TopN:                0,
		byTagKey:            false,
	}
}

// ReportTSISummary is returned by a report-tsi Run() command and is used to access cardinality information
type Summary struct {
	TotalCardinality             int64
	OrgCardinality               map[influxdb.ID]int64
	BucketByOrgCardinality       map[influxdb.ID]map[influxdb.ID]int64
	BucketMeasurementCardinality map[influxdb.ID]map[string]int64
}

func newSummary() *Summary {
	return &Summary{
		OrgCardinality:               make(map[influxdb.ID]int64),
		BucketByOrgCardinality:       make(map[influxdb.ID]map[influxdb.ID]int64),
		BucketMeasurementCardinality: make(map[influxdb.ID]map[string]int64),
	}
}

// Run runs the report-tsi tool which can be used to find the cardinality
// any org or bucket. Run returns a *ReportTSISummary, which contains maps for finding
// the cardinality of a bucket or org based on its influxdb.ID
func (report *ReportCommand) Run(print bool) (*Summary, error) {
	report.start = time.Now()
	report.Stdout = os.Stdout

	if report.SeriesDirPath == "" {
		report.SeriesDirPath = filepath.Join(report.DataPath, "_series")
	}

	sFile := tsdb.NewSeriesFile(report.SeriesDirPath)

	// TODO: do we actually want the seriesfile logging?
	// sFile.WithLogger(report.Logger)

	if err := sFile.Open(context.Background()); err != nil {
		report.Logger.Error("failed to open series")
		return nil, err
	}
	defer sFile.Close()
	report.sfile = sFile

	indexPath := filepath.Join(report.DataPath, "index")
	report.indexFile = NewIndex(sFile, NewConfig(), WithPath(indexPath))
	if err := report.indexFile.Open(context.Background()); err != nil {
		return nil, err
	}
	defer report.indexFile.Close()

	summary, err := report.calculateOrgBucketCardinality()
	if err != nil {
		return nil, err
	}

	if print {
		report.printCardinalitySummary(summary)
	}

	return summary, nil
}

type cardinality struct {
	name  []byte
	short []uint32
	set   *tsdb.SeriesIDSet
}

func (c *cardinality) add(x uint64) {
	if c.set != nil {
		c.set.AddNoLock(tsdb.NewSeriesID(x))
		return
	}

	c.short = append(c.short, uint32(x)) // Series IDs never get beyond 2^32

	// Cheaper to store in bitmap.
	if len(c.short) > useBitmapN {
		c.set = tsdb.NewSeriesIDSet()
		for i := 0; i < len(c.short); i++ {
			c.set.AddNoLock(tsdb.NewSeriesID(uint64(c.short[i])))
		}
		c.short = nil
		return
	}
}

func (c *cardinality) cardinality() int64 {
	if c == nil || (c.short == nil && c.set == nil) {
		return 0
	}

	if c.short != nil {
		return int64(len(c.short))
	}
	return int64(c.set.Cardinality())
}

func (report *ReportCommand) calculateCardinalities() error {
	idx := report.indexFile
	itr, err := idx.MeasurementIterator()
	if err != nil {
		return err
	} else if itr == nil {
		return nil
	}
	defer itr.Close()

OUTER:
	for {
		name, err := itr.Next()
		if err != nil {
			return err
		} else if name == nil {
			break OUTER
		}

		// decode org and bucket from measurement name
		var a [16]byte
		copy(a[:], name[:16])
		org, bucket := tsdb.DecodeName(a)

		if report.OrgID != nil && *report.OrgID != org {
			continue
		} else if report.BucketID != nil && *report.BucketID != bucket {
			continue
		}

		sitr, err := idx.MeasurementSeriesIDIterator(name)
		if err != nil {
			return err
		} else if sitr == nil {
			continue
		}

		var bucketCard *cardinality

		// initialize map of bucket to measurements
		if _, ok := report.byBucketMeasurement[bucket]; !ok {
			report.byBucketMeasurement[bucket] = make(map[string]*cardinality)
		}

		if _, ok := report.byOrgBucket[org]; !ok {
			report.byOrgBucket[org] = make(map[influxdb.ID]*cardinality)
		}

		// initialize total cardinality tracking struct for this bucket
		if c, ok := report.byOrgBucket[org][bucket]; !ok {
			bucketCard = &cardinality{name: []byte(bucket.String())}
			report.byOrgBucket[org][bucket] = bucketCard
		} else {
			bucketCard = c
		}

		var e tsdb.SeriesIDElem
		for e, err = sitr.Next(); err == nil && e.SeriesID.ID != 0; e, err = sitr.Next() {
			id := e.SeriesID.ID

			if id > math.MaxUint32 {
				panic(fmt.Sprintf("series ID is too large: %d (max %d). Corrupted series file?", e.SeriesID, uint32(math.MaxUint32)))
			}

			// add cardinality to bucket
			bucketCard.add(id)

			// retrieve tags associated with series id so we can get
			// associated measurement
			_, tags := report.sfile.Series(e.SeriesID)
			if len(tags) == 0 {
				panic(fmt.Sprintf("empty series key"))
			}

			// measurement name should be first tag
			mName := string(tags[0].Value)

			// update measurement-level cardinality if tracking by measurement
			if report.ByMeasurement {
				var mCard *cardinality
				if cardForM, ok := report.byBucketMeasurement[bucket][mName]; !ok {
					mCard = &cardinality{name: []byte(mName)}
					report.byBucketMeasurement[bucket][mName] = mCard
				} else {
					mCard = cardForM
				}

				mCard.add(id)
			}
		}

		sitr.Close()

		if err != nil {
			return err
		}

	}
	return nil
}

func (report *ReportCommand) calculateOrgBucketCardinality() (*Summary, error) {
	if err := report.calculateCardinalities(); err != nil {
		return nil, err
	}

	var totalCard int64
	// Generate a new summary
	summary := newSummary()
	for orgID, bucketMap := range report.byOrgBucket {
		summary.BucketByOrgCardinality[orgID] = make(map[influxdb.ID]int64)
		orgTotal := int64(0)
		for bucketID, bucketCard := range bucketMap {
			count := bucketCard.cardinality()
			summary.BucketByOrgCardinality[orgID][bucketID] = count
			summary.BucketMeasurementCardinality[bucketID] = make(map[string]int64)
			orgTotal += count
			totalCard += count
		}
		summary.OrgCardinality[orgID] = orgTotal
	}

	summary.TotalCardinality = totalCard

	for bucketID, bucketMeasurement := range report.byBucketMeasurement {
		for mName, mCard := range bucketMeasurement {
			summary.BucketMeasurementCardinality[bucketID][mName] = mCard.cardinality()
		}
	}

	return summary, nil
}

func (report *ReportCommand) printCardinalitySummary(summary *Summary) {
	tw := tabwriter.NewWriter(report.Stdout, 4, 4, 1, '\t', 0)
	fmt.Fprint(tw, "\n")

	fmt.Fprintf(tw, "Total: %d\n", summary.TotalCardinality)
	// sort total org and bucket  and limit to top n values
	sortedOrgs := sortKeys(summary.OrgCardinality, report.TopN)

	for i, orgResult := range sortedOrgs {
		orgID, _ := influxdb.IDFromString(orgResult.id)
		sortedBuckets := sortKeys(summary.BucketByOrgCardinality[*orgID], report.TopN)
		// if we specify a bucket, we do not print the org cardinality
		fmt.Fprintln(tw, "===============")
		if report.BucketID == nil {
			fmt.Fprintf(tw, "Org %s total: %d\n", orgResult.id, orgResult.card)
		}

		for _, bucketResult := range sortedBuckets {
			fmt.Fprintf(tw, "\tBucket %s total: %d\n", bucketResult.id, bucketResult.card)

			if report.ByMeasurement {
				bucketID, _ := influxdb.IDFromString(bucketResult.id)
				sortedMeasurements := sortMeasurements(summary.BucketMeasurementCardinality[*bucketID], report.TopN)

				for _, measResult := range sortedMeasurements {
					fmt.Fprintf(tw, "\t\t_m=%s\t%d\n", measResult.id, measResult.card)
				}
			}
		}
		if i == len(sortedOrgs)-1 {
			fmt.Fprintln(tw, "===============")
		}
	}
	fmt.Fprint(tw, "\n\n")

	elapsed := time.Since(report.start)
	fmt.Fprintf(tw, "Finished in %v\n", elapsed)

	tw.Flush()
}

// sortKeys is a quick helper to return the sorted set of a map's keys
// sortKeys will only return report.topN keys if the flag is set
type result struct {
	id   string
	card int64
}

type resultList []result

func (a resultList) Len() int           { return len(a) }
func (a resultList) Less(i, j int) bool { return a[i].card < a[j].card }
func (a resultList) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }

func sortKeys(vals map[influxdb.ID]int64, topN int) resultList {
	sorted := make(resultList, 0)
	for k, v := range vals {
		sorted = append(sorted, result{k.String(), v})
	}
	sort.Sort(sort.Reverse(sorted))

	if topN == 0 {
		return sorted
	}
	if topN > len(sorted) {
		topN = len(sorted)
	}
	return sorted[:topN]
}

func sortMeasurements(vals map[string]int64, topN int) resultList {
	sorted := make(resultList, 0)
	for k, v := range vals {
		sorted = append(sorted, result{k, v})
	}
	sort.Sort(sort.Reverse(sorted))

	if topN == 0 {
		return sorted
	}
	if topN > len(sorted) {
		topN = len(sorted)
	}
	return sorted[:topN]
}
<|endoftext|>


<commit_before>package jsonpointer

import (
	"compress/gzip"
	"io/ioutil"
	"os"
	"reflect"
	"testing"

	"github.com/dustin/gojson"
)

var ptests = []struct {
	path string
	exp  interface{}
}{
	{"/foo", []interface{}{"bar", "baz"}},
	{"/foo/0", "bar"},
	{"/", 0.0},
	{"/a~1b", 1.0},
	{"/c%d", 2.0},
	{"/e^f", 3.0},
	{"/g|h", 4.0},
	{"/i\\j", 5.0},
	{"/k\"l", 6.0},
	{"/ ", 7.0},
	{"/m~0n", 8.0},
	{"/g~1n~1r", "has slash, will travel"},
	{"/g/n/r", "where's tito?"},
}

func TestPointerRoot(t *testing.T) {
	got, err := Find([]byte(objSrc), "")
	if err != nil {
		t.Fatalf("Error finding root: %v", err)
	}
	if !reflect.DeepEqual([]byte(objSrc), got) {
		t.Fatalf("Error finding root, found\n%s\n, wanted\n%s",
			got, objSrc)
	}
}

func TestPointerMissing(t *testing.T) {
	got, err := Find([]byte(objSrc), "/missing")
	if err != nil {
		t.Fatalf("Error finding missing item: %v", err)
	}
	if got != nil {
		t.Fatalf("Expected nil looking for /missing, got %v",
			got)
	}
}

func TestManyPointers(t *testing.T) {
	pointers := []string{}
	exp := map[string]interface{}{}
	for _, test := range ptests {
		pointers = append(pointers, test.path)
		exp[test.path] = test.exp
	}

	rv, err := FindMany([]byte(objSrc), pointers)
	if err != nil {
		t.Fatalf("Error finding many: %v", err)
	}

	got := map[string]interface{}{}
	for k, v := range rv {
		var val interface{}
		err = json.Unmarshal(v, &val)
		if err != nil {
			t.Fatalf("Error unmarshaling %s: %v", v, err)
		}
		got[k] = val
	}

	if !reflect.DeepEqual(got, exp) {
		for k, v := range exp {
			if !reflect.DeepEqual(got[k], v) {
				t.Errorf("At %v, expected %#v, got %#v", k, v, got[k])
			}
		}
		t.Fail()
	}
}

func TestManyPointersMissing(t *testing.T) {
	got, err := FindMany([]byte(objSrc), []string{"/missing"})
	if err != nil {
		t.Fatalf("Error finding missing item: %v", err)
	}
	if len(got) != 0 {
		t.Fatalf("Expected empty looking for many /missing, got %v",
			got)
	}

}

func TestPointer(t *testing.T) {

	for _, test := range ptests {
		got, err := Find([]byte(objSrc), test.path)
		var val interface{}
		if err == nil {
			err = json.Unmarshal([]byte(got), &val)
		}
		if err != nil {
			t.Errorf("Got an error on key %v: %v", test.path, err)
			t.Fail()
		} else if !reflect.DeepEqual(val, test.exp) {
			t.Errorf("On %#v, expected %+v (%T), got %+v (%T)",
				test.path, test.exp, test.exp, val, val)
			t.Fail()
		} else {
			t.Logf("Success - got %s for %#v", got, test.path)
		}
	}
}

func TestPointerCoder(t *testing.T) {
	tests := map[string][]string{
		"/":        []string{""},
		"/a":       []string{"a"},
		"/a~1b":    []string{"a/b"},
		"/m~0n":    []string{"m~n"},
		"/ ":       []string{" "},
		"/g~1n~1r": []string{"g/n/r"},
		"/g/n/r":   []string{"g", "n", "r"},
	}

	for k, v := range tests {
		parsed := parsePointer(k)
		encoded := encodePointer(v)

		if k != encoded {
			t.Errorf("Expected to encode %#v as %#v, got %#v",
				v, k, encoded)
			t.Fail()
		}
		if !arreq(v, parsed) {
			t.Errorf("Expected to decode %#v as %#v, got %#v",
				k, v, parsed)
			t.Fail()
		}
	}
}

func BenchmarkPointer(b *testing.B) {
	obj := []byte(objSrc)
	for i := 0; i < b.N; i++ {
		for _, test := range tests {
			Find(obj, test.path)
		}
	}
}

func BenchmarkManyPointer(b *testing.B) {
	pointers := []string{}
	for _, test := range ptests {
		pointers = append(pointers, test.path)
	}
	obj := []byte(objSrc)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		FindMany(obj, pointers)
	}
}

var codeJSON []byte

func codeInit() {
	f, err := os.Open("testdata/code.json.gz")
	if err != nil {
		panic(err)
	}
	defer f.Close()
	gz, err := gzip.NewReader(f)
	if err != nil {
		panic(err)
	}
	data, err := ioutil.ReadAll(gz)
	if err != nil {
		panic(err)
	}

	codeJSON = data
}

func BenchmarkPointerLarge(b *testing.B) {
	keys := []string{
		"/tree/kids/0/kids/0/name",
		"/tree/kids/0/name",
		"/tree/kids/0/kids/0/kids/0/kids/0/kids/0/name",
	}
	if codeJSON == nil {
		b.StopTimer()
		codeInit()
		b.StartTimer()
	}
	b.SetBytes(int64(len(codeJSON)))

	for i := 0; i < b.N; i++ {
		found, err := FindMany(codeJSON, keys)
		if err != nil || len(found) != 3 {
			b.Fatalf("Didn't find all the things from %v/%v",
				found, err)
		}
	}
}

func BenchmarkPointerLargeShallow(b *testing.B) {
	keys := []string{
		"/tree/kids/0/kids/0/kids/0/kids/0/kids/0/name",
	}
	if codeJSON == nil {
		b.StopTimer()
		codeInit()
		b.StartTimer()
	}
	b.SetBytes(int64(len(codeJSON)))

	for i := 0; i < b.N; i++ {
		found, err := FindMany(codeJSON, keys)
		if err != nil || len(found) != 1 {
			b.Fatalf("Didn't find all the things: %v/%v",
				found, err)
		}
	}
}

func BenchmarkPointerSlow(b *testing.B) {
	keys := []string{
		"/tree/kids/0/kids/0/kids/0/kids/0/kids/0/name",
	}
	if codeJSON == nil {
		b.StopTimer()
		codeInit()
		b.StartTimer()
	}
	b.SetBytes(int64(len(codeJSON)))

	for i := 0; i < b.N; i++ {
		m := map[string]interface{}{}
		err := json.Unmarshal(codeJSON, &m)
		if err != nil {
			b.Fatalf("Error parsing JSON: %v", err)
		}
		Get(m, keys[0])
	}
}
<commit_msg>Added benchmarking for missing prefixes.<commit_after>package jsonpointer

import (
	"compress/gzip"
	"io/ioutil"
	"os"
	"reflect"
	"testing"

	"github.com/dustin/gojson"
)

var ptests = []struct {
	path string
	exp  interface{}
}{
	{"/foo", []interface{}{"bar", "baz"}},
	{"/foo/0", "bar"},
	{"/", 0.0},
	{"/a~1b", 1.0},
	{"/c%d", 2.0},
	{"/e^f", 3.0},
	{"/g|h", 4.0},
	{"/i\\j", 5.0},
	{"/k\"l", 6.0},
	{"/ ", 7.0},
	{"/m~0n", 8.0},
	{"/g~1n~1r", "has slash, will travel"},
	{"/g/n/r", "where's tito?"},
}

func TestPointerRoot(t *testing.T) {
	got, err := Find([]byte(objSrc), "")
	if err != nil {
		t.Fatalf("Error finding root: %v", err)
	}
	if !reflect.DeepEqual([]byte(objSrc), got) {
		t.Fatalf("Error finding root, found\n%s\n, wanted\n%s",
			got, objSrc)
	}
}

func TestPointerMissing(t *testing.T) {
	got, err := Find([]byte(objSrc), "/missing")
	if err != nil {
		t.Fatalf("Error finding missing item: %v", err)
	}
	if got != nil {
		t.Fatalf("Expected nil looking for /missing, got %v",
			got)
	}
}

func TestManyPointers(t *testing.T) {
	pointers := []string{}
	exp := map[string]interface{}{}
	for _, test := range ptests {
		pointers = append(pointers, test.path)
		exp[test.path] = test.exp
	}

	rv, err := FindMany([]byte(objSrc), pointers)
	if err != nil {
		t.Fatalf("Error finding many: %v", err)
	}

	got := map[string]interface{}{}
	for k, v := range rv {
		var val interface{}
		err = json.Unmarshal(v, &val)
		if err != nil {
			t.Fatalf("Error unmarshaling %s: %v", v, err)
		}
		got[k] = val
	}

	if !reflect.DeepEqual(got, exp) {
		for k, v := range exp {
			if !reflect.DeepEqual(got[k], v) {
				t.Errorf("At %v, expected %#v, got %#v", k, v, got[k])
			}
		}
		t.Fail()
	}
}

func TestManyPointersMissing(t *testing.T) {
	got, err := FindMany([]byte(objSrc), []string{"/missing"})
	if err != nil {
		t.Fatalf("Error finding missing item: %v", err)
	}
	if len(got) != 0 {
		t.Fatalf("Expected empty looking for many /missing, got %v",
			got)
	}

}

func TestPointer(t *testing.T) {

	for _, test := range ptests {
		got, err := Find([]byte(objSrc), test.path)
		var val interface{}
		if err == nil {
			err = json.Unmarshal([]byte(got), &val)
		}
		if err != nil {
			t.Errorf("Got an error on key %v: %v", test.path, err)
			t.Fail()
		} else if !reflect.DeepEqual(val, test.exp) {
			t.Errorf("On %#v, expected %+v (%T), got %+v (%T)",
				test.path, test.exp, test.exp, val, val)
			t.Fail()
		} else {
			t.Logf("Success - got %s for %#v", got, test.path)
		}
	}
}

func TestPointerCoder(t *testing.T) {
	tests := map[string][]string{
		"/":        []string{""},
		"/a":       []string{"a"},
		"/a~1b":    []string{"a/b"},
		"/m~0n":    []string{"m~n"},
		"/ ":       []string{" "},
		"/g~1n~1r": []string{"g/n/r"},
		"/g/n/r":   []string{"g", "n", "r"},
	}

	for k, v := range tests {
		parsed := parsePointer(k)
		encoded := encodePointer(v)

		if k != encoded {
			t.Errorf("Expected to encode %#v as %#v, got %#v",
				v, k, encoded)
			t.Fail()
		}
		if !arreq(v, parsed) {
			t.Errorf("Expected to decode %#v as %#v, got %#v",
				k, v, parsed)
			t.Fail()
		}
	}
}

func BenchmarkPointer(b *testing.B) {
	obj := []byte(objSrc)
	for i := 0; i < b.N; i++ {
		for _, test := range tests {
			Find(obj, test.path)
		}
	}
}

func BenchmarkManyPointer(b *testing.B) {
	pointers := []string{}
	for _, test := range ptests {
		pointers = append(pointers, test.path)
	}
	obj := []byte(objSrc)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		FindMany(obj, pointers)
	}
}

var codeJSON []byte

func init() {
	f, err := os.Open("testdata/code.json.gz")
	if err != nil {
		panic(err)
	}
	defer f.Close()
	gz, err := gzip.NewReader(f)
	if err != nil {
		panic(err)
	}
	data, err := ioutil.ReadAll(gz)
	if err != nil {
		panic(err)
	}

	codeJSON = data
}

func BenchmarkPointerLarge(b *testing.B) {
	keys := []string{
		"/tree/kids/0/kids/0/name",
		"/tree/kids/0/name",
		"/tree/kids/0/kids/0/kids/0/kids/0/kids/0/name",
	}
	b.SetBytes(int64(len(codeJSON)))

	for i := 0; i < b.N; i++ {
		found, err := FindMany(codeJSON, keys)
		if err != nil || len(found) != 3 {
			b.Fatalf("Didn't find all the things from %v/%v",
				found, err)
		}
	}
}

func BenchmarkPointerLargeShallow(b *testing.B) {
	keys := []string{
		"/tree/kids/0/kids/0/kids/0/kids/0/kids/0/name",
	}
	b.SetBytes(int64(len(codeJSON)))

	for i := 0; i < b.N; i++ {
		found, err := FindMany(codeJSON, keys)
		if err != nil || len(found) != 1 {
			b.Fatalf("Didn't find all the things: %v/%v",
				found, err)
		}
	}
}

func BenchmarkPointerLargeMissing(b *testing.B) {
	keys := []string{
		"/this/does/not/exist",
	}
	b.SetBytes(int64(len(codeJSON)))

	for i := 0; i < b.N; i++ {
		found, err := FindMany(codeJSON, keys)
		if err != nil || len(found) != 0 {
			b.Fatalf("Didn't find all the things: %v/%v",
				found, err)
		}
	}
}

func BenchmarkPointerSlow(b *testing.B) {
	keys := []string{
		"/tree/kids/0/kids/0/kids/0/kids/0/kids/0/name",
	}
	b.SetBytes(int64(len(codeJSON)))

	for i := 0; i < b.N; i++ {
		m := map[string]interface{}{}
		err := json.Unmarshal(codeJSON, &m)
		if err != nil {
			b.Fatalf("Error parsing JSON: %v", err)
		}
		Get(m, keys[0])
	}
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 conn

import (
	"errors"
	"github.com/jacobsa/aws"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestSigner(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type SignerTest struct {
}

func init() { RegisterTestSuite(&SignerTest{}) }

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *SignerTest) CallsFunction() {
	// Function
	var reqArg Request
	var hostArg string

	sts := func(r Request, host string) (string, error) {
		reqArg = r
		hostArg = host
		return "", nil
	}

	// Signer
	signer := newSigner(aws.AccessKey{}, "some_host", sts)

	// Call
	req := Request{"foo": "bar"}
	signer.SignRequest(req)

	ExpectEq(req, reqArg)
	ExpectEq("some_host", hostArg)
}

func (t *SignerTest) FunctionReturnsError() {
	// Function
	sts := func(r Request, h string) (string, error) {
		return "", errors.New("taco")
	}

	// Signer
	signer := newSigner(aws.AccessKey{}, "", sts)

	// Call
	req := Request{}
	err := signer.SignRequest(req)

	ExpectThat(err, Error(HasSubstr("computeStringToSign")))
	ExpectThat(err, Error(HasSubstr("taco")))
}

func (t *SignerTest) FunctionReturnsString() {
	ExpectEq("TODO", "")
}
<commit_msg>SignerTest.FunctionReturnsString<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 conn

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"errors"
	"github.com/jacobsa/aws"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestSigner(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type SignerTest struct {
}

func init() { RegisterTestSuite(&SignerTest{}) }

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *SignerTest) CallsFunction() {
	// Function
	var reqArg Request
	var hostArg string

	sts := func(r Request, host string) (string, error) {
		reqArg = r
		hostArg = host
		return "", nil
	}

	// Signer
	signer := newSigner(aws.AccessKey{}, "some_host", sts)

	// Call
	req := Request{"foo": "bar"}
	signer.SignRequest(req)

	ExpectEq(req, reqArg)
	ExpectEq("some_host", hostArg)
}

func (t *SignerTest) FunctionReturnsError() {
	// Function
	sts := func(r Request, h string) (string, error) {
		return "", errors.New("taco")
	}

	// Signer
	signer := newSigner(aws.AccessKey{}, "", sts)

	// Call
	req := Request{}
	err := signer.SignRequest(req)

	ExpectThat(err, Error(HasSubstr("computeStringToSign")))
	ExpectThat(err, Error(HasSubstr("taco")))
}

func (t *SignerTest) FunctionReturnsString() {
	// Function
	sts := func(r Request, h string) (string, error) {
		return "taco", nil
	}

	// Signer
	key := aws.AccessKey{Id: "queso", Secret: "burrito"}
	signer := newSigner(key, "", sts)

	// Expected output
	h := hmac.New(sha1.New, []byte("burrito"))
	_, err := h.Write([]byte("taco"))
	AssertEq(nil, err)

	buf := new(bytes.Buffer)
	encoder := base64.NewEncoder(base64.StdEncoding, buf)
	_, err = encoder.Write(h.Sum(nil))
	AssertEq(nil, err)
	AssertEq(nil, encoder.Close())

	expected := buf.String()

	// Call
	req := Request{
		"foo": "bar",
	}

	err = signer.SignRequest(req)
	AssertEq(nil, err)

	ExpectEq("bar", req["foo"])
	ExpectEq(expected, req["Signature"])
}
<|endoftext|>


<commit_before>/*
Package compose provides a Go wrapper around Docker Compose, useful for integration testing.

	// Define Compose config.
	var composeYML =`
	test_mockserver:
	  container_name: ms
	  image: jamesdbloom/mockserver
	  ports:
	    - "10000:1080"
	    - "${SOME_ENV_VAR}" # This is replaced with the value of SOME_ENV_VAR.
	test_postgres:
	  container_name: pg
	  image: postgres
	  ports:
	    - "5432"
	`

	// Start containers.
	c, err := compose.Start(composeYML, true, true)
	if err != nil {
		panic(err)
	}
	defer c.Kill()

	// Build MockServer public URL.
	mockServerURL := fmt.Sprintf(
		"http://%v:%v",
		compose.MustInferDockerHost(),
		c.Containers["ms"].MustGetFirstPublicPort(1080, "tcp"))

	// Wait for MockServer to start accepting connections.
	MustConnectWithDefaults(func() error {
		_, err := http.Get(mockServerURL)
		return err
	})
	...
*/
package compose

import (
	"fmt"
	"log"
	"os"
	"regexp"
	"strings"
)

// Compose is the main type exported by the package, used to interact with a running Docker Compose configuration.
type Compose struct {
	fileName   string
	Containers map[string]*Container
}

var (
	logger           = log.New(os.Stdout, "go-compose: ", log.LstdFlags)
	replaceEnvRegexp = regexp.MustCompile("\\$\\{[^\\}]+\\}")
	composeUpRegexp  = regexp.MustCompile("(?m:^docker start <- \\(u'(.*)'\\)$)")
)

const (
	composeProjectName = "compose"
)

// Start starts a Docker Compose configuration.
// If forcePull is true, it attempts do pull newer versions of the images.
// If rmFirst is true, it attempts to kill and delete containers before starting new ones.
func Start(dockerComposeYML string, forcePull, rmFirst bool) (*Compose, error) {
	logger.Println("initializing...")
	dockerComposeYML = replaceEnv(dockerComposeYML)

	fName, err := writeTmp(dockerComposeYML)
	if err != nil {
		return nil, err
	}

	ids, err := composeStart(fName, forcePull, rmFirst)
	if err != nil {
		return nil, err
	}

	containers := make(map[string]*Container)

	for _, id := range ids {
		container, err := Inspect(id)
		if err != nil {
			return nil, err
		}
		if !container.State.Running {
			return nil, fmt.Errorf("compose: container '%v' is not running", container.Name)
		}
		containers[container.Name[1:]] = container
	}

	return &Compose{fileName: fName, Containers: containers}, nil
}

// MustStart is like Start, but panics on error.
func MustStart(dockerComposeYML string, forcePull, killFirst bool) *Compose {
	compose, err := Start(dockerComposeYML, forcePull, killFirst)
	if err != nil {
		panic(err)
	}
	return compose
}

// Kill kills any running containers for the current configuration.
func (c *Compose) Kill() error {
	return composeKill(c.fileName)
}

// MustKill is like Kill, but panics on error.
func (c *Compose) MustKill() {
	if err := c.Kill(); err != nil {
		panic(err)
	}
}

func replaceEnv(dockerComposeYML string) string {
	return replaceEnvRegexp.ReplaceAllStringFunc(dockerComposeYML, replaceEnvFunc)
}

func replaceEnvFunc(s string) string {
	return os.Getenv(strings.TrimSpace(s[2 : len(s)-1]))
}

func composeStart(fName string, forcePull, rmFirst bool) ([]string, error) {
	if forcePull {
		logger.Println("pulling images...")
		if _, err := composeRun(fName, "pull"); err != nil {
			return nil, fmt.Errorf("compose: error pulling images: %v", err)
		}
	}

	if rmFirst {
		if err := composeKill(fName); err != nil {
			return nil, err
		}
		if err := composeRm(fName); err != nil {
			return nil, err
		}
	}

	logger.Println("starting containers...")
	out, err := composeRun(fName, "--verbose", "up", "-d")
	if err != nil {
		return nil, fmt.Errorf("compose: error starting containers: %v", err)
	}
	logger.Println("containers started")

	matches := composeUpRegexp.FindAllStringSubmatch(out, -1)
	ids := make([]string, 0, len(matches))
	for _, match := range matches {
		ids = append(ids, match[1])
	}

	return ids, nil
}

func composeKill(fName string) error {
	logger.Println("killing stale containers...")
	_, err := composeRun(fName, "kill")
	if err != nil {
		return fmt.Errorf("compose: error killing stale containers: %v", err)
	}
	return err
}

func composeRm(fName string) error {
	logger.Println("removing stale containers...")
	_, err := composeRun(fName, "rm", "--force")
	if err != nil {
		return fmt.Errorf("compose: error removing stale containers: %v", err)
	}
	return err
}

func composeRun(fName string, otherArgs ...string) (string, error) {
	args := []string{"-f", fName, "-p", composeProjectName}
	args = append(args, otherArgs...)
	return runCmd("docker-compose", args...)
}
<commit_msg>Fix CircleCI config.<commit_after>/*
Package compose provides a Go wrapper around Docker Compose, useful for integration testing.

	// Define Compose config.
	var composeYML =`
	test_mockserver:
	  container_name: ms
	  image: jamesdbloom/mockserver
	  ports:
	    - "10000:1080"
	    - "${SOME_ENV_VAR}" # This is replaced with the value of SOME_ENV_VAR.
	test_postgres:
	  container_name: pg
	  image: postgres
	  ports:
	    - "5432"
	`

	// Start containers.
	c, err := compose.Start(composeYML, true, true)
	if err != nil {
		panic(err)
	}
	defer c.Kill()

	// Build MockServer public URL.
	mockServerURL := fmt.Sprintf(
		"http://%v:%v",
		compose.MustInferDockerHost(),
		c.Containers["ms"].MustGetFirstPublicPort(1080, "tcp"))

	// Wait for MockServer to start accepting connections.
	MustConnectWithDefaults(func() error {
		_, err := http.Get(mockServerURL)
		return err
	})
	...
*/
package compose

import (
	"fmt"
	"log"
	"os"
	"regexp"
	"strings"
	"os/exec"
)

// Compose is the main type exported by the package, used to interact with a running Docker Compose configuration.
type Compose struct {
	fileName   string
	Containers map[string]*Container
}

var (
	logger           = log.New(os.Stdout, "go-compose: ", log.LstdFlags)
	replaceEnvRegexp = regexp.MustCompile("\\$\\{[^\\}]+\\}")
	composeUpRegexp  = regexp.MustCompile("(?m:^docker start <- \\(u'(.*)'\\)$)")
)

const (
	composeProjectName = "compose"
)

// Start starts a Docker Compose configuration.
// If forcePull is true, it attempts do pull newer versions of the images.
// If rmFirst is true, it attempts to kill and delete containers before starting new ones.
func Start(dockerComposeYML string, forcePull, rmFirst bool) (*Compose, error) {
	logger.Println("initializing...")
	dockerComposeYML = replaceEnv(dockerComposeYML)

	fName, err := writeTmp(dockerComposeYML)
	if err != nil {
		return nil, err
	}

	ids, err := composeStart(fName, forcePull, rmFirst)
	if err != nil {
		return nil, err
	}

	containers := make(map[string]*Container)

	cmd := exec.Command("docker", "ps")
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	cmd.Run()

	for _, id := range ids {
		container, err := Inspect(id)
		if err != nil {
			return nil, err
		}
		if !container.State.Running {
			return nil, fmt.Errorf("compose: container '%v' is not running", container.Name)
		}
		containers[container.Name[1:]] = container
	}

	return &Compose{fileName: fName, Containers: containers}, nil
}

// MustStart is like Start, but panics on error.
func MustStart(dockerComposeYML string, forcePull, killFirst bool) *Compose {
	compose, err := Start(dockerComposeYML, forcePull, killFirst)
	if err != nil {
		panic(err)
	}
	return compose
}

// Kill kills any running containers for the current configuration.
func (c *Compose) Kill() error {
	return composeKill(c.fileName)
}

// MustKill is like Kill, but panics on error.
func (c *Compose) MustKill() {
	if err := c.Kill(); err != nil {
		panic(err)
	}
}

func replaceEnv(dockerComposeYML string) string {
	return replaceEnvRegexp.ReplaceAllStringFunc(dockerComposeYML, replaceEnvFunc)
}

func replaceEnvFunc(s string) string {
	return os.Getenv(strings.TrimSpace(s[2 : len(s)-1]))
}

func composeStart(fName string, forcePull, rmFirst bool) ([]string, error) {
	if forcePull {
		logger.Println("pulling images...")
		if _, err := composeRun(fName, "pull"); err != nil {
			return nil, fmt.Errorf("compose: error pulling images: %v", err)
		}
	}

	if rmFirst {
		if err := composeKill(fName); err != nil {
			return nil, err
		}
		if err := composeRm(fName); err != nil {
			return nil, err
		}
	}

	logger.Println("starting containers...")
	out, err := composeRun(fName, "--verbose", "up", "-d")
	if err != nil {
		return nil, fmt.Errorf("compose: error starting containers: %v", err)
	}
	logger.Println("containers started")

	matches := composeUpRegexp.FindAllStringSubmatch(out, -1)
	ids := make([]string, 0, len(matches))
	for _, match := range matches {
		ids = append(ids, match[1])
	}

	return ids, nil
}

func composeKill(fName string) error {
	logger.Println("killing stale containers...")
	_, err := composeRun(fName, "kill")
	if err != nil {
		return fmt.Errorf("compose: error killing stale containers: %v", err)
	}
	return err
}

func composeRm(fName string) error {
	logger.Println("removing stale containers...")
	_, err := composeRun(fName, "rm", "--force")
	if err != nil {
		return fmt.Errorf("compose: error removing stale containers: %v", err)
	}
	return err
}

func composeRun(fName string, otherArgs ...string) (string, error) {
	args := []string{"-f", fName, "-p", composeProjectName}
	args = append(args, otherArgs...)
	return runCmd("docker-compose", args...)
}
<|endoftext|>


<commit_before>package collector

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"

	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/karimra/gnmic/outputs"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"google.golang.org/grpc"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/encoding/prototext"
	"google.golang.org/protobuf/proto"
)

const (
	defaultTargetReceivebuffer = 1000
)

// Config is the collector config
type Config struct {
	PrometheusAddress   string
	Debug               bool
	Format              string
	TargetReceiveBuffer uint
}

// Collector //
type Collector struct {
	Config        *Config
	Subscriptions map[string]*SubscriptionConfig
	Outputs       map[string][]outputs.Output
	DialOpts      []grpc.DialOption
	//
	m          *sync.Mutex
	Targets    map[string]*Target
	Logger     *log.Logger
	httpServer *http.Server

	ctx      context.Context
	cancelFn context.CancelFunc
}

// NewCollector //
func NewCollector(ctx context.Context,
	config *Config,
	targetConfigs map[string]*TargetConfig,
	subscriptions map[string]*SubscriptionConfig,
	outputs map[string][]outputs.Output,
	dialOpts []grpc.DialOption,
	logger *log.Logger,
) *Collector {
	nctx, cancel := context.WithCancel(ctx)
	grpcMetrics := grpc_prometheus.NewClientMetrics()
	reg := prometheus.NewRegistry()
	reg.MustRegister(prometheus.NewGoCollector())
	reg.MustRegister(prometheus.NewProcessCollector(prometheus.ProcessCollectorOpts{}))
	grpcMetrics.EnableClientHandlingTimeHistogram()
	reg.MustRegister(grpcMetrics)
	handler := http.NewServeMux()
	handler.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))
	httpServer := &http.Server{
		Handler: handler,
		Addr:    config.PrometheusAddress,
	}
	dialOpts = append(dialOpts, grpc.WithStreamInterceptor(grpcMetrics.StreamClientInterceptor()))
	if config.TargetReceiveBuffer == 0 {
		config.TargetReceiveBuffer = defaultTargetReceivebuffer
	}
	c := &Collector{
		Config:        config,
		Subscriptions: subscriptions,
		Outputs:       outputs,
		DialOpts:      dialOpts,
		m:             new(sync.Mutex),
		Targets:       make(map[string]*Target),
		Logger:        logger,
		httpServer:    httpServer,
		ctx:           nctx,
		cancelFn:      cancel,
	}

	for _, tc := range targetConfigs {
		c.InitTarget(tc)
	}
	return c
}

// InitTarget initializes a target based on *TargetConfig
func (c *Collector) InitTarget(tc *TargetConfig) {
	if tc.BufferSize == 0 {
		tc.BufferSize = c.Config.TargetReceiveBuffer
	}
	t := NewTarget(tc)
	//
	t.Subscriptions = make([]*SubscriptionConfig, 0, len(tc.Subscriptions))
	for _, subName := range tc.Subscriptions {
		if sub, ok := c.Subscriptions[subName]; ok {
			t.Subscriptions = append(t.Subscriptions, sub)
		}
	}
	if len(t.Subscriptions) == 0 {
		t.Subscriptions = make([]*SubscriptionConfig, 0, len(c.Subscriptions))
		for _, sub := range c.Subscriptions {
			t.Subscriptions = append(t.Subscriptions, sub)
		}
	}
	//
	t.Outputs = make([]outputs.Output, 0, len(tc.Outputs))
	for _, outName := range tc.Outputs {
		if outs, ok := c.Outputs[outName]; ok {
			t.Outputs = append(t.Outputs, outs...)
		}
	}
	if len(t.Outputs) == 0 {
		t.Outputs = make([]outputs.Output, 0, len(c.Outputs))
		for _, o := range c.Outputs {
			t.Outputs = append(t.Outputs, o...)
		}
	}
	c.m.Lock()
	defer c.m.Unlock()
	c.Targets[t.Config.Name] = t
}

// Subscribe //
func (c *Collector) Subscribe(tName string) error {
	if t, ok := c.Targets[tName]; ok {
		if err := t.CreateGNMIClient(c.ctx, c.DialOpts...); err != nil {
			return err
		}
		c.Logger.Printf("target '%s' gNMI client created", t.Config.Name)
		for _, sc := range t.Subscriptions {
			req, err := sc.CreateSubscribeRequest()
			if err != nil {
				return err
			}
			c.Logger.Printf("sending gNMI SubscribeRequest: subscribe='%+v', mode='%+v', encoding='%+v', to %s",
				req, req.GetSubscribe().GetMode(), req.GetSubscribe().GetEncoding(), t.Config.Name)
			go t.Subscribe(c.ctx, req, sc.Name)
		}
		return nil
	}
	return fmt.Errorf("unknown target name: %s", tName)
}

// Start start the prometheus server as well as a goroutine per target selecting on the response chan, the error chan and the ctx.Done() chan
func (c *Collector) Start() {
	go func() {
		if err := c.httpServer.ListenAndServe(); err != nil {
			c.Logger.Printf("Unable to start prometheus http server: %v", err)
			return
		}
	}()
	defer func() {
		for _, outputs := range c.Outputs {
			for _, o := range outputs {
				o.Close()
			}
		}
	}()
	wg := new(sync.WaitGroup)
	wg.Add(len(c.Targets))
	for _, t := range c.Targets {
		go func(t *Target) {
			defer wg.Done()
			numOnceSubscriptions := t.numberOfOnceSubscriptions()
			remainingOnceSubscriptions := numOnceSubscriptions
			numSubscriptions := len(t.Subscriptions)
			for {
				select {
				case rsp := <-t.SubscribeResponses:
					if c.Config.Debug {
						c.Logger.Printf("received gNMI Subscribe Response: %+v", rsp)
					}
					if c.subscriptionMode(rsp.SubscriptionName) == "ONCE" {
						t.Export(rsp.Response, outputs.Meta{"source": t.Config.Name, "format": c.Config.Format, "subscription-name": rsp.SubscriptionName})
					} else {
						go t.Export(rsp.Response, outputs.Meta{"source": t.Config.Name, "format": c.Config.Format, "subscription-name": rsp.SubscriptionName})
					}
					if remainingOnceSubscriptions > 0 {
						if c.subscriptionMode(rsp.SubscriptionName) == "ONCE" {
							switch rsp.Response.Response.(type) {
							case *gnmi.SubscribeResponse_SyncResponse:
								remainingOnceSubscriptions--
							}
						}
					}
					if remainingOnceSubscriptions == 0 && numSubscriptions == numOnceSubscriptions {
						return
					}
				case err := <-t.Errors:
					if err == io.EOF {
						c.Logger.Printf("target '%s' closed stream(EOF)", t.Config.Name)
						return
					}
					c.Logger.Printf("target '%s' error: %v", t.Config.Name, err)
				}
			}
		}(t)
	}
	wg.Wait()
}

// FormatMsg formats the gnmi.SubscribeResponse and returns a []byte and an error
func FormatMsg(meta map[string]string, rsp proto.Message, multiline bool, indent string) ([]byte, error) {
	if rsp == nil {
		return nil, nil
	}
	switch rsp := rsp.ProtoReflect().Interface().(type) {
	case *gnmi.SubscribeResponse:
		switch rsp := rsp.Response.(type) {
		case *gnmi.SubscribeResponse_Update:
			msg := new(msg)
			msg.Timestamp = rsp.Update.Timestamp
			t := time.Unix(0, rsp.Update.Timestamp)
			msg.Time = &t
			if meta == nil {
				meta = make(map[string]string)
			}
			msg.Prefix = gnmiPathToXPath(rsp.Update.Prefix)
			if s, ok := meta["source"]; ok {
				msg.Source = s
			}
			if s, ok := meta["system-name"]; ok {
				msg.SystemName = s
			}
			if s, ok := meta["subscription-name"]; ok {
				msg.SubscriptionName = s
			}
			for i, upd := range rsp.Update.Update {
				pathElems := make([]string, 0, len(upd.Path.Elem))
				for _, pElem := range upd.Path.Elem {
					pathElems = append(pathElems, pElem.GetName())
				}
				value, err := getValue(upd.Val)
				if err != nil {
					return nil, err
				}
				msg.Updates = append(msg.Updates,
					&update{
						Path:   gnmiPathToXPath(upd.Path),
						Values: make(map[string]interface{}),
					})
				msg.Updates[i].Values[strings.Join(pathElems, "/")] = value
			}
			for _, del := range rsp.Update.Delete {
				msg.Deletes = append(msg.Deletes, gnmiPathToXPath(del))
			}
			if multiline {
				return json.MarshalIndent(msg, "", indent)
			}
			return json.Marshal(msg)
		}
	}
	return nil, nil
}

// TargetPoll sends a gnmi.SubscribeRequest_Poll to targetName and returns the response and an error,
// it uses the targetName and the subscriptionName strings to find the gnmi.GNMI_SubscribeClient
func (c *Collector) TargetPoll(targetName, subscriptionName string) (*gnmi.SubscribeResponse, error) {
	if sub, ok := c.Subscriptions[subscriptionName]; ok {
		if strings.ToUpper(sub.Mode) != "POLL" {
			return nil, fmt.Errorf("subscription '%s' is not a POLL subscription", subscriptionName)
		}
		if t, ok := c.Targets[targetName]; ok {
			if subClient, ok := t.SubscribeClients[subscriptionName]; ok {
				err := subClient.Send(&gnmi.SubscribeRequest{
					Request: &gnmi.SubscribeRequest_Poll{
						Poll: &gnmi.Poll{},
					},
				})
				if err != nil {
					return nil, err
				}
				return subClient.Recv()
			}
		}
		return nil, fmt.Errorf("unknown target name '%s'", targetName)
	}
	return nil, fmt.Errorf("unknown subscription name '%s'", subscriptionName)
}

// PolledSubscriptionsTargets returns a map of target name to a list of subscription names that have Mode == POLL
func (c *Collector) PolledSubscriptionsTargets() map[string][]string {
	result := make(map[string][]string)
	for tn, target := range c.Targets {
		for _, sub := range target.Subscriptions {
			if strings.ToUpper(sub.Mode) == "POLL" {
				if result[tn] == nil {
					result[tn] = make([]string, 0)
				}
				result[tn] = append(result[tn], sub.Name)
			}
		}
	}
	return result
}

func (c *Collector) subscriptionMode(name string) string {
	if sub, ok := c.Subscriptions[name]; ok {
		return strings.ToUpper(sub.Mode)
	}
	return ""
}

// Marshal //
func Marshal(msg proto.Message, format string, meta map[string]string, multiline bool, indent string) ([]byte, error) {
	b := make([]byte, 0)
	var err error
	switch format {
	case "json":
		b, err = FormatMsg(meta, msg, multiline, indent)
	case "proto":
		b, err = proto.Marshal(msg)
	case "protojson":
		b, err = protojson.MarshalOptions{Multiline: multiline, Indent: indent}.Marshal(msg)
	case "prototext":
		b, err = prototext.MarshalOptions{Multiline: multiline, Indent: indent}.Marshal(msg)
	case "event":
		switch msg := msg.ProtoReflect().Interface().(type) {
		case *gnmi.SubscribeResponse:
			var subscriptionName string
			var ok bool
			if subscriptionName, ok = meta["subscription-name"]; !ok {
				subscriptionName = "default"
			}
			switch msg.Response.(type) {
			case *gnmi.SubscribeResponse_Update:
				events, err := ResponseToEventMsgs(subscriptionName, msg, meta)
				if err != nil {
					return nil, fmt.Errorf("failed converting response to events: %v", err)
				}
				if multiline {
					b, err = json.MarshalIndent(events, "", indent)
				} else {
					b, err = json.Marshal(events)
				}
				if err != nil {
					return nil, fmt.Errorf("failed marshaling events: %v", err)
				}
			}
		}
	}
	if err != nil {
		return nil, fmt.Errorf("failed marshaling event: %v", err)
	}
	return b, nil
}
<commit_msg>add MarshalOptions struct<commit_after>package collector

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"

	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/karimra/gnmic/outputs"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"google.golang.org/grpc"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/encoding/prototext"
	"google.golang.org/protobuf/proto"
)

const (
	defaultTargetReceivebuffer = 1000
)

// Config is the collector config
type Config struct {
	PrometheusAddress   string
	Debug               bool
	Format              string
	TargetReceiveBuffer uint
}

// Collector //
type Collector struct {
	Config        *Config
	Subscriptions map[string]*SubscriptionConfig
	Outputs       map[string][]outputs.Output
	DialOpts      []grpc.DialOption
	//
	m          *sync.Mutex
	Targets    map[string]*Target
	Logger     *log.Logger
	httpServer *http.Server

	ctx      context.Context
	cancelFn context.CancelFunc
}

// NewCollector //
func NewCollector(ctx context.Context,
	config *Config,
	targetConfigs map[string]*TargetConfig,
	subscriptions map[string]*SubscriptionConfig,
	outputs map[string][]outputs.Output,
	dialOpts []grpc.DialOption,
	logger *log.Logger,
) *Collector {
	nctx, cancel := context.WithCancel(ctx)
	grpcMetrics := grpc_prometheus.NewClientMetrics()
	reg := prometheus.NewRegistry()
	reg.MustRegister(prometheus.NewGoCollector())
	reg.MustRegister(prometheus.NewProcessCollector(prometheus.ProcessCollectorOpts{}))
	grpcMetrics.EnableClientHandlingTimeHistogram()
	reg.MustRegister(grpcMetrics)
	handler := http.NewServeMux()
	handler.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))
	httpServer := &http.Server{
		Handler: handler,
		Addr:    config.PrometheusAddress,
	}
	dialOpts = append(dialOpts, grpc.WithStreamInterceptor(grpcMetrics.StreamClientInterceptor()))
	if config.TargetReceiveBuffer == 0 {
		config.TargetReceiveBuffer = defaultTargetReceivebuffer
	}
	c := &Collector{
		Config:        config,
		Subscriptions: subscriptions,
		Outputs:       outputs,
		DialOpts:      dialOpts,
		m:             new(sync.Mutex),
		Targets:       make(map[string]*Target),
		Logger:        logger,
		httpServer:    httpServer,
		ctx:           nctx,
		cancelFn:      cancel,
	}

	for _, tc := range targetConfigs {
		c.InitTarget(tc)
	}
	return c
}

// InitTarget initializes a target based on *TargetConfig
func (c *Collector) InitTarget(tc *TargetConfig) {
	if tc.BufferSize == 0 {
		tc.BufferSize = c.Config.TargetReceiveBuffer
	}
	t := NewTarget(tc)
	//
	t.Subscriptions = make([]*SubscriptionConfig, 0, len(tc.Subscriptions))
	for _, subName := range tc.Subscriptions {
		if sub, ok := c.Subscriptions[subName]; ok {
			t.Subscriptions = append(t.Subscriptions, sub)
		}
	}
	if len(t.Subscriptions) == 0 {
		t.Subscriptions = make([]*SubscriptionConfig, 0, len(c.Subscriptions))
		for _, sub := range c.Subscriptions {
			t.Subscriptions = append(t.Subscriptions, sub)
		}
	}
	//
	t.Outputs = make([]outputs.Output, 0, len(tc.Outputs))
	for _, outName := range tc.Outputs {
		if outs, ok := c.Outputs[outName]; ok {
			t.Outputs = append(t.Outputs, outs...)
		}
	}
	if len(t.Outputs) == 0 {
		t.Outputs = make([]outputs.Output, 0, len(c.Outputs))
		for _, o := range c.Outputs {
			t.Outputs = append(t.Outputs, o...)
		}
	}
	c.m.Lock()
	defer c.m.Unlock()
	c.Targets[t.Config.Name] = t
}

// Subscribe //
func (c *Collector) Subscribe(tName string) error {
	if t, ok := c.Targets[tName]; ok {
		if err := t.CreateGNMIClient(c.ctx, c.DialOpts...); err != nil {
			return err
		}
		c.Logger.Printf("target '%s' gNMI client created", t.Config.Name)
		for _, sc := range t.Subscriptions {
			req, err := sc.CreateSubscribeRequest()
			if err != nil {
				return err
			}
			c.Logger.Printf("sending gNMI SubscribeRequest: subscribe='%+v', mode='%+v', encoding='%+v', to %s",
				req, req.GetSubscribe().GetMode(), req.GetSubscribe().GetEncoding(), t.Config.Name)
			go t.Subscribe(c.ctx, req, sc.Name)
		}
		return nil
	}
	return fmt.Errorf("unknown target name: %s", tName)
}

// Start start the prometheus server as well as a goroutine per target selecting on the response chan, the error chan and the ctx.Done() chan
func (c *Collector) Start() {
	go func() {
		if err := c.httpServer.ListenAndServe(); err != nil {
			c.Logger.Printf("Unable to start prometheus http server: %v", err)
			return
		}
	}()
	defer func() {
		for _, outputs := range c.Outputs {
			for _, o := range outputs {
				o.Close()
			}
		}
	}()
	wg := new(sync.WaitGroup)
	wg.Add(len(c.Targets))
	for _, t := range c.Targets {
		go func(t *Target) {
			defer wg.Done()
			numOnceSubscriptions := t.numberOfOnceSubscriptions()
			remainingOnceSubscriptions := numOnceSubscriptions
			numSubscriptions := len(t.Subscriptions)
			for {
				select {
				case rsp := <-t.SubscribeResponses:
					if c.Config.Debug {
						c.Logger.Printf("received gNMI Subscribe Response: %+v", rsp)
					}
					if c.subscriptionMode(rsp.SubscriptionName) == "ONCE" {
						t.Export(rsp.Response, outputs.Meta{"source": t.Config.Name, "format": c.Config.Format, "subscription-name": rsp.SubscriptionName})
					} else {
						go t.Export(rsp.Response, outputs.Meta{"source": t.Config.Name, "format": c.Config.Format, "subscription-name": rsp.SubscriptionName})
					}
					if remainingOnceSubscriptions > 0 {
						if c.subscriptionMode(rsp.SubscriptionName) == "ONCE" {
							switch rsp.Response.Response.(type) {
							case *gnmi.SubscribeResponse_SyncResponse:
								remainingOnceSubscriptions--
							}
						}
					}
					if remainingOnceSubscriptions == 0 && numSubscriptions == numOnceSubscriptions {
						return
					}
				case err := <-t.Errors:
					if err == io.EOF {
						c.Logger.Printf("target '%s' closed stream(EOF)", t.Config.Name)
						return
					}
					c.Logger.Printf("target '%s' error: %v", t.Config.Name, err)
				}
			}
		}(t)
	}
	wg.Wait()
}

// FormatMsg formats the gnmi.SubscribeResponse and returns a []byte and an error
func (o *MarshalOptions) FormatJSON(rsp proto.Message, meta map[string]string) ([]byte, error) {
	if rsp == nil {
		return nil, nil
	}
	switch rsp := rsp.ProtoReflect().Interface().(type) {
	case *gnmi.SubscribeResponse:
		switch rsp := rsp.Response.(type) {
		case *gnmi.SubscribeResponse_Update:
			msg := new(subscribeRspMsg)
			msg.Timestamp = rsp.Update.Timestamp
			t := time.Unix(0, rsp.Update.Timestamp)
			msg.Time = &t
			if meta == nil {
				meta = make(map[string]string)
			}
			msg.Prefix = gnmiPathToXPath(rsp.Update.Prefix)
			if s, ok := meta["source"]; ok {
				msg.Source = s
			}
			if s, ok := meta["system-name"]; ok {
				msg.SystemName = s
			}
			if s, ok := meta["subscription-name"]; ok {
				msg.SubscriptionName = s
			}
			for i, upd := range rsp.Update.Update {
				pathElems := make([]string, 0, len(upd.Path.Elem))
				for _, pElem := range upd.Path.Elem {
					pathElems = append(pathElems, pElem.GetName())
				}
				value, err := getValue(upd.Val)
				if err != nil {
					return nil, err
				}
				msg.Updates = append(msg.Updates,
					&update{
						Path:   gnmiPathToXPath(upd.Path),
						Values: make(map[string]interface{}),
					})
				msg.Updates[i].Values[strings.Join(pathElems, "/")] = value
			}
			for _, del := range rsp.Update.Delete {
				msg.Deletes = append(msg.Deletes, gnmiPathToXPath(del))
			}
			if o.Multiline {
				return json.MarshalIndent(msg, o.Prefix, o.Indent)
			}
			return json.Marshal(msg)
		}
	}
	return nil, nil
}

// TargetPoll sends a gnmi.SubscribeRequest_Poll to targetName and returns the response and an error,
// it uses the targetName and the subscriptionName strings to find the gnmi.GNMI_SubscribeClient
func (c *Collector) TargetPoll(targetName, subscriptionName string) (*gnmi.SubscribeResponse, error) {
	if sub, ok := c.Subscriptions[subscriptionName]; ok {
		if strings.ToUpper(sub.Mode) != "POLL" {
			return nil, fmt.Errorf("subscription '%s' is not a POLL subscription", subscriptionName)
		}
		if t, ok := c.Targets[targetName]; ok {
			if subClient, ok := t.SubscribeClients[subscriptionName]; ok {
				err := subClient.Send(&gnmi.SubscribeRequest{
					Request: &gnmi.SubscribeRequest_Poll{
						Poll: &gnmi.Poll{},
					},
				})
				if err != nil {
					return nil, err
				}
				return subClient.Recv()
			}
		}
		return nil, fmt.Errorf("unknown target name '%s'", targetName)
	}
	return nil, fmt.Errorf("unknown subscription name '%s'", subscriptionName)
}

// PolledSubscriptionsTargets returns a map of target name to a list of subscription names that have Mode == POLL
func (c *Collector) PolledSubscriptionsTargets() map[string][]string {
	result := make(map[string][]string)
	for tn, target := range c.Targets {
		for _, sub := range target.Subscriptions {
			if strings.ToUpper(sub.Mode) == "POLL" {
				if result[tn] == nil {
					result[tn] = make([]string, 0)
				}
				result[tn] = append(result[tn], sub.Name)
			}
		}
	}
	return result
}

func (c *Collector) subscriptionMode(name string) string {
	if sub, ok := c.Subscriptions[name]; ok {
		return strings.ToUpper(sub.Mode)
	}
	return ""
}

// Marshal //
func (o *MarshalOptions) Marshal(msg proto.Message, meta map[string]string) ([]byte, error) {
	b := make([]byte, 0)
	var err error
	switch o.Format {
	default: // json
		b, err = o.FormatJSON(msg, meta)
	case "proto":
		b, err = proto.Marshal(msg)
	case "protojson":
		b, err = protojson.MarshalOptions{Multiline: o.Multiline, Indent: o.Indent}.Marshal(msg)
	case "prototext":
		b, err = prototext.MarshalOptions{Multiline: o.Multiline, Indent: o.Indent}.Marshal(msg)
	case "event":
		switch msg := msg.ProtoReflect().Interface().(type) {
		case *gnmi.SubscribeResponse:
			var subscriptionName string
			var ok bool
			if subscriptionName, ok = meta["subscription-name"]; !ok {
				subscriptionName = "default"
			}
			switch msg.Response.(type) {
			case *gnmi.SubscribeResponse_Update:
				events, err := ResponseToEventMsgs(subscriptionName, msg, meta)
				if err != nil {
					return nil, fmt.Errorf("failed converting response to events: %v", err)
				}
				if o.Multiline {
					b, err = json.MarshalIndent(events, "", o.Indent)
				} else {
					b, err = json.Marshal(events)
				}
				if err != nil {
					return nil, fmt.Errorf("failed marshaling events: %v", err)
				}
			}
		}
	}
	if err != nil {
		return nil, fmt.Errorf("failed marshaling event: %v", err)
	}
	return b, nil
}

type MarshalOptions struct {
	Multiline bool
	Indent    string
	Prefix    string
	Format    string // TODO
}
<|endoftext|>