<commit_before>package main

import (
	"bufio"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/jpillora/backoff"
	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
)

type LogEvent struct {
	Pod       *v1.Pod
	Container *v1.Container
	Timestamp *time.Time
	Message   string
}

type LogEventFunc func(LogEvent)

func NewContainerTailer(
	clientset *kubernetes.Clientset,
	pod v1.Pod,
	container v1.Container,
	eventFunc LogEventFunc) *ContainerTailer {
	return &ContainerTailer{
		clientset: clientset,
		pod:       pod,
		container: container,
		eventFunc: eventFunc,
	}
}

type ContainerTailer struct {
	clientset *kubernetes.Clientset
	pod       v1.Pod
	container v1.Container
	stop      bool
	eventFunc LogEventFunc
}

func (ct *ContainerTailer) Stop() {
	ct.stop = true
}

func (ct *ContainerTailer) Run() error {
	for !ct.stop {
		stream, err := ct.getStream()
		if err != nil {
			return err
		}
		if stream == nil {
			break
		}
		if err := ct.runStream(stream); err != nil {
			return err
		}
	}
	return nil
}

func (ct *ContainerTailer) runStream(stream io.ReadCloser) error {
	defer func() {
		_ = stream.Close()
	}()

	r := bufio.NewReader(stream)
	for {
		line, err := r.ReadString('\n')
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		ct.receiveLine(line)
	}
	return nil
}

func (ct *ContainerTailer) receiveLine(s string) {
	if len(s) > 0 && s[len(s)-1] == '\n' {
		s = s[0 : len(s)-1]
	}
	for len(s) > 0 && s[len(s)-1] == '\r' {
		s = s[0 : len(s)-1]
	}

	parts := strings.SplitN(s, " ", 2)

	var timestamp *time.Time
	if t, err := time.Parse(time.RFC3339Nano, parts[0]); err == nil {
		timestamp = &t
	}

	ct.eventFunc(LogEvent{
		Pod:       &ct.pod,
		Container: &ct.container,
		Timestamp: timestamp,
		Message:   parts[1],
	})
}

func (ct *ContainerTailer) getStream() (io.ReadCloser, error) {
	sinceSeconds := int64(1)

	boff := &backoff.Backoff{}
	for {
		stream, err := ct.clientset.Core().Pods(ct.pod.Namespace).GetLogs(ct.pod.Name, &v1.PodLogOptions{
			Container:    ct.container.Name,
			Follow:       true,
			Timestamps:   true,
			SinceSeconds: &sinceSeconds,
		}).Stream()
		if err != nil {
			if status, ok := err.(errors.APIStatus); ok {
				// This will happen if the pod isn't ready for log-reading yet
				switch status.Status().Code {
				case http.StatusBadRequest:
					time.Sleep(boff.Duration())
					continue
				case http.StatusNotFound:
					return nil, nil
				}
			}
			return nil, err
		}
		return stream, nil
	}
}
<commit_msg>Fix a potential out-of-bounds array index error.<commit_after>package main

import (
	"bufio"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/jpillora/backoff"
	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
)

type LogEvent struct {
	Pod       *v1.Pod
	Container *v1.Container
	Timestamp *time.Time
	Message   string
}

type LogEventFunc func(LogEvent)

func NewContainerTailer(
	clientset *kubernetes.Clientset,
	pod v1.Pod,
	container v1.Container,
	eventFunc LogEventFunc) *ContainerTailer {
	return &ContainerTailer{
		clientset: clientset,
		pod:       pod,
		container: container,
		eventFunc: eventFunc,
	}
}

type ContainerTailer struct {
	clientset *kubernetes.Clientset
	pod       v1.Pod
	container v1.Container
	stop      bool
	eventFunc LogEventFunc
}

func (ct *ContainerTailer) Stop() {
	ct.stop = true
}

func (ct *ContainerTailer) Run() error {
	for !ct.stop {
		stream, err := ct.getStream()
		if err != nil {
			return err
		}
		if stream == nil {
			break
		}
		if err := ct.runStream(stream); err != nil {
			return err
		}
	}
	return nil
}

func (ct *ContainerTailer) runStream(stream io.ReadCloser) error {
	defer func() {
		_ = stream.Close()
	}()

	r := bufio.NewReader(stream)
	for {
		line, err := r.ReadString('\n')
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		ct.receiveLine(line)
	}
	return nil
}

func (ct *ContainerTailer) receiveLine(s string) {
	if len(s) > 0 && s[len(s)-1] == '\n' {
		s = s[0 : len(s)-1]
	}
	for len(s) > 0 && s[len(s)-1] == '\r' {
		s = s[0 : len(s)-1]
	}

	parts := strings.SplitN(s, " ", 2)
	if len(parts) < 2 {
		// TODO: Warn?
		return
	}

	var timestamp *time.Time
	if t, err := time.Parse(time.RFC3339Nano, parts[0]); err == nil {
		timestamp = &t
	}

	ct.eventFunc(LogEvent{
		Pod:       &ct.pod,
		Container: &ct.container,
		Timestamp: timestamp,
		Message:   parts[1],
	})
}

func (ct *ContainerTailer) getStream() (io.ReadCloser, error) {
	sinceSeconds := int64(1)

	boff := &backoff.Backoff{}
	for {
		stream, err := ct.clientset.Core().Pods(ct.pod.Namespace).GetLogs(ct.pod.Name, &v1.PodLogOptions{
			Container:    ct.container.Name,
			Follow:       true,
			Timestamps:   true,
			SinceSeconds: &sinceSeconds,
		}).Stream()
		if err != nil {
			if status, ok := err.(errors.APIStatus); ok {
				// This will happen if the pod isn't ready for log-reading yet
				switch status.Status().Code {
				case http.StatusBadRequest:
					time.Sleep(boff.Duration())
					continue
				case http.StatusNotFound:
					return nil, nil
				}
			}
			return nil, err
		}
		return stream, nil
	}
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	"time"

	_ "github.com/mattn/go-sqlite3"

	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/osarch"
)

var dbImageSourceProtocol = map[int]string{
	0: "lxd",
	1: "direct",
	2: "simplestreams",
}

func dbImagesGet(db *sql.DB, public bool) ([]string, error) {
	q := "SELECT fingerprint FROM images"
	if public == true {
		q = "SELECT fingerprint FROM images WHERE public=1"
	}

	var fp string
	inargs := []interface{}{}
	outfmt := []interface{}{fp}
	dbResults, err := dbQueryScan(db, q, inargs, outfmt)
	if err != nil {
		return []string{}, err
	}

	results := []string{}
	for _, r := range dbResults {
		results = append(results, r[0].(string))
	}

	return results, nil
}

func dbImagesGetExpired(db *sql.DB, expiry int64) ([]string, error) {
	q := `SELECT fingerprint FROM images WHERE cached=1 AND creation_date<=strftime('%s', date('now', '-` + fmt.Sprintf("%d", expiry) + ` day'))`

	var fp string
	inargs := []interface{}{}
	outfmt := []interface{}{fp}
	dbResults, err := dbQueryScan(db, q, inargs, outfmt)
	if err != nil {
		return []string{}, err
	}

	results := []string{}
	for _, r := range dbResults {
		results = append(results, r[0].(string))
	}

	return results, nil
}

func dbImageSourceInsert(db *sql.DB, imageId int, server string, protocol string, certificate string, alias string) error {
	stmt := `INSERT INTO images_source (image_id, server, protocol, certificate, alias) values (?, ?, ?, ?, ?)`

	protocolInt := -1
	for protoInt, protoString := range dbImageSourceProtocol {
		if protoString == protocol {
			protocolInt = protoInt
		}
	}

	if protocolInt == -1 {
		return fmt.Errorf("Invalid protocol: %s", protocol)
	}

	_, err := dbExec(db, stmt, imageId, server, protocolInt, certificate, alias)
	return err
}

func dbImageSourceGet(db *sql.DB, imageId int) (int, api.ImageSource, error) {
	q := `SELECT id, server, protocol, certificate, alias FROM images_source WHERE image_id=?`

	id := 0
	protocolInt := -1
	result := api.ImageSource{}

	arg1 := []interface{}{imageId}
	arg2 := []interface{}{&id, &result.Server, &protocolInt, &result.Certificate, &result.Alias}
	err := dbQueryRowScan(db, q, arg1, arg2)
	if err != nil {
		if err == sql.ErrNoRows {
			return -1, api.ImageSource{}, NoSuchObjectError
		}

		return -1, api.ImageSource{}, err
	}

	protocol, found := dbImageSourceProtocol[protocolInt]
	if !found {
		return -1, api.ImageSource{}, fmt.Errorf("Invalid protocol: %d", protocolInt)
	}

	result.Protocol = protocol

	return id, result, nil

}

// dbImageGet gets an ImageBaseInfo object from the database.
// The argument fingerprint will be queried with a LIKE query, means you can
// pass a shortform and will get the full fingerprint.
// There can never be more than one image with a given fingerprint, as it is
// enforced by a UNIQUE constraint in the schema.
func dbImageGet(db *sql.DB, fingerprint string, public bool, strictMatching bool) (int, *api.Image, error) {
	var err error
	var create, expire, used, upload *time.Time // These hold the db-returned times

	// The object we'll actually return
	image := api.Image{}
	id := -1
	arch := -1

	// These two humongous things will be filled by the call to DbQueryRowScan
	outfmt := []interface{}{&id, &image.Fingerprint, &image.Filename,
		&image.Size, &image.Cached, &image.Public, &image.AutoUpdate, &arch,
		&create, &expire, &used, &upload}

	var query string

	var inargs []interface{}
	if strictMatching {
		inargs = []interface{}{fingerprint}
		query = `
        SELECT
            id, fingerprint, filename, size, cached, public, auto_update, architecture,
            creation_date, expiry_date, last_use_date, upload_date
        FROM
            images
        WHERE fingerprint = ?`
	} else {
		inargs = []interface{}{fingerprint + "%"}
		query = `
        SELECT
            id, fingerprint, filename, size, cached, public, auto_update, architecture,
            creation_date, expiry_date, last_use_date, upload_date
        FROM
            images
        WHERE fingerprint LIKE ?`
	}

	if public {
		query = query + " AND public=1"
	}

	err = dbQueryRowScan(db, query, inargs, outfmt)

	if err != nil {
		return -1, nil, err // Likely: there are no rows for this fingerprint
	}

	// Some of the dates can be nil in the DB, let's process them.
	if create != nil {
		image.CreatedAt = *create
	} else {
		image.CreatedAt = time.Time{}
	}

	if expire != nil {
		image.ExpiresAt = *expire
	} else {
		image.ExpiresAt = time.Time{}
	}

	if used != nil {
		image.LastUsedAt = *used
	} else {
		image.LastUsedAt = time.Time{}
	}

	image.Architecture, _ = osarch.ArchitectureName(arch)

	// The upload date is enforced by NOT NULL in the schema, so it can never be nil.
	image.UploadedAt = *upload

	// Get the properties
	q := "SELECT key, value FROM images_properties where image_id=?"
	var key, value, name, desc string
	inargs = []interface{}{id}
	outfmt = []interface{}{key, value}
	results, err := dbQueryScan(db, q, inargs, outfmt)
	if err != nil {
		return -1, nil, err
	}

	properties := map[string]string{}
	for _, r := range results {
		key = r[0].(string)
		value = r[1].(string)
		properties[key] = value
	}

	image.Properties = properties

	// Get the aliases
	q = "SELECT name, description FROM images_aliases WHERE image_id=?"
	inargs = []interface{}{id}
	outfmt = []interface{}{name, desc}
	results, err = dbQueryScan(db, q, inargs, outfmt)
	if err != nil {
		return -1, nil, err
	}

	aliases := []api.ImageAlias{}
	for _, r := range results {
		name = r[0].(string)
		desc = r[0].(string)
		a := api.ImageAlias{Name: name, Description: desc}
		aliases = append(aliases, a)
	}

	image.Aliases = aliases

	_, source, err := dbImageSourceGet(db, id)
	if err == nil {
		image.UpdateSource = &source
	}

	return id, &image, nil
}

func dbImageDelete(db *sql.DB, id int) error {
	tx, err := dbBegin(db)
	if err != nil {
		return err
	}

	_, _ = tx.Exec("DELETE FROM images_aliases WHERE image_id=?", id)
	_, _ = tx.Exec("DELETE FROM images_properties WHERE image_id=?", id)
	_, _ = tx.Exec("DELETE FROM images_source WHERE image_id=?", id)
	_, _ = tx.Exec("DELETE FROM images WHERE id=?", id)

	if err := txCommit(tx); err != nil {
		return err
	}

	return nil
}

func dbImageAliasGet(db *sql.DB, name string, isTrustedClient bool) (int, api.ImageAliasesEntry, error) {
	q := `SELECT images_aliases.id, images.fingerprint, images_aliases.description
			 FROM images_aliases
			 INNER JOIN images
			 ON images_aliases.image_id=images.id
			 WHERE images_aliases.name=?`
	if !isTrustedClient {
		q = q + ` AND images.public=1`
	}

	var fingerprint, description string
	id := -1
	entry := api.ImageAliasesEntry{}

	arg1 := []interface{}{name}
	arg2 := []interface{}{&id, &fingerprint, &description}
	err := dbQueryRowScan(db, q, arg1, arg2)
	if err != nil {
		if err == sql.ErrNoRows {
			return -1, entry, NoSuchObjectError
		}

		return -1, entry, err
	}

	entry.Name = name
	entry.Target = fingerprint
	entry.Description = description

	return id, entry, nil
}

func dbImageAliasRename(db *sql.DB, id int, name string) error {
	_, err := dbExec(db, "UPDATE images_aliases SET name=? WHERE id=?", name, id)
	return err
}

func dbImageAliasDelete(db *sql.DB, name string) error {
	_, err := dbExec(db, "DELETE FROM images_aliases WHERE name=?", name)
	return err
}

func dbImageAliasesMove(db *sql.DB, source int, destination int) error {
	_, err := dbExec(db, "UPDATE images_aliases SET image_id=? WHERE image_id=?", destination, source)
	return err
}

// Insert an alias ento the database.
func dbImageAliasAdd(db *sql.DB, name string, imageID int, desc string) error {
	stmt := `INSERT INTO images_aliases (name, image_id, description) values (?, ?, ?)`
	_, err := dbExec(db, stmt, name, imageID, desc)
	return err
}

func dbImageAliasUpdate(db *sql.DB, id int, imageID int, desc string) error {
	stmt := `UPDATE images_aliases SET image_id=?, description=? WHERE id=?`
	_, err := dbExec(db, stmt, imageID, desc, id)
	return err
}

func dbImageLastAccessUpdate(db *sql.DB, fingerprint string, date time.Time) error {
	stmt := `UPDATE images SET last_use_date=? WHERE fingerprint=?`
	_, err := dbExec(db, stmt, date, fingerprint)
	return err
}

func dbImageLastAccessInit(db *sql.DB, fingerprint string) error {
	stmt := `UPDATE images SET cached=1, last_use_date=strftime("%s") WHERE fingerprint=?`
	_, err := dbExec(db, stmt, fingerprint)
	return err
}

func dbImageUpdate(db *sql.DB, id int, fname string, sz int64, public bool, autoUpdate bool, architecture string, createdAt time.Time, expiresAt time.Time, properties map[string]string) error {
	arch, err := osarch.ArchitectureId(architecture)
	if err != nil {
		arch = 0
	}

	tx, err := dbBegin(db)
	if err != nil {
		return err
	}

	publicInt := 0
	if public {
		publicInt = 1
	}

	autoUpdateInt := 0
	if autoUpdate {
		autoUpdateInt = 1
	}

	stmt, err := tx.Prepare(`UPDATE images SET filename=?, size=?, public=?, auto_update=?, architecture=?, creation_date=?, expiry_date=? WHERE id=?`)
	if err != nil {
		tx.Rollback()
		return err
	}
	defer stmt.Close()

	_, err = stmt.Exec(fname, sz, publicInt, autoUpdateInt, arch, createdAt, expiresAt, id)
	if err != nil {
		tx.Rollback()
		return err
	}

	_, err = tx.Exec(`DELETE FROM images_properties WHERE image_id=?`, id)

	stmt, err = tx.Prepare(`INSERT INTO images_properties (image_id, type, key, value) VALUES (?, ?, ?, ?)`)
	if err != nil {
		tx.Rollback()
		return err
	}

	for key, value := range properties {
		_, err = stmt.Exec(id, 0, key, value)
		if err != nil {
			tx.Rollback()
			return err
		}
	}

	if err := txCommit(tx); err != nil {
		return err
	}

	return nil
}

func dbImageInsert(db *sql.DB, fp string, fname string, sz int64, public bool, autoUpdate bool, architecture string, createdAt time.Time, expiresAt time.Time, properties map[string]string) error {
	arch, err := osarch.ArchitectureId(architecture)
	if err != nil {
		arch = 0
	}

	tx, err := dbBegin(db)
	if err != nil {
		return err
	}

	publicInt := 0
	if public {
		publicInt = 1
	}

	autoUpdateInt := 0
	if autoUpdate {
		autoUpdateInt = 1
	}

	stmt, err := tx.Prepare(`INSERT INTO images (fingerprint, filename, size, public, auto_update, architecture, creation_date, expiry_date, upload_date) VALUES (?, ?, ?, ?, ?, ?, ?, ?, strftime("%s"))`)
	if err != nil {
		tx.Rollback()
		return err
	}
	defer stmt.Close()

	result, err := stmt.Exec(fp, fname, sz, publicInt, autoUpdateInt, arch, createdAt, expiresAt)
	if err != nil {
		tx.Rollback()
		return err
	}

	if len(properties) > 0 {
		id64, err := result.LastInsertId()
		if err != nil {
			tx.Rollback()
			return err
		}
		id := int(id64)

		pstmt, err := tx.Prepare(`INSERT INTO images_properties (image_id, type, key, value) VALUES (?, 0, ?, ?)`)
		if err != nil {
			tx.Rollback()
			return err
		}
		defer pstmt.Close()

		for k, v := range properties {

			// we can assume, that there is just one
			// value per key
			_, err = pstmt.Exec(id, k, v)
			if err != nil {
				tx.Rollback()
				return err
			}
		}

	}

	if err := txCommit(tx); err != nil {
		return err
	}

	return nil
}
<commit_msg>Fix partial image fingerprint matches<commit_after>package main

import (
	"database/sql"
	"fmt"
	"time"

	_ "github.com/mattn/go-sqlite3"

	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/osarch"
)

var dbImageSourceProtocol = map[int]string{
	0: "lxd",
	1: "direct",
	2: "simplestreams",
}

func dbImagesGet(db *sql.DB, public bool) ([]string, error) {
	q := "SELECT fingerprint FROM images"
	if public == true {
		q = "SELECT fingerprint FROM images WHERE public=1"
	}

	var fp string
	inargs := []interface{}{}
	outfmt := []interface{}{fp}
	dbResults, err := dbQueryScan(db, q, inargs, outfmt)
	if err != nil {
		return []string{}, err
	}

	results := []string{}
	for _, r := range dbResults {
		results = append(results, r[0].(string))
	}

	return results, nil
}

func dbImagesGetExpired(db *sql.DB, expiry int64) ([]string, error) {
	q := `SELECT fingerprint FROM images WHERE cached=1 AND creation_date<=strftime('%s', date('now', '-` + fmt.Sprintf("%d", expiry) + ` day'))`

	var fp string
	inargs := []interface{}{}
	outfmt := []interface{}{fp}
	dbResults, err := dbQueryScan(db, q, inargs, outfmt)
	if err != nil {
		return []string{}, err
	}

	results := []string{}
	for _, r := range dbResults {
		results = append(results, r[0].(string))
	}

	return results, nil
}

func dbImageSourceInsert(db *sql.DB, imageId int, server string, protocol string, certificate string, alias string) error {
	stmt := `INSERT INTO images_source (image_id, server, protocol, certificate, alias) values (?, ?, ?, ?, ?)`

	protocolInt := -1
	for protoInt, protoString := range dbImageSourceProtocol {
		if protoString == protocol {
			protocolInt = protoInt
		}
	}

	if protocolInt == -1 {
		return fmt.Errorf("Invalid protocol: %s", protocol)
	}

	_, err := dbExec(db, stmt, imageId, server, protocolInt, certificate, alias)
	return err
}

func dbImageSourceGet(db *sql.DB, imageId int) (int, api.ImageSource, error) {
	q := `SELECT id, server, protocol, certificate, alias FROM images_source WHERE image_id=?`

	id := 0
	protocolInt := -1
	result := api.ImageSource{}

	arg1 := []interface{}{imageId}
	arg2 := []interface{}{&id, &result.Server, &protocolInt, &result.Certificate, &result.Alias}
	err := dbQueryRowScan(db, q, arg1, arg2)
	if err != nil {
		if err == sql.ErrNoRows {
			return -1, api.ImageSource{}, NoSuchObjectError
		}

		return -1, api.ImageSource{}, err
	}

	protocol, found := dbImageSourceProtocol[protocolInt]
	if !found {
		return -1, api.ImageSource{}, fmt.Errorf("Invalid protocol: %d", protocolInt)
	}

	result.Protocol = protocol

	return id, result, nil

}

// dbImageGet gets an ImageBaseInfo object from the database.
// The argument fingerprint will be queried with a LIKE query, means you can
// pass a shortform and will get the full fingerprint.
// There can never be more than one image with a given fingerprint, as it is
// enforced by a UNIQUE constraint in the schema.
func dbImageGet(db *sql.DB, fingerprint string, public bool, strictMatching bool) (int, *api.Image, error) {
	var err error
	var create, expire, used, upload *time.Time // These hold the db-returned times

	// The object we'll actually return
	image := api.Image{}
	id := -1
	arch := -1

	// These two humongous things will be filled by the call to DbQueryRowScan
	outfmt := []interface{}{&id, &image.Fingerprint, &image.Filename,
		&image.Size, &image.Cached, &image.Public, &image.AutoUpdate, &arch,
		&create, &expire, &used, &upload}

	var query string

	var inargs []interface{}
	if strictMatching {
		inargs = []interface{}{fingerprint}
		query = `
        SELECT
            id, fingerprint, filename, size, cached, public, auto_update, architecture,
            creation_date, expiry_date, last_use_date, upload_date
        FROM
            images
        WHERE fingerprint = ?`
	} else {
		inargs = []interface{}{fingerprint + "%"}
		query = `
        SELECT
            id, fingerprint, filename, size, cached, public, auto_update, architecture,
            creation_date, expiry_date, last_use_date, upload_date
        FROM
            images
        WHERE fingerprint LIKE ?`
	}

	if public {
		query = query + " AND public=1"
	}

	err = dbQueryRowScan(db, query, inargs, outfmt)
	if err != nil {
		return -1, nil, err // Likely: there are no rows for this fingerprint
	}

	// Validate we only have a single match
	if !strictMatching {
		query = "SELECT COUNT(id) FROM images WHERE fingerprint LIKE ?"
		count := 0
		outfmt := []interface{}{&count}

		err = dbQueryRowScan(db, query, inargs, outfmt)
		if err != nil {
			return -1, nil, err
		}

		if count > 1 {
			return -1, nil, fmt.Errorf("Partial fingerprint matches more than one image")
		}
	}

	// Some of the dates can be nil in the DB, let's process them.
	if create != nil {
		image.CreatedAt = *create
	} else {
		image.CreatedAt = time.Time{}
	}

	if expire != nil {
		image.ExpiresAt = *expire
	} else {
		image.ExpiresAt = time.Time{}
	}

	if used != nil {
		image.LastUsedAt = *used
	} else {
		image.LastUsedAt = time.Time{}
	}

	image.Architecture, _ = osarch.ArchitectureName(arch)

	// The upload date is enforced by NOT NULL in the schema, so it can never be nil.
	image.UploadedAt = *upload

	// Get the properties
	q := "SELECT key, value FROM images_properties where image_id=?"
	var key, value, name, desc string
	inargs = []interface{}{id}
	outfmt = []interface{}{key, value}
	results, err := dbQueryScan(db, q, inargs, outfmt)
	if err != nil {
		return -1, nil, err
	}

	properties := map[string]string{}
	for _, r := range results {
		key = r[0].(string)
		value = r[1].(string)
		properties[key] = value
	}

	image.Properties = properties

	// Get the aliases
	q = "SELECT name, description FROM images_aliases WHERE image_id=?"
	inargs = []interface{}{id}
	outfmt = []interface{}{name, desc}
	results, err = dbQueryScan(db, q, inargs, outfmt)
	if err != nil {
		return -1, nil, err
	}

	aliases := []api.ImageAlias{}
	for _, r := range results {
		name = r[0].(string)
		desc = r[0].(string)
		a := api.ImageAlias{Name: name, Description: desc}
		aliases = append(aliases, a)
	}

	image.Aliases = aliases

	_, source, err := dbImageSourceGet(db, id)
	if err == nil {
		image.UpdateSource = &source
	}

	return id, &image, nil
}

func dbImageDelete(db *sql.DB, id int) error {
	tx, err := dbBegin(db)
	if err != nil {
		return err
	}

	_, _ = tx.Exec("DELETE FROM images_aliases WHERE image_id=?", id)
	_, _ = tx.Exec("DELETE FROM images_properties WHERE image_id=?", id)
	_, _ = tx.Exec("DELETE FROM images_source WHERE image_id=?", id)
	_, _ = tx.Exec("DELETE FROM images WHERE id=?", id)

	if err := txCommit(tx); err != nil {
		return err
	}

	return nil
}

func dbImageAliasGet(db *sql.DB, name string, isTrustedClient bool) (int, api.ImageAliasesEntry, error) {
	q := `SELECT images_aliases.id, images.fingerprint, images_aliases.description
			 FROM images_aliases
			 INNER JOIN images
			 ON images_aliases.image_id=images.id
			 WHERE images_aliases.name=?`
	if !isTrustedClient {
		q = q + ` AND images.public=1`
	}

	var fingerprint, description string
	id := -1
	entry := api.ImageAliasesEntry{}

	arg1 := []interface{}{name}
	arg2 := []interface{}{&id, &fingerprint, &description}
	err := dbQueryRowScan(db, q, arg1, arg2)
	if err != nil {
		if err == sql.ErrNoRows {
			return -1, entry, NoSuchObjectError
		}

		return -1, entry, err
	}

	entry.Name = name
	entry.Target = fingerprint
	entry.Description = description

	return id, entry, nil
}

func dbImageAliasRename(db *sql.DB, id int, name string) error {
	_, err := dbExec(db, "UPDATE images_aliases SET name=? WHERE id=?", name, id)
	return err
}

func dbImageAliasDelete(db *sql.DB, name string) error {
	_, err := dbExec(db, "DELETE FROM images_aliases WHERE name=?", name)
	return err
}

func dbImageAliasesMove(db *sql.DB, source int, destination int) error {
	_, err := dbExec(db, "UPDATE images_aliases SET image_id=? WHERE image_id=?", destination, source)
	return err
}

// Insert an alias ento the database.
func dbImageAliasAdd(db *sql.DB, name string, imageID int, desc string) error {
	stmt := `INSERT INTO images_aliases (name, image_id, description) values (?, ?, ?)`
	_, err := dbExec(db, stmt, name, imageID, desc)
	return err
}

func dbImageAliasUpdate(db *sql.DB, id int, imageID int, desc string) error {
	stmt := `UPDATE images_aliases SET image_id=?, description=? WHERE id=?`
	_, err := dbExec(db, stmt, imageID, desc, id)
	return err
}

func dbImageLastAccessUpdate(db *sql.DB, fingerprint string, date time.Time) error {
	stmt := `UPDATE images SET last_use_date=? WHERE fingerprint=?`
	_, err := dbExec(db, stmt, date, fingerprint)
	return err
}

func dbImageLastAccessInit(db *sql.DB, fingerprint string) error {
	stmt := `UPDATE images SET cached=1, last_use_date=strftime("%s") WHERE fingerprint=?`
	_, err := dbExec(db, stmt, fingerprint)
	return err
}

func dbImageUpdate(db *sql.DB, id int, fname string, sz int64, public bool, autoUpdate bool, architecture string, createdAt time.Time, expiresAt time.Time, properties map[string]string) error {
	arch, err := osarch.ArchitectureId(architecture)
	if err != nil {
		arch = 0
	}

	tx, err := dbBegin(db)
	if err != nil {
		return err
	}

	publicInt := 0
	if public {
		publicInt = 1
	}

	autoUpdateInt := 0
	if autoUpdate {
		autoUpdateInt = 1
	}

	stmt, err := tx.Prepare(`UPDATE images SET filename=?, size=?, public=?, auto_update=?, architecture=?, creation_date=?, expiry_date=? WHERE id=?`)
	if err != nil {
		tx.Rollback()
		return err
	}
	defer stmt.Close()

	_, err = stmt.Exec(fname, sz, publicInt, autoUpdateInt, arch, createdAt, expiresAt, id)
	if err != nil {
		tx.Rollback()
		return err
	}

	_, err = tx.Exec(`DELETE FROM images_properties WHERE image_id=?`, id)

	stmt, err = tx.Prepare(`INSERT INTO images_properties (image_id, type, key, value) VALUES (?, ?, ?, ?)`)
	if err != nil {
		tx.Rollback()
		return err
	}

	for key, value := range properties {
		_, err = stmt.Exec(id, 0, key, value)
		if err != nil {
			tx.Rollback()
			return err
		}
	}

	if err := txCommit(tx); err != nil {
		return err
	}

	return nil
}

func dbImageInsert(db *sql.DB, fp string, fname string, sz int64, public bool, autoUpdate bool, architecture string, createdAt time.Time, expiresAt time.Time, properties map[string]string) error {
	arch, err := osarch.ArchitectureId(architecture)
	if err != nil {
		arch = 0
	}

	tx, err := dbBegin(db)
	if err != nil {
		return err
	}

	publicInt := 0
	if public {
		publicInt = 1
	}

	autoUpdateInt := 0
	if autoUpdate {
		autoUpdateInt = 1
	}

	stmt, err := tx.Prepare(`INSERT INTO images (fingerprint, filename, size, public, auto_update, architecture, creation_date, expiry_date, upload_date) VALUES (?, ?, ?, ?, ?, ?, ?, ?, strftime("%s"))`)
	if err != nil {
		tx.Rollback()
		return err
	}
	defer stmt.Close()

	result, err := stmt.Exec(fp, fname, sz, publicInt, autoUpdateInt, arch, createdAt, expiresAt)
	if err != nil {
		tx.Rollback()
		return err
	}

	if len(properties) > 0 {
		id64, err := result.LastInsertId()
		if err != nil {
			tx.Rollback()
			return err
		}
		id := int(id64)

		pstmt, err := tx.Prepare(`INSERT INTO images_properties (image_id, type, key, value) VALUES (?, 0, ?, ?)`)
		if err != nil {
			tx.Rollback()
			return err
		}
		defer pstmt.Close()

		for k, v := range properties {

			// we can assume, that there is just one
			// value per key
			_, err = pstmt.Exec(id, k, v)
			if err != nil {
				tx.Rollback()
				return err
			}
		}

	}

	if err := txCommit(tx); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package server

import (
	"bufio"
	"bytes"
	"io"
	"net/http"

	"github.com/influxdata/chronograf"
)

const (
	ErrNotFlusher = "Expected http.ResponseWriter to be an http.Flusher, but wasn't"
)

// URLPrefixer is a wrapper for an http.Handler that will prefix all occurrences of a relative URL with the configured Prefix
type URLPrefixer struct {
	Prefix string            // the prefix to be appended after any detected Attrs
	Next   http.Handler      // the http.Handler which will generate the content to be modified by this handler
	Attrs  [][]byte          // a list of attrs that should have their URLs prefixed. For example `src="` or `href="` would be valid
	Logger chronograf.Logger // The logger where prefixing errors will be dispatched to
}

type wrapResponseWriter struct {
	http.ResponseWriter
	Substitute *io.PipeWriter

	headerWritten bool
	dupHeader     *http.Header
}

func (wrw *wrapResponseWriter) Write(p []byte) (int, error) {
	return wrw.Substitute.Write(p)
}

func (wrw *wrapResponseWriter) WriteHeader(code int) {
	if !wrw.headerWritten {
		wrw.ResponseWriter.Header().Set("Content-Type", wrw.dupHeader.Get("Content-Type"))
		header := wrw.ResponseWriter.Header()
		// Filter out content length header to prevent stopping writing
		if wrw.dupHeader != nil {
			for k, v := range *wrw.dupHeader {
				if k == "Content-Length" {
					continue
				}
				header[k] = v
			}
		}

		wrw.headerWritten = true
	}
	wrw.ResponseWriter.WriteHeader(code)
}

// Header() copies the Header map from the underlying ResponseWriter to prevent
// modifications to it by callers
func (wrw *wrapResponseWriter) Header() http.Header {
	if wrw.dupHeader == nil {
		h := http.Header{}
		origHeader := wrw.ResponseWriter.Header()
		for k, v := range origHeader {
			h[k] = v
		}
		wrw.dupHeader = &h
	}
	return *wrw.dupHeader
}

// ChunkSize is the number of bytes per chunked transfer-encoding
const ChunkSize int = 512

// ServeHTTP implements an http.Handler that prefixes relative URLs from the
// Next handler with the configured prefix. It does this by examining the
// stream through the ResponseWriter, and appending the Prefix after any of the
// Attrs detected in the stream.
func (up *URLPrefixer) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
	// extract the flusher for flushing chunks
	flusher, ok := rw.(http.Flusher)

	if !ok {
		up.Logger.Info(ErrNotFlusher)
		up.Next.ServeHTTP(rw, r)
		return
	}

	// chunked transfer because we're modifying the response on the fly, so we
	// won't know the final content-length
	rw.Header().Set("Connection", "Keep-Alive")
	rw.Header().Set("Transfer-Encoding", "chunked")

	writtenCount := 0 // number of bytes written to rw
	nextRead, nextWrite := io.Pipe()
	go func() {
		defer nextWrite.Close()
		up.Next.ServeHTTP(&wrapResponseWriter{ResponseWriter: rw, Substitute: nextWrite}, r)
	}()

	// setup a buffer which is the max length of our target attrs
	b := make([]byte, up.maxlen(up.Attrs...))
	io.ReadFull(nextRead, b) // prime the buffer with the start of the input
	buf := bytes.NewBuffer(b)

	// Read next handler's response byte by byte
	src := bufio.NewScanner(nextRead)
	src.Split(bufio.ScanBytes)
	for {
		window := buf.Bytes()

		// advance a byte if window is not a src attr
		if matchlen, match := up.match(window, up.Attrs...); matchlen == 0 {
			if src.Scan() {
				// shift the next byte into buf
				rw.Write(buf.Next(1))
				writtenCount++
				buf.Write(src.Bytes())

				if writtenCount >= ChunkSize {
					flusher.Flush()
					writtenCount = 0
				}
			} else {
				if err := src.Err(); err != nil {
					up.Logger.
						WithField("component", "prefixer").
						Error("Error encountered while scanning: err:", err)
				}
				rw.Write(window)
				flusher.Flush()
				break
			}
			continue
		} else {
			buf.Next(matchlen) // advance to the relative URL
			for i := 0; i < matchlen; i++ {
				src.Scan()
				buf.Write(src.Bytes())
			}
			rw.Write(match)               // add the src attr to the output
			io.WriteString(rw, up.Prefix) // write the prefix
		}
	}
}

// match compares the subject against a list of targets. If there is a match
// between any of them a non-zero value is returned. The returned value is the
// length of the match. It is assumed that subject's length > length of all
// targets. The matching []byte is also returned as the second return parameter
func (up *URLPrefixer) match(subject []byte, targets ...[]byte) (int, []byte) {
	for _, target := range targets {
		if bytes.Equal(subject[:len(target)], target) {
			return len(target), target
		}
	}
	return 0, []byte{}
}

// maxlen returns the length of the largest []byte provided to it as an argument
func (up *URLPrefixer) maxlen(targets ...[]byte) int {
	max := 0
	for _, tgt := range targets {
		if tlen := len(tgt); tlen > max {
			max = tlen
		}
	}
	return max
}

// NewDefaultURLPrefixer returns a URLPrefixer that will prefix any src and
// href attributes found in HTML as well as any url() directives found in CSS
// with the provided prefix. Additionally, it will prefix any `data-basepath`
// attributes as well for informing front end logic about any prefixes. `next`
// is the next http.Handler that will have its output prefixed
func NewDefaultURLPrefixer(prefix string, next http.Handler, lg chronograf.Logger) *URLPrefixer {
	return &URLPrefixer{
		Prefix: prefix,
		Next:   next,
		Logger: lg,
		Attrs: [][]byte{
			[]byte(`src="`),
			[]byte(`href="`),
			[]byte(`url(`),
			[]byte(`data-basepath="`), // for forwarding basepath to frontend
		},
	}
}
<commit_msg>Update url prefixer to ignore svgs<commit_after>package server

import (
	"bufio"
	"bytes"
	"io"
	"net/http"
	"regexp"

	"github.com/influxdata/chronograf"
)

const (
	ErrNotFlusher = "Expected http.ResponseWriter to be an http.Flusher, but wasn't"
)

// URLPrefixer is a wrapper for an http.Handler that will prefix all occurrences of a relative URL with the configured Prefix
type URLPrefixer struct {
	Prefix string            // the prefix to be appended after any detected Attrs
	Next   http.Handler      // the http.Handler which will generate the content to be modified by this handler
	Attrs  [][]byte          // a list of attrs that should have their URLs prefixed. For example `src="` or `href="` would be valid
	Logger chronograf.Logger // The logger where prefixing errors will be dispatched to
}

type wrapResponseWriter struct {
	http.ResponseWriter
	Substitute *io.PipeWriter

	headerWritten bool
	dupHeader     *http.Header
}

func (wrw *wrapResponseWriter) Write(p []byte) (int, error) {
	return wrw.Substitute.Write(p)
}

func (wrw *wrapResponseWriter) WriteHeader(code int) {
	if !wrw.headerWritten {
		wrw.ResponseWriter.Header().Set("Content-Type", wrw.dupHeader.Get("Content-Type"))
		header := wrw.ResponseWriter.Header()
		// Filter out content length header to prevent stopping writing
		if wrw.dupHeader != nil {
			for k, v := range *wrw.dupHeader {
				if k == "Content-Length" {
					continue
				}
				header[k] = v
			}
		}

		wrw.headerWritten = true
	}
	wrw.ResponseWriter.WriteHeader(code)
}

// Header() copies the Header map from the underlying ResponseWriter to prevent
// modifications to it by callers
func (wrw *wrapResponseWriter) Header() http.Header {
	if wrw.dupHeader == nil {
		h := http.Header{}
		origHeader := wrw.ResponseWriter.Header()
		for k, v := range origHeader {
			h[k] = v
		}
		wrw.dupHeader = &h
	}
	return *wrw.dupHeader
}

// ChunkSize is the number of bytes per chunked transfer-encoding
const ChunkSize int = 512

// ServeHTTP implements an http.Handler that prefixes relative URLs from the
// Next handler with the configured prefix. It does this by examining the
// stream through the ResponseWriter, and appending the Prefix after any of the
// Attrs detected in the stream.
func (up *URLPrefixer) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
	// extract the flusher for flushing chunks
	flusher, ok := rw.(http.Flusher)

	if !ok {
		up.Logger.Info(ErrNotFlusher)
		up.Next.ServeHTTP(rw, r)
		return
	}

	isSVG, _ := regexp.Match(".svg$", []byte(r.URL.String()))
	if isSVG {
		up.Next.ServeHTTP(rw, r)
		return
	}

	// chunked transfer because we're modifying the response on the fly, so we
	// won't know the final content-length
	rw.Header().Set("Connection", "Keep-Alive")
	rw.Header().Set("Transfer-Encoding", "chunked")

	writtenCount := 0 // number of bytes written to rw
	nextRead, nextWrite := io.Pipe()
	go func() {
		defer nextWrite.Close()
		up.Next.ServeHTTP(&wrapResponseWriter{ResponseWriter: rw, Substitute: nextWrite}, r)
	}()

	// setup a buffer which is the max length of our target attrs
	b := make([]byte, up.maxlen(up.Attrs...))
	io.ReadFull(nextRead, b) // prime the buffer with the start of the input
	buf := bytes.NewBuffer(b)

	// Read next handler's response byte by byte
	src := bufio.NewScanner(nextRead)
	src.Split(bufio.ScanBytes)
	for {
		window := buf.Bytes()

		// advance a byte if window is not a src attr
		if matchlen, match := up.match(window, up.Attrs...); matchlen == 0 {
			if src.Scan() {
				// shift the next byte into buf
				rw.Write(buf.Next(1))
				writtenCount++
				buf.Write(src.Bytes())

				if writtenCount >= ChunkSize {
					flusher.Flush()
					writtenCount = 0
				}
			} else {
				if err := src.Err(); err != nil {
					up.Logger.
						WithField("component", "prefixer").
						Error("Error encountered while scanning: err:", err)
				}
				rw.Write(window)
				flusher.Flush()
				break
			}
			continue
		} else {
			buf.Next(matchlen) // advance to the relative URL
			for i := 0; i < matchlen; i++ {
				src.Scan()
				buf.Write(src.Bytes())
			}
			rw.Write(match)               // add the src attr to the output
			io.WriteString(rw, up.Prefix) // write the prefix
		}
	}
}

// match compares the subject against a list of targets. If there is a match
// between any of them a non-zero value is returned. The returned value is the
// length of the match. It is assumed that subject's length > length of all
// targets. The matching []byte is also returned as the second return parameter
func (up *URLPrefixer) match(subject []byte, targets ...[]byte) (int, []byte) {
	for _, target := range targets {
		if bytes.Equal(subject[:len(target)], target) {
			return len(target), target
		}
	}
	return 0, []byte{}
}

// maxlen returns the length of the largest []byte provided to it as an argument
func (up *URLPrefixer) maxlen(targets ...[]byte) int {
	max := 0
	for _, tgt := range targets {
		if tlen := len(tgt); tlen > max {
			max = tlen
		}
	}
	return max
}

// NewDefaultURLPrefixer returns a URLPrefixer that will prefix any src and
// href attributes found in HTML as well as any url() directives found in CSS
// with the provided prefix. Additionally, it will prefix any `data-basepath`
// attributes as well for informing front end logic about any prefixes. `next`
// is the next http.Handler that will have its output prefixed
func NewDefaultURLPrefixer(prefix string, next http.Handler, lg chronograf.Logger) *URLPrefixer {
	return &URLPrefixer{
		Prefix: prefix,
		Next:   next,
		Logger: lg,
		Attrs: [][]byte{
			[]byte(`src="`),
			[]byte(`href="`),
			[]byte(`url(`),
			[]byte(`data-basepath="`), // for forwarding basepath to frontend
		},
	}
}
<|endoftext|>


<commit_before>package session

import (
	"sync/atomic"

	"github.com/eclipse/paho.mqtt.golang/packets"

	log "github.com/Sirupsen/logrus"

	"github.com/iotalking/mqtt-broker/store"
)

//处理 CONNECT 消息
//连接消息
func (this *Session) onConnect(msg *packets.ConnectPacket) (err error) {
	//如果已经连接，则是违规，应该断开
	if this.IsConnected() {
		log.Error("connected client dup CONNECT")
		return packets.ConnErrors[packets.ErrProtocolViolation]
	}
	//判断保留标志位是否为0,如果不为0，返回协议错误
	if msg.ReservedBit&0x1 == 0x1 {
		log.Error("client ReservedBit != 0")
		return packets.ConnErrors[packets.ErrProtocolViolation]
	}
	//检查是否有相同clientId的连接，如果有，则要断开原有连接
	this.clientId = msg.ClientIdentifier
	this.mgr.DisconectSessionByClientId(this.clientId)

	conack := packets.NewControlPacket(packets.Connack).(*packets.ConnackPacket)
	conack.ReturnCode = packets.Accepted
	conack.SessionPresent = true

	log.Debug("sessionMgr.OnConnected")
	this.mgr.OnConnected(this)

	log.Debug("session.Send")

	err = this.Send(conack)
	atomic.StoreInt32(&this.connected, 1)

	return
}

//处理 CONNACK 消息
//连接应答消息
//此消息只有作为客户端时才会收到
func (this *Session) onConnack(msg *packets.ConnackPacket) error {
	atomic.StoreInt32(&this.connected, 1)
	//如果是客户端要启动pingTimer定时器，定时发送定时包
	//同时启动pingrespTimer定时器，检测服务器返回PINGRESP包有没有超时
	//如果客户pingrespTimer时间内没有收到PINGRESP，则断开连接

	return nil
}

func (this *Session) sendToSubcriber(msg *packets.PublishPacket) error {
	subMgr := this.mgr.GetSubscriptionMgr()

	sessionQosMap, err := subMgr.GetSessions(msg.TopicName)
	if err != nil {
		log.Error("subMgr.GetSessions error:", err)
		return err
	}
	if len(sessionQosMap) == 0 {
		log.Debug("sessionQosMap is empty for topic:", msg.TopicName)
	}
	for v, qos := range sessionQosMap {
		s := v.(*Session)
		if s.IsConnected() && !s.IsClosed() {
			nmsg := msg.Copy()
			//qos取订阅和原消息的最小值
			if qos < msg.Qos {
				nmsg.Qos = qos
			}
			err = s.Send(msg)
			if err != nil {
				log.Errorf("session[%s].Publish Send error:", s.clientId, err)
				break
			}
		} else {
			log.Debug("session is disconnected or not connect")
		}

	}
	return nil
}

//处理 PUBLISH 消息
//发布消息
func (this *Session) onPublish(msg *packets.PublishPacket) (err error) {
	log.Debug("session onpublish")
	//从mgr获取和tipic匹配的sessions
	//TODO:获取匹配的sessions

	switch msg.Qos {
	case 0:
		if msg.Retain {
			//清空离线消息
			err = this.mgr.storeMgr.RemoveMsgByTopic(msg.TopicName)
			if err != nil {
				log.Errorf("this.mgr.storeMgr.RemoveMsgByTopic error:", err)
				if err == store.ErrNoExsit {
					err = nil
				}
			}
			log.Debugf("retain msg has clear for client:%s", this.clientId)
			return
		}

		if this.isServer {
			this.sendToSubcriber(msg)
		} else {
			//回调接入到消息的函数
			this.callbackOnMessage(msg)
		}
	case 1:
		smsg := &store.Msg{
			ClientId: this.clientId,
			MsgId:    msg.MessageID,
			Topic:    msg.TopicName,
			Qos:      msg.Qos,
			Body:     msg.Payload,
		}
		if msg.Retain {
			err = this.mgr.GetStoreMgr().SaveRetainMsg(smsg)
			if err != nil {
				log.Error(err)
				err = packets.ConnErrors[packets.ErrRefusedServerUnavailable]
				return
			}
		}
		ack := packets.NewControlPacket(packets.Puback).(*packets.PubackPacket)
		ack.MessageID = msg.MessageID
		this.Send(ack)

		if this.isServer {
			this.sendToSubcriber(msg)
		} else {
			//回调接入到消息的函数
			this.callbackOnMessage(msg)
		}
	case 2:
		//保存消息，等待PUBREL后再发送
		smsg := &store.Msg{
			ClientId: this.clientId,
			MsgId:    msg.MessageID,
			Topic:    msg.TopicName,
			Qos:      msg.Qos,
			Body:     msg.Payload,
		}
		err = this.mgr.GetStoreMgr().SaveByClientIdMsgId(smsg)
		if err != nil {
			log.Error(err)
			err = packets.ConnErrors[packets.ErrRefusedServerUnavailable]
			return
		}
		if msg.Retain {
			err = this.mgr.GetStoreMgr().SaveRetainMsg(smsg)
			if err != nil {
				log.Error(err)
				err = packets.ConnErrors[packets.ErrRefusedServerUnavailable]
				return
			}
		}

		ack := packets.NewControlPacket(packets.Pubrec).(*packets.PubrecPacket)
		ack.MessageID = msg.MessageID
		this.Send(ack)

	}
	return err
}

//处理 PUBACK 消息
//向qos=1的主题消息发布者应答
func (this *Session) onPuback(msg *packets.PubackPacket) error {

	this.removeInflightMsg(msg.Details().MessageID, packets.Publish)

	return nil
}

//处理 PUBREC 消息
func (this *Session) onPubrec(msg *packets.PubrecPacket) error {

	this.removeInflightMsg(msg.Details().MessageID, packets.Publish)
	relmsg := packets.NewControlPacket(packets.Pubrel).(*packets.PubrelPacket)
	relmsg.MessageID = msg.MessageID
	this.Send(relmsg)
	return nil
}

//处理 PUBREL 消息
//向qos=2的主题消息发布者应答服务器第二步
//对PUBREC的响应
func (this *Session) onPubrel(msg *packets.PubrelPacket) error {
	storeMgr := this.mgr.GetStoreMgr()
	smsg, err := storeMgr.GetMsgByClientIdMsgId(this.clientId, msg.MessageID)
	if err != nil {
		log.Error(err)
		return nil
	}
	storeMgr.RemoveMsgById(smsg.Id)

	this.removeInflightMsg(msg.Details().MessageID, packets.Pubrec)
	compmsg := packets.NewControlPacket(packets.Pubcomp).(*packets.PubcompPacket)
	compmsg.MessageID = msg.MessageID
	this.Send(compmsg)

	pubMsg := packets.NewControlPacket(packets.Publish).(*packets.PublishPacket)
	pubMsg.Qos = smsg.Qos
	pubMsg.Payload = smsg.Body
	pubMsg.TopicName = smsg.Topic

	if this.isServer {
		this.sendToSubcriber(pubMsg)
	} else {
		//回调接入到消息的函数
		this.callbackOnMessage(pubMsg)
	}

	return nil
}

//处理 PUBCOMP 消息
//向qos=2的主题消息发布者应答服务器第三步(发布完成)
//对PUBREL的响应
//它是QoS 2等级协议交换的第四个也是最后一个报文
func (this *Session) onPubcomp(msg *packets.PubcompPacket) error {
	this.removeInflightMsg(msg.Details().MessageID, packets.Pubrel)
	return nil
}

//处理 SUBSCRIBE - 订阅主题消息
//客户端向服务端发送SUBSCRIBE报文用于创建一个或多个订阅。
//每个订阅注册客户端关心的一个或多个主题。为了将应用消息转发给与那些订阅匹配的主题，服务端发送PUBLISH报文给客户端。
//SUBSCRIBE报文也（为每个订阅）指定了最大的QoS等级，服务端根据这个发送应用消息给客户端。
func (this *Session) onSubscribe(msg *packets.SubscribePacket) error {
	log.Debug("session.onSubscribe")
	submgr := this.mgr.GetSubscriptionMgr()

	err := submgr.Add(msg.Topics, msg.Qoss, this)

	suback := packets.NewControlPacket(packets.Suback).(*packets.SubackPacket)
	suback.MessageID = msg.MessageID

	suback.ReturnCodes = msg.Qoss
	for i, _ := range msg.Qoss {
		if err != nil {
			suback.ReturnCodes[i] = 0x80
		}

	}

	return this.Send(suback)
}

//处理 SUBACK – 订阅确认
//服务端发送SUBACK报文给客户端，用于确认它已收到并且正在处理SUBSCRIBE报文。
//SUBACK报文包含一个返回码清单，它们指定了SUBSCRIBE请求的每个订阅被授予的最大QoS等级。
func (this *Session) onSuback(msg *packets.SubackPacket) error {
	//TODO

	return nil
}

//处理 UNSUBACK –取消订阅
//客户端发送UNSUBSCRIBE报文给服务端，用于取消订阅主题。
func (this *Session) onUnsubscribe(msg *packets.UnsubscribePacket) error {
	submgr := this.mgr.GetSubscriptionMgr()
	for _, s := range msg.Topics {
		submgr.Remove(s, this)
	}
	ack := &packets.UnsubackPacket{
		MessageID: msg.MessageID,
	}
	this.Send(ack)
	return nil
}

//处理 UNSUBSCRIBE –取消订阅确认
//服务端发送UNSUBACK报文给客户端。
func (this *Session) onUnsuback(msg *packets.UnsubackPacket) error {
	//TODO
	return nil
}

//处理 PINGREQ
func (this *Session) onPingreq(msg *packets.PingreqPacket) error {
	pingresp := packets.NewControlPacket(packets.Pingresp).(*packets.PingrespPacket)
	return this.Send(pingresp)
}

//处理 PINGRESP
//服务端发送PINGRESP报文响应客户端的PINGREQ报文。表示服务端还活着。
//保持连接（Keep Alive）处理中用到这个报文
func (this *Session) onPingresp(msg *packets.PingrespPacket) error {
	//TODO
	return nil
}

//处理 DISCONNECT
//从sessionMgr中删除session
//断开网络
func (this *Session) onDisconnect(msg *packets.DisconnectPacket) error {
	log.Debugf("session(%s) onDisconnect", this.clientId, this.sentMsgChan)
	this.mgrOnDisconnected()

	log.Debugf("session(%s) onDisconnect end", this.clientId)
	return nil
}

//处理Publish的参数
//qos=0
//Channel.Send后直接给调用者结果
//qos=1
//要等待服务端返回PUBACK
//如果超时要重发
//qos=2
//客户端时：
//1.发送PUBLISH,如果超时没有收到PUBREC，要重发PUBLISH
//2.发送PUBREL,如果超时没有收到PUBCOMP，要重发PUBREL
//3.收到PUBCOMP,给发布者返回结果
func (this *Session) onPublishData(msg *packets.PublishPacket) (err error) {
	log.Debug("onPublishData")
	msgId := this.lastMsgId + 1
	this.lastMsgId++
	msg.MessageID = uint16(msgId)
	switch msg.Qos {
	case 0:
		err = this.Send(msg)

	case 1, 2:
		//都有多步流程
		err = this.Send(msg)
	}
	return
}
<commit_msg>1.设置SessionPresent,TODO:未处理有保存session时的情况<commit_after>package session

import (
	"sync/atomic"

	"github.com/eclipse/paho.mqtt.golang/packets"

	log "github.com/Sirupsen/logrus"

	"github.com/iotalking/mqtt-broker/store"
)

//处理 CONNECT 消息
//连接消息
func (this *Session) onConnect(msg *packets.ConnectPacket) (err error) {
	//如果已经连接，则是违规，应该断开
	if this.IsConnected() {
		log.Error("connected client dup CONNECT")
		return packets.ConnErrors[packets.ErrProtocolViolation]
	}
	//判断保留标志位是否为0,如果不为0，返回协议错误
	if msg.ReservedBit&0x1 == 0x1 {
		log.Error("client ReservedBit != 0")
		return packets.ConnErrors[packets.ErrProtocolViolation]
	}
	//检查是否有相同clientId的连接，如果有，则要断开原有连接
	this.clientId = msg.ClientIdentifier
	this.mgr.DisconectSessionByClientId(this.clientId)

	conack := packets.NewControlPacket(packets.Connack).(*packets.ConnackPacket)
	conack.ReturnCode = packets.Accepted
	if msg.CleanSession {
		conack.SessionPresent = false
	} else {
		//如果服务器有保存session,则 SessionPresent设成true
		//如果服务器没有保存session,则SessionPresent设成false
		conack.SessionPresent = false
	}

	log.Debug("sessionMgr.OnConnected")
	this.mgr.OnConnected(this)

	log.Debug("session.Send")

	err = this.Send(conack)
	atomic.StoreInt32(&this.connected, 1)

	return
}

//处理 CONNACK 消息
//连接应答消息
//此消息只有作为客户端时才会收到
func (this *Session) onConnack(msg *packets.ConnackPacket) error {
	atomic.StoreInt32(&this.connected, 1)
	//如果是客户端要启动pingTimer定时器，定时发送定时包
	//同时启动pingrespTimer定时器，检测服务器返回PINGRESP包有没有超时
	//如果客户pingrespTimer时间内没有收到PINGRESP，则断开连接

	return nil
}

func (this *Session) sendToSubcriber(msg *packets.PublishPacket) error {
	subMgr := this.mgr.GetSubscriptionMgr()

	sessionQosMap, err := subMgr.GetSessions(msg.TopicName)
	if err != nil {
		log.Error("subMgr.GetSessions error:", err)
		return err
	}
	if len(sessionQosMap) == 0 {
		log.Debug("sessionQosMap is empty for topic:", msg.TopicName)
	}
	for v, qos := range sessionQosMap {
		s := v.(*Session)
		if s.IsConnected() && !s.IsClosed() {
			nmsg := msg.Copy()
			//qos取订阅和原消息的最小值
			if qos < msg.Qos {
				nmsg.Qos = qos
			}
			err = s.Send(msg)
			if err != nil {
				log.Errorf("session[%s].Publish Send error:", s.clientId, err)
				break
			}
		} else {
			log.Debug("session is disconnected or not connect")
		}

	}
	return nil
}

//处理 PUBLISH 消息
//发布消息
func (this *Session) onPublish(msg *packets.PublishPacket) (err error) {
	log.Debug("session onpublish")
	//从mgr获取和tipic匹配的sessions
	//TODO:获取匹配的sessions

	switch msg.Qos {
	case 0:
		if msg.Retain {
			//清空离线消息
			err = this.mgr.storeMgr.RemoveMsgByTopic(msg.TopicName)
			if err != nil {
				log.Errorf("this.mgr.storeMgr.RemoveMsgByTopic error:", err)
				if err == store.ErrNoExsit {
					err = nil
				}
			}
			log.Debugf("retain msg has clear for client:%s", this.clientId)
			return
		}

		if this.isServer {
			this.sendToSubcriber(msg)
		} else {
			//回调接入到消息的函数
			this.callbackOnMessage(msg)
		}
	case 1:
		smsg := &store.Msg{
			ClientId: this.clientId,
			MsgId:    msg.MessageID,
			Topic:    msg.TopicName,
			Qos:      msg.Qos,
			Body:     msg.Payload,
		}
		if msg.Retain {
			err = this.mgr.GetStoreMgr().SaveRetainMsg(smsg)
			if err != nil {
				log.Error(err)
				err = packets.ConnErrors[packets.ErrRefusedServerUnavailable]
				return
			}
		}
		ack := packets.NewControlPacket(packets.Puback).(*packets.PubackPacket)
		ack.MessageID = msg.MessageID
		this.Send(ack)

		if this.isServer {
			this.sendToSubcriber(msg)
		} else {
			//回调接入到消息的函数
			this.callbackOnMessage(msg)
		}
	case 2:
		//保存消息，等待PUBREL后再发送
		smsg := &store.Msg{
			ClientId: this.clientId,
			MsgId:    msg.MessageID,
			Topic:    msg.TopicName,
			Qos:      msg.Qos,
			Body:     msg.Payload,
		}
		err = this.mgr.GetStoreMgr().SaveByClientIdMsgId(smsg)
		if err != nil {
			log.Error(err)
			err = packets.ConnErrors[packets.ErrRefusedServerUnavailable]
			return
		}
		if msg.Retain {
			err = this.mgr.GetStoreMgr().SaveRetainMsg(smsg)
			if err != nil {
				log.Error(err)
				err = packets.ConnErrors[packets.ErrRefusedServerUnavailable]
				return
			}
		}

		ack := packets.NewControlPacket(packets.Pubrec).(*packets.PubrecPacket)
		ack.MessageID = msg.MessageID
		this.Send(ack)

	}
	return err
}

//处理 PUBACK 消息
//向qos=1的主题消息发布者应答
func (this *Session) onPuback(msg *packets.PubackPacket) error {

	this.removeInflightMsg(msg.Details().MessageID, packets.Publish)

	return nil
}

//处理 PUBREC 消息
func (this *Session) onPubrec(msg *packets.PubrecPacket) error {

	this.removeInflightMsg(msg.Details().MessageID, packets.Publish)
	relmsg := packets.NewControlPacket(packets.Pubrel).(*packets.PubrelPacket)
	relmsg.MessageID = msg.MessageID
	this.Send(relmsg)
	return nil
}

//处理 PUBREL 消息
//向qos=2的主题消息发布者应答服务器第二步
//对PUBREC的响应
func (this *Session) onPubrel(msg *packets.PubrelPacket) error {
	storeMgr := this.mgr.GetStoreMgr()
	smsg, err := storeMgr.GetMsgByClientIdMsgId(this.clientId, msg.MessageID)
	if err != nil {
		log.Error(err)
		return nil
	}
	storeMgr.RemoveMsgById(smsg.Id)

	this.removeInflightMsg(msg.Details().MessageID, packets.Pubrec)
	compmsg := packets.NewControlPacket(packets.Pubcomp).(*packets.PubcompPacket)
	compmsg.MessageID = msg.MessageID
	this.Send(compmsg)

	pubMsg := packets.NewControlPacket(packets.Publish).(*packets.PublishPacket)
	pubMsg.Qos = smsg.Qos
	pubMsg.Payload = smsg.Body
	pubMsg.TopicName = smsg.Topic

	if this.isServer {
		this.sendToSubcriber(pubMsg)
	} else {
		//回调接入到消息的函数
		this.callbackOnMessage(pubMsg)
	}

	return nil
}

//处理 PUBCOMP 消息
//向qos=2的主题消息发布者应答服务器第三步(发布完成)
//对PUBREL的响应
//它是QoS 2等级协议交换的第四个也是最后一个报文
func (this *Session) onPubcomp(msg *packets.PubcompPacket) error {
	this.removeInflightMsg(msg.Details().MessageID, packets.Pubrel)
	return nil
}

//处理 SUBSCRIBE - 订阅主题消息
//客户端向服务端发送SUBSCRIBE报文用于创建一个或多个订阅。
//每个订阅注册客户端关心的一个或多个主题。为了将应用消息转发给与那些订阅匹配的主题，服务端发送PUBLISH报文给客户端。
//SUBSCRIBE报文也（为每个订阅）指定了最大的QoS等级，服务端根据这个发送应用消息给客户端。
func (this *Session) onSubscribe(msg *packets.SubscribePacket) error {
	log.Debug("session.onSubscribe")
	submgr := this.mgr.GetSubscriptionMgr()

	err := submgr.Add(msg.Topics, msg.Qoss, this)

	suback := packets.NewControlPacket(packets.Suback).(*packets.SubackPacket)
	suback.MessageID = msg.MessageID

	suback.ReturnCodes = msg.Qoss
	for i, _ := range msg.Qoss {
		if err != nil {
			suback.ReturnCodes[i] = 0x80
		}

	}

	return this.Send(suback)
}

//处理 SUBACK – 订阅确认
//服务端发送SUBACK报文给客户端，用于确认它已收到并且正在处理SUBSCRIBE报文。
//SUBACK报文包含一个返回码清单，它们指定了SUBSCRIBE请求的每个订阅被授予的最大QoS等级。
func (this *Session) onSuback(msg *packets.SubackPacket) error {
	//TODO

	return nil
}

//处理 UNSUBACK –取消订阅
//客户端发送UNSUBSCRIBE报文给服务端，用于取消订阅主题。
func (this *Session) onUnsubscribe(msg *packets.UnsubscribePacket) error {
	submgr := this.mgr.GetSubscriptionMgr()
	for _, s := range msg.Topics {
		submgr.Remove(s, this)
	}
	ack := &packets.UnsubackPacket{
		MessageID: msg.MessageID,
	}
	this.Send(ack)
	return nil
}

//处理 UNSUBSCRIBE –取消订阅确认
//服务端发送UNSUBACK报文给客户端。
func (this *Session) onUnsuback(msg *packets.UnsubackPacket) error {
	//TODO
	return nil
}

//处理 PINGREQ
func (this *Session) onPingreq(msg *packets.PingreqPacket) error {
	pingresp := packets.NewControlPacket(packets.Pingresp).(*packets.PingrespPacket)
	return this.Send(pingresp)
}

//处理 PINGRESP
//服务端发送PINGRESP报文响应客户端的PINGREQ报文。表示服务端还活着。
//保持连接（Keep Alive）处理中用到这个报文
func (this *Session) onPingresp(msg *packets.PingrespPacket) error {
	//TODO
	return nil
}

//处理 DISCONNECT
//从sessionMgr中删除session
//断开网络
func (this *Session) onDisconnect(msg *packets.DisconnectPacket) error {
	log.Debugf("session(%s) onDisconnect", this.clientId, this.sentMsgChan)
	this.mgrOnDisconnected()

	log.Debugf("session(%s) onDisconnect end", this.clientId)
	return nil
}

//处理Publish的参数
//qos=0
//Channel.Send后直接给调用者结果
//qos=1
//要等待服务端返回PUBACK
//如果超时要重发
//qos=2
//客户端时：
//1.发送PUBLISH,如果超时没有收到PUBREC，要重发PUBLISH
//2.发送PUBREL,如果超时没有收到PUBCOMP，要重发PUBREL
//3.收到PUBCOMP,给发布者返回结果
func (this *Session) onPublishData(msg *packets.PublishPacket) (err error) {
	log.Debug("onPublishData")
	msgId := this.lastMsgId + 1
	this.lastMsgId++
	msg.MessageID = uint16(msgId)
	switch msg.Qos {
	case 0:
		err = this.Send(msg)

	case 1, 2:
		//都有多步流程
		err = this.Send(msg)
	}
	return
}
<|endoftext|>


<commit_before>package glutton

import (
	"bufio"
	"log"
	"net"
	"regexp"
	"strings"
)

// ConnID is used to relate logs in a connection
var connID int64

// Based on https://github.com/CymmetriaResearch/MTPot/blob/master/mirai_conf.json
var miraiCom = map[string]string{
	"ECCHI":                                           "ECCHI: applet not found",
	"ps":                                              "1 pts/21   00:00:00 init",
	"cat /proc/mounts":                                "tmpfs /run tmpfs rw,nosuid,noexec,relatime,size=1635616k,mode=755 0 0",
	"echo -e \\x6b\\x61\\x6d\\x69/dev > /dev/.nippon": "",
	"cat /dev/.nippon":                                "kami/dev",
	"rm /dev/.nippon":                                 "",
	"echo -e \\x6b\\x61\\x6d\\x69/run > /run/.nippon": "",
	"cat /run/.nippon":                                "kami/run",
	"rm /run/.nippon":                                 "",
}

func writeMsg(conn net.Conn, msg string) error {
	_, err := conn.Write([]byte(msg))
	log.Printf("[%v] [TCP] [TELNET -> %v] Payload: %v", connID, conn.RemoteAddr(), msg)
	return err
}

func readMsg(conn net.Conn) (string, error) {
	message, err := bufio.NewReader(conn).ReadString('\n')
	if err != nil {
		return "", err
	}
	log.Printf("[%v] [TCP] [%v -> TELNET] Payload: %v", connID, conn.RemoteAddr(), message)
	return message, err
}

func handleTelnet(id int64, conn net.Conn) error {
	defer conn.Close()
	connID = id

	writeMsg(conn, "Username: ")
	_, err := readMsg(conn)
	if err != nil {
		return err
	}
	writeMsg(conn, "Password: ")
	_, err = readMsg(conn)
	if err != nil {
		return err
	}

	writeMsg(conn, "welcome\n> ")
	for {
		msg, err := readMsg(conn)
		if err != nil {
			return err
		}
		for _, cmd := range strings.Split(msg, ";") {
			if resp := miraiCom[strings.TrimSpace(cmd)]; resp != "" {
				writeMsg(conn, resp+"\n")
			} else {
				re := regexp.MustCompile(`\/bin\/busybox (?P<applet>[A-Z]+)`)
				match := re.FindStringSubmatch(cmd)
				if len(match) > 1 {
					writeMsg(conn, match[1]+": applet not found\n")
				}
			}
		}
		writeMsg(conn, "> ")
	}
}
<commit_msg>slight mirai adjustments<commit_after>package glutton

import (
	"bufio"
	"log"
	"net"
	"regexp"
	"strings"
)

// ConnID is used to relate logs in a connection
var connID int64

// Based on https://github.com/CymmetriaResearch/MTPot/blob/master/mirai_conf.json
var miraiCom = map[string]string{
	"ECCHI":                                           "ECCHI: applet not found",
	"ps":                                              "1 pts/21   00:00:00 init",
	"cat /proc/mounts":                                "tmpfs /run tmpfs rw,nosuid,noexec,relatime,size=3231524k,mode=755 0 0",
	"echo -e \\x6b\\x61\\x6d\\x69/dev > /dev/.nippon": "",
	"cat /dev/.nippon":                                "kami/dev",
	"rm /dev/.nippon":                                 "",
	"echo -e \\x6b\\x61\\x6d\\x69/run > /run/.nippon": "",
	"cat /run/.nippon":                                "kami/run",
	"rm /run/.nippon":                                 "",
}

func writeMsg(conn net.Conn, msg string) error {
	_, err := conn.Write([]byte(msg))
	log.Printf("[%v] [TCP] [TELNET -> %v] Payload: %v", connID, conn.RemoteAddr(), msg)
	return err
}

func readMsg(conn net.Conn) (string, error) {
	message, err := bufio.NewReader(conn).ReadString('\n')
	if err != nil {
		return "", err
	}
	log.Printf("[%v] [TCP] [%v -> TELNET] Payload: %v", connID, conn.RemoteAddr(), message)
	return message, err
}

func handleTelnet(id int64, conn net.Conn) error {
	defer conn.Close()
	connID = id

	writeMsg(conn, "Username: ")
	_, err := readMsg(conn)
	if err != nil {
		return err
	}
	writeMsg(conn, "Password: ")
	_, err = readMsg(conn)
	if err != nil {
		return err
	}

	writeMsg(conn, "welcome\n> ")
	for {
		msg, err := readMsg(conn)
		if err != nil {
			return err
		}
		for _, cmd := range strings.Split(msg, ";") {
			if resp := miraiCom[strings.TrimSpace(cmd)]; resp != "" {
				writeMsg(conn, resp+"\n")
			} else {
				re := regexp.MustCompile(`\/bin\/busybox (?P<applet>[A-Z]+)`)
				match := re.FindStringSubmatch(cmd)
				if len(match) > 1 {
					writeMsg(conn, match[1]+": applet not found\n")
				}
			}
		}
		writeMsg(conn, "> ")
	}
}
<|endoftext|>


<commit_before>package mapper

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"reflect"
	"strings"
	"unsafe"

	"github.com/jonbodner/proteus/logger"
)

func ptrConverter(c context.Context, isPtr bool, sType reflect.Type, out reflect.Value, err error) (interface{}, error) {
	if err != nil {
		return nil, err
	}
	if isPtr {
		out2 := reflect.New(sType)
		k := out.Type().Kind()
		logger.Log(c, logger.DEBUG, fmt.Sprintln("kind of out", k))
		if (k == reflect.Interface || k == reflect.Ptr) && out.IsNil() {
			out2 = reflect.NewAt(sType, unsafe.Pointer(nil))
		} else {
			out2.Elem().Set(out)
		}
		return out2.Interface(), nil
	}
	k := out.Type().Kind()
	if (k == reflect.Ptr || k == reflect.Interface) && out.IsNil() {
		return nil, fmt.Errorf("Attempting to return nil for non-pointer type %v", sType)
	}
	return out.Interface(), nil
}

func MakeBuilder(c context.Context, sType reflect.Type) (Builder, error) {
	if sType == nil {
		return nil, errors.New("sType cannot be nil")
	}

	isPtr := false
	if sType.Kind() == reflect.Ptr {
		sType = sType.Elem()
		isPtr = true
	}

	//only going to handle scalar types here
	if sType.Kind() == reflect.Slice {
		sType = sType.Elem()
	}

	if sType.Kind() == reflect.Map {
		if sType.Key().Kind() != reflect.String {
			return nil, errors.New("Only maps with string keys are supported")
		}
		return func(cols []string, vals []interface{}) (interface{}, error) {
			out, err := buildMap(c, sType, cols, vals)
			return ptrConverter(c, isPtr, sType, out, err)
		}, nil
	} else if sType.Kind() == reflect.Struct {
		//build map of col names to field names (makes this 2N instead of N^2)
		colFieldMap := map[string]fieldInfo{}
		buildColFieldMap(sType, fieldInfo{}, colFieldMap)
		return func(cols []string, vals []interface{}) (interface{}, error) {
			out, err := buildStruct(c, sType, cols, vals, colFieldMap)
			return ptrConverter(c, isPtr, sType, out, err)
		}, nil

	} else {
		// assume primitive
		return func(cols []string, vals []interface{}) (interface{}, error) {
			out, err := buildPrimitive(c, sType, cols, vals)
			return ptrConverter(c, isPtr, sType, out, err)
		}, nil
	}
}

func buildColFieldMap(sType reflect.Type, parentFieldInfo fieldInfo, colFieldMap map[string]fieldInfo) {
	for i := 0; i < sType.NumField(); i++ {
		sf := sType.Field(i)
		if tagVal := sf.Tag.Get("prof"); tagVal != "" {
			childFieldInfo := fieldInfo{
				name:      append([]string(nil), parentFieldInfo.name...),
				fieldType: append([]reflect.Type(nil), parentFieldInfo.fieldType...),
				pos:       append([]int(nil), parentFieldInfo.pos...),
			}
			childFieldInfo.name = append(childFieldInfo.name, sf.Name)
			childFieldInfo.fieldType = append(childFieldInfo.fieldType, sf.Type)
			childFieldInfo.pos = append(childFieldInfo.pos, i)
			//if this is a struct, recurse
			// we are going to use an prof value as a prefix
			// then we go through each of the fields in the struct
			// if any of them aren't structs and have prof tags, we
			// prepend the parent struct and store off a fieldi
			// only if this doesn't implement a scanner. If it does, then go with the scanner
			// another special case: time.Time isn't recursed into
			if sf.Type.Kind() == reflect.Struct && !reflect.PtrTo(sf.Type).Implements(scannerType) && sf.Type.Name() != "Time" && sf.Type.PkgPath() != "time" {
				buildColFieldMap(sf.Type, childFieldInfo, colFieldMap)
			} else {
				colFieldMap[strings.SplitN(tagVal, ",", 2)[0]] = childFieldInfo
			}
			continue
		}

		//last chance, if it's an anonymous struct, you can recurse into it without having a prof tag
		if sf.Type.Kind() == reflect.Struct && sf.Anonymous {
			childFieldInfo := fieldInfo{
				name:      append([]string(nil), parentFieldInfo.name...),
				fieldType: append([]reflect.Type(nil), parentFieldInfo.fieldType...),
				pos:       append([]int(nil), parentFieldInfo.pos...),
			}
			childFieldInfo.name = append(childFieldInfo.name, "")
			childFieldInfo.fieldType = append(childFieldInfo.fieldType, sf.Type)
			childFieldInfo.pos = append(childFieldInfo.pos, i)
			buildColFieldMap(sf.Type, childFieldInfo, colFieldMap)
		}
	}
}

type Builder func(cols []string, vals []interface{}) (interface{}, error)

type fieldInfo struct {
	name      []string
	fieldType []reflect.Type
	pos       []int
}

func buildMap(c context.Context, sType reflect.Type, cols []string, vals []interface{}) (reflect.Value, error) {
	out := reflect.MakeMap(sType)
	for k, v := range cols {
		curVal := vals[k]
		rv := reflect.ValueOf(curVal)
		if rv.Elem().IsNil() {
			continue
		}
		if rv.Elem().Elem().Type().ConvertibleTo(sType.Elem()) {
			out.SetMapIndex(reflect.ValueOf(v), rv.Elem().Elem().Convert(sType.Elem()))
		} else {
			return out, fmt.Errorf("Unable to assign value %v of type %v to map value of type %v with key %s", rv.Elem().Elem(), rv.Elem().Elem().Type(), sType.Elem(), v)
		}
	}
	return out, nil
}

var (
	scannerType = reflect.TypeOf((*sql.Scanner)(nil)).Elem()
)

func buildStruct(c context.Context, sType reflect.Type, cols []string, vals []interface{}, colFieldMap map[string]fieldInfo) (reflect.Value, error) {
	logger.Log(c, logger.DEBUG, fmt.Sprintf("sType: %s cols: %v vals: %+v colFieldMap: %+v", sType, cols, vals, colFieldMap))
	out := reflect.New(sType).Elem()
	for k, v := range cols {
		if sf, ok := colFieldMap[v]; ok {
			curVal := vals[k]
			rv := reflect.ValueOf(curVal)
			err := buildStructInner(c, sType, out, sf, curVal, rv, 0)
			if err != nil {
				return out, err
			}
		}
	}
	return out, nil
}

func buildStructInner(c context.Context, sType reflect.Type, out reflect.Value, sf fieldInfo, curVal interface{}, rv reflect.Value, depth int) error {
	field := out.Field(sf.pos[depth])
	curFieldType := sf.fieldType[depth]
	fmt.Printf("sf==%+v, depth==%d, field == %+v, curFieldType == %+v, curVal == %+v\n", sf, depth, field, curFieldType, curVal)
	if curFieldType.Kind() == reflect.Ptr {
		logger.Log(c, logger.DEBUG, fmt.Sprintln("isPtr", sf, rv, rv.Type(), rv.Elem(), curVal, sType))
		if rv.Elem().IsNil() {
			logger.Log(c, logger.DEBUG, fmt.Sprintln("nil", sf, rv, curVal, sType))
			return nil
		}
		logger.Log(c, logger.DEBUG, fmt.Sprintln("isPtr Not Nil", rv.Elem().Type()))
		if curFieldType.Implements(scannerType) {
			toScan := (reflect.New(curFieldType).Elem().Interface()).(sql.Scanner)
			err := toScan.Scan(rv.Elem().Elem().Interface())
			if err != nil {
				return err
			}
			field.Set(reflect.ValueOf(toScan))
		} else if rv.Elem().Elem().Type().ConvertibleTo(curFieldType.Elem()) {
			field.Set(reflect.New(curFieldType.Elem()))
			field.Elem().Set(rv.Elem().Elem().Convert(curFieldType.Elem()))
		} else {
			logger.Log(c, logger.ERROR, fmt.Sprintln("can't find the field"))
			return fmt.Errorf("Unable to assign pointer to value %v of type %v to struct field %s of type %v", rv.Elem().Elem(), rv.Elem().Elem().Type(), sf.name[depth], curFieldType)
		}
	} else {
		if reflect.PtrTo(curFieldType).Implements(scannerType) {
			toScan := (reflect.New(curFieldType).Interface()).(sql.Scanner)
			err := toScan.Scan(rv.Elem().Interface())
			if err != nil {
				return err
			}
			field.Set(reflect.ValueOf(toScan).Elem())
		} else if field.Kind() == reflect.Struct && field.Type().Name() != "Time" && field.Type().PkgPath() != "time" {
			err := buildStructInner(c, field.Type(), field, sf, curVal, rv, depth+1)
			if err != nil {
				return err
			}
		} else if rv.Elem().IsNil() {
			logger.Log(c, logger.ERROR, fmt.Sprintln("Attempting to assign a nil to a non-pointer field"))
			return fmt.Errorf("Unable to assign nil value to non-pointer struct field %s of type %v", sf.name[depth], curFieldType)
		} else if rv.Elem().Elem().Type().ConvertibleTo(curFieldType) {
			field.Set(rv.Elem().Elem().Convert(curFieldType))
		} else {
			logger.Log(c, logger.ERROR, fmt.Sprintln("can't find the field"))
			return fmt.Errorf("Unable to assign value %v of type %v to struct field %s of type %v", rv.Elem().Elem(), rv.Elem().Elem().Type(), sf.name[depth], curFieldType)
		}
	}
	return nil
}

func buildPrimitive(c context.Context, sType reflect.Type, cols []string, vals []interface{}) (reflect.Value, error) {
	out := reflect.New(sType).Elem()
	//vals[0] is of type *interface{}, because everything in vals is of type *interface{}
	rv := reflect.ValueOf(vals[0])
	if rv.Elem().IsNil() {
		return rv.Elem(), nil
	}
	if rv.Elem().Elem().Type().ConvertibleTo(sType) {
		out.Set(rv.Elem().Elem().Convert(sType))
	} else {
		return out, fmt.Errorf("Unable to assign value %v of type %v to return type of type %v", rv.Elem().Elem(), rv.Elem().Elem().Type(), sType)
	}
	return out, nil
}
<commit_msg>remove debug code<commit_after>package mapper

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"reflect"
	"strings"
	"unsafe"

	"github.com/jonbodner/proteus/logger"
)

func ptrConverter(c context.Context, isPtr bool, sType reflect.Type, out reflect.Value, err error) (interface{}, error) {
	if err != nil {
		return nil, err
	}
	if isPtr {
		out2 := reflect.New(sType)
		k := out.Type().Kind()
		logger.Log(c, logger.DEBUG, fmt.Sprintln("kind of out", k))
		if (k == reflect.Interface || k == reflect.Ptr) && out.IsNil() {
			out2 = reflect.NewAt(sType, unsafe.Pointer(nil))
		} else {
			out2.Elem().Set(out)
		}
		return out2.Interface(), nil
	}
	k := out.Type().Kind()
	if (k == reflect.Ptr || k == reflect.Interface) && out.IsNil() {
		return nil, fmt.Errorf("Attempting to return nil for non-pointer type %v", sType)
	}
	return out.Interface(), nil
}

func MakeBuilder(c context.Context, sType reflect.Type) (Builder, error) {
	if sType == nil {
		return nil, errors.New("sType cannot be nil")
	}

	isPtr := false
	if sType.Kind() == reflect.Ptr {
		sType = sType.Elem()
		isPtr = true
	}

	//only going to handle scalar types here
	if sType.Kind() == reflect.Slice {
		sType = sType.Elem()
	}

	if sType.Kind() == reflect.Map {
		if sType.Key().Kind() != reflect.String {
			return nil, errors.New("Only maps with string keys are supported")
		}
		return func(cols []string, vals []interface{}) (interface{}, error) {
			out, err := buildMap(c, sType, cols, vals)
			return ptrConverter(c, isPtr, sType, out, err)
		}, nil
	} else if sType.Kind() == reflect.Struct {
		//build map of col names to field names (makes this 2N instead of N^2)
		colFieldMap := map[string]fieldInfo{}
		buildColFieldMap(sType, fieldInfo{}, colFieldMap)
		return func(cols []string, vals []interface{}) (interface{}, error) {
			out, err := buildStruct(c, sType, cols, vals, colFieldMap)
			return ptrConverter(c, isPtr, sType, out, err)
		}, nil

	} else {
		// assume primitive
		return func(cols []string, vals []interface{}) (interface{}, error) {
			out, err := buildPrimitive(c, sType, cols, vals)
			return ptrConverter(c, isPtr, sType, out, err)
		}, nil
	}
}

func buildColFieldMap(sType reflect.Type, parentFieldInfo fieldInfo, colFieldMap map[string]fieldInfo) {
	for i := 0; i < sType.NumField(); i++ {
		sf := sType.Field(i)
		if tagVal := sf.Tag.Get("prof"); tagVal != "" {
			childFieldInfo := fieldInfo{
				name:      append([]string(nil), parentFieldInfo.name...),
				fieldType: append([]reflect.Type(nil), parentFieldInfo.fieldType...),
				pos:       append([]int(nil), parentFieldInfo.pos...),
			}
			childFieldInfo.name = append(childFieldInfo.name, sf.Name)
			childFieldInfo.fieldType = append(childFieldInfo.fieldType, sf.Type)
			childFieldInfo.pos = append(childFieldInfo.pos, i)
			//if this is a struct, recurse
			// we are going to use an prof value as a prefix
			// then we go through each of the fields in the struct
			// if any of them aren't structs and have prof tags, we
			// prepend the parent struct and store off a fieldi
			// only if this doesn't implement a scanner. If it does, then go with the scanner
			// another special case: time.Time isn't recursed into
			if sf.Type.Kind() == reflect.Struct && !reflect.PtrTo(sf.Type).Implements(scannerType) && sf.Type.Name() != "Time" && sf.Type.PkgPath() != "time" {
				buildColFieldMap(sf.Type, childFieldInfo, colFieldMap)
			} else {
				colFieldMap[strings.SplitN(tagVal, ",", 2)[0]] = childFieldInfo
			}
			continue
		}

		//last chance, if it's an anonymous struct, you can recurse into it without having a prof tag
		if sf.Type.Kind() == reflect.Struct && sf.Anonymous {
			childFieldInfo := fieldInfo{
				name:      append([]string(nil), parentFieldInfo.name...),
				fieldType: append([]reflect.Type(nil), parentFieldInfo.fieldType...),
				pos:       append([]int(nil), parentFieldInfo.pos...),
			}
			childFieldInfo.name = append(childFieldInfo.name, "")
			childFieldInfo.fieldType = append(childFieldInfo.fieldType, sf.Type)
			childFieldInfo.pos = append(childFieldInfo.pos, i)
			buildColFieldMap(sf.Type, childFieldInfo, colFieldMap)
		}
	}
}

type Builder func(cols []string, vals []interface{}) (interface{}, error)

type fieldInfo struct {
	name      []string
	fieldType []reflect.Type
	pos       []int
}

func buildMap(c context.Context, sType reflect.Type, cols []string, vals []interface{}) (reflect.Value, error) {
	out := reflect.MakeMap(sType)
	for k, v := range cols {
		curVal := vals[k]
		rv := reflect.ValueOf(curVal)
		if rv.Elem().IsNil() {
			continue
		}
		if rv.Elem().Elem().Type().ConvertibleTo(sType.Elem()) {
			out.SetMapIndex(reflect.ValueOf(v), rv.Elem().Elem().Convert(sType.Elem()))
		} else {
			return out, fmt.Errorf("Unable to assign value %v of type %v to map value of type %v with key %s", rv.Elem().Elem(), rv.Elem().Elem().Type(), sType.Elem(), v)
		}
	}
	return out, nil
}

var (
	scannerType = reflect.TypeOf((*sql.Scanner)(nil)).Elem()
)

func buildStruct(c context.Context, sType reflect.Type, cols []string, vals []interface{}, colFieldMap map[string]fieldInfo) (reflect.Value, error) {
	logger.Log(c, logger.DEBUG, fmt.Sprintf("sType: %s cols: %v vals: %+v colFieldMap: %+v", sType, cols, vals, colFieldMap))
	out := reflect.New(sType).Elem()
	for k, v := range cols {
		if sf, ok := colFieldMap[v]; ok {
			curVal := vals[k]
			rv := reflect.ValueOf(curVal)
			err := buildStructInner(c, sType, out, sf, curVal, rv, 0)
			if err != nil {
				return out, err
			}
		}
	}
	return out, nil
}

func buildStructInner(c context.Context, sType reflect.Type, out reflect.Value, sf fieldInfo, curVal interface{}, rv reflect.Value, depth int) error {
	field := out.Field(sf.pos[depth])
	curFieldType := sf.fieldType[depth]
	if curFieldType.Kind() == reflect.Ptr {
		logger.Log(c, logger.DEBUG, fmt.Sprintln("isPtr", sf, rv, rv.Type(), rv.Elem(), curVal, sType))
		if rv.Elem().IsNil() {
			logger.Log(c, logger.DEBUG, fmt.Sprintln("nil", sf, rv, curVal, sType))
			return nil
		}
		logger.Log(c, logger.DEBUG, fmt.Sprintln("isPtr Not Nil", rv.Elem().Type()))
		if curFieldType.Implements(scannerType) {
			toScan := (reflect.New(curFieldType).Elem().Interface()).(sql.Scanner)
			err := toScan.Scan(rv.Elem().Elem().Interface())
			if err != nil {
				return err
			}
			field.Set(reflect.ValueOf(toScan))
		} else if rv.Elem().Elem().Type().ConvertibleTo(curFieldType.Elem()) {
			field.Set(reflect.New(curFieldType.Elem()))
			field.Elem().Set(rv.Elem().Elem().Convert(curFieldType.Elem()))
		} else {
			logger.Log(c, logger.ERROR, fmt.Sprintln("can't find the field"))
			return fmt.Errorf("Unable to assign pointer to value %v of type %v to struct field %s of type %v", rv.Elem().Elem(), rv.Elem().Elem().Type(), sf.name[depth], curFieldType)
		}
	} else {
		if reflect.PtrTo(curFieldType).Implements(scannerType) {
			toScan := (reflect.New(curFieldType).Interface()).(sql.Scanner)
			err := toScan.Scan(rv.Elem().Interface())
			if err != nil {
				return err
			}
			field.Set(reflect.ValueOf(toScan).Elem())
		} else if field.Kind() == reflect.Struct && field.Type().Name() != "Time" && field.Type().PkgPath() != "time" {
			err := buildStructInner(c, field.Type(), field, sf, curVal, rv, depth+1)
			if err != nil {
				return err
			}
		} else if rv.Elem().IsNil() {
			logger.Log(c, logger.ERROR, fmt.Sprintln("Attempting to assign a nil to a non-pointer field"))
			return fmt.Errorf("Unable to assign nil value to non-pointer struct field %s of type %v", sf.name[depth], curFieldType)
		} else if rv.Elem().Elem().Type().ConvertibleTo(curFieldType) {
			field.Set(rv.Elem().Elem().Convert(curFieldType))
		} else {
			logger.Log(c, logger.ERROR, fmt.Sprintln("can't find the field"))
			return fmt.Errorf("Unable to assign value %v of type %v to struct field %s of type %v", rv.Elem().Elem(), rv.Elem().Elem().Type(), sf.name[depth], curFieldType)
		}
	}
	return nil
}

func buildPrimitive(c context.Context, sType reflect.Type, cols []string, vals []interface{}) (reflect.Value, error) {
	out := reflect.New(sType).Elem()
	//vals[0] is of type *interface{}, because everything in vals is of type *interface{}
	rv := reflect.ValueOf(vals[0])
	if rv.Elem().IsNil() {
		return rv.Elem(), nil
	}
	if rv.Elem().Elem().Type().ConvertibleTo(sType) {
		out.Set(rv.Elem().Elem().Convert(sType))
	} else {
		return out, fmt.Errorf("Unable to assign value %v of type %v to return type of type %v", rv.Elem().Elem(), rv.Elem().Elem().Type(), sType)
	}
	return out, nil
}
<|endoftext|>


<commit_before>/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package etcd

import (
	"context"
	"sync"

	"github.com/apache/servicecomb-service-center/pkg/log"
	"github.com/apache/servicecomb-service-center/pkg/util"
	pb "github.com/apache/servicecomb-service-center/syncer/proto"
	"github.com/apache/servicecomb-service-center/syncer/servicecenter"
	"github.com/coreos/etcd/clientv3"
	"github.com/gogo/protobuf/proto"
)

var (
	// mappingsKey the key of instances mapping in etcd
	mappingsKey = "/syncer/v1/mappings"
)

type storage struct {
	client *clientv3.Client
	data   *pb.SyncData
	lock   sync.RWMutex
}

func NewStorage(client *clientv3.Client) servicecenter.Storage {
	storage := &storage{
		client: client,
		data:   &pb.SyncData{},
	}
	return storage
}

// getPrefixKey Get data from etcd based on the prefix key
func (s *storage) getPrefixKey(prefix string, handler func(key, val []byte) (next bool)) {
	resp, err := s.client.Get(context.Background(), prefix, clientv3.WithPrefix())
	if err != nil {
		log.Errorf(err, "Get mapping from etcd failed: %s", err)
		return
	}

	for _, kv := range resp.Kvs {
		if !handler(kv.Key, kv.Value) {
			break
		}
	}
}

// UpdateData Update data to storage
func (s *storage) UpdateData(data *pb.SyncData) {
	s.lock.Lock()
	s.data = data
	s.lock.Unlock()
}

// GetData Get data from storage
func (s *storage) GetData() (data *pb.SyncData) {
	s.lock.RLock()
	data = s.data
	s.lock.RUnlock()
	return
}

// cleanExpired clean the expired instances in the maps
func (s *storage) cleanExpired(sources pb.SyncMapping, actives pb.SyncMapping) {
next:
	for _, entry := range sources {
		for _, act := range actives {
			if act.CurInstanceID == entry.CurInstanceID {
				continue next
			}
		}
		key := mappingsKey + "/" + entry.ClusterName + "/" + entry.OrgInstanceID
		if _, err := s.client.Delete(context.Background(), key); err != nil {
			log.Errorf(err, "Delete instance clusterName=%s instanceID=%s failed", entry.ClusterName, entry.OrgInstanceID)
		}
	}
}

// UpdateMapByCluster update map to storage by clusterName of other cluster
func (s *storage) UpdateMapByCluster(clusterName string, mapping pb.SyncMapping) {
	newMaps := make(pb.SyncMapping, 0, len(mapping))
	for _, val := range mapping {
		key := mappingsKey + "/" + clusterName + "/" + val.OrgInstanceID
		data, err := proto.Marshal(val)
		if err != nil {
			log.Errorf(err, "Proto marshal failed: %s", err)
			continue
		}
		_, err = s.client.Put(context.Background(), key, util.BytesToStringWithNoCopy(data))
		if err != nil {
			log.Errorf(err, "Save mapping to etcd failed: %s", err)
		}
		newMaps = append(newMaps, val)
	}
	s.cleanExpired(s.GetMapByCluster(clusterName), newMaps)
}

// GetMapByCluster get map by clusterName of other cluster
func (s *storage) GetMapByCluster(clusterName string) (mapping pb.SyncMapping) {
	maps := make(pb.SyncMapping, 0, 10)
	s.getPrefixKey(mappingsKey+"/"+clusterName, func(key, val []byte) (next bool) {
		next = true
		item := &pb.MappingEntry{}
		if err := proto.Unmarshal(val, item); err != nil {
			log.Errorf(err, "Proto unmarshal '%s' failed: %s", val, err)
			return
		}

		maps = append(maps, item)
		return
	})
	return maps
}

// UpdateMaps update all maps to etcd
func (s *storage) UpdateMaps(maps pb.SyncMapping) {
	srcMaps := s.GetMaps()
	mappings := make(pb.SyncMapping, 0, len(maps))
	for _, val := range maps {
		key := mappingsKey + "/" + val.ClusterName + "/" + val.OrgInstanceID
		data, err := proto.Marshal(val)
		if err != nil {
			log.Errorf(err, "Proto marshal failed: %s", err)
			continue
		}
		_, err = s.client.Put(context.Background(), key, util.BytesToStringWithNoCopy(data))
		if err != nil {
			log.Errorf(err, "Save mapping to etcd failed: %s", err)
			continue
		}
		mappings = append(mappings, val)
	}
	s.cleanExpired(srcMaps, mappings)
}

// GetMaps Get maps from storage
func (s *storage) GetMaps() (mapping pb.SyncMapping) {
	maps := make(pb.SyncMapping, 0, 10)
	s.getPrefixKey(mappingsKey, func(key, val []byte) (next bool) {
		next = true
		item := &pb.MappingEntry{}
		if err := proto.Unmarshal(val, item); err != nil {
			log.Errorf(err, "Proto unmarshal '%s' failed: %s", val, err)
			return
		}
		maps = append(maps, item)
		return
	})
	return maps
}
<commit_msg>save syncdata to etcd<commit_after>/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package etcd

import (
	"context"
	"sync"

	"github.com/apache/servicecomb-service-center/pkg/log"
	"github.com/apache/servicecomb-service-center/pkg/util"
	pb "github.com/apache/servicecomb-service-center/syncer/proto"
	"github.com/apache/servicecomb-service-center/syncer/servicecenter"
	"github.com/coreos/etcd/clientv3"
	"github.com/gogo/protobuf/proto"
)

var (
	// mappingsKey the key of instances mapping in etcd
	mappingsKey = "/syncer/v1/mappings"
	// servicesKey the key of service in etcd
	servicesKey = "/syncer/v1/services"
)

type storage struct {
	client *clientv3.Client
	data   *pb.SyncData
	lock   sync.RWMutex
}

func NewStorage(client *clientv3.Client) servicecenter.Storage {
	storage := &storage{
		client: client,
		data:   &pb.SyncData{},
	}
	return storage
}

// getPrefixKey Get data from etcd based on the prefix key
func (s *storage) getPrefixKey(prefix string, handler func(key, val []byte) (next bool)) {
	resp, err := s.client.Get(context.Background(), prefix, clientv3.WithPrefix())
	if err != nil {
		log.Errorf(err, "Get mapping from etcd failed: %s", err)
		return
	}

	for _, kv := range resp.Kvs {
		if !handler(kv.Key, kv.Value) {
			break
		}
	}
}

// UpdateData Update data to storage
func (s *storage) UpdateData(data *pb.SyncData) {
	s.UpdateServices(data.Services)
}

// GetData Get data from storage
func (s *storage) GetData() (data *pb.SyncData) {
	data = &pb.SyncData{Services: s.GetServices()}
	return
}

// cleanExpired clean the expired instances in the maps
func (s *storage) cleanExpired(sources pb.SyncMapping, actives pb.SyncMapping) {
next:
	for _, entry := range sources {
		for _, act := range actives {
			if act.CurInstanceID == entry.CurInstanceID {
				continue next
			}
		}
		key := mappingsKey + "/" + entry.ClusterName + "/" + entry.OrgInstanceID
		if _, err := s.client.Delete(context.Background(), key); err != nil {
			log.Errorf(err, "Delete instance clusterName=%s instanceID=%s failed", entry.ClusterName, entry.OrgInstanceID)
		}
	}
}

// UpdateServices Update services to storage
func (s *storage) UpdateServices(services []*pb.SyncService) {
	for _, val := range services {
		key := servicesKey + "/" + val.Service.ServiceId
		data, err := proto.Marshal(val)
		if err != nil {
			log.Errorf(err, "Proto marshal failed: %s", err)
			continue
		}
		_, err = s.client.Put(context.Background(), key, util.BytesToStringWithNoCopy(data))
		if err != nil {
			log.Errorf(err, "Save service to etcd failed: %s", err)
		}
	}
}

// GetServices Get services from storage
func (s *storage) GetServices() (services []*pb.SyncService) {
	services = make([]*pb.SyncService, 0, 10)
	s.getPrefixKey(servicesKey, func(key, val []byte) (next bool) {
		next = true
		item := &pb.SyncService{}
		if err := proto.Unmarshal(val, item); err != nil {
			log.Errorf(err, "Proto unmarshal '%s' failed: %s", val, err)
			return
		}
		services = append(services, item)
		return
	})
	return
}

// UpdateMapByCluster update map to storage by clusterName of other cluster
func (s *storage) UpdateMapByCluster(clusterName string, mapping pb.SyncMapping) {
	newMaps := make(pb.SyncMapping, 0, len(mapping))
	for _, val := range mapping {
		key := mappingsKey + "/" + clusterName + "/" + val.OrgInstanceID
		data, err := proto.Marshal(val)
		if err != nil {
			log.Errorf(err, "Proto marshal failed: %s", err)
			continue
		}
		_, err = s.client.Put(context.Background(), key, util.BytesToStringWithNoCopy(data))
		if err != nil {
			log.Errorf(err, "Save mapping to etcd failed: %s", err)
		}
		newMaps = append(newMaps, val)
	}
	s.cleanExpired(s.GetMapByCluster(clusterName), newMaps)
}

// GetMapByCluster get map by clusterName of other cluster
func (s *storage) GetMapByCluster(clusterName string) (mapping pb.SyncMapping) {
	maps := make(pb.SyncMapping, 0, 10)
	s.getPrefixKey(mappingsKey+"/"+clusterName, func(key, val []byte) (next bool) {
		next = true
		item := &pb.MappingEntry{}
		if err := proto.Unmarshal(val, item); err != nil {
			log.Errorf(err, "Proto unmarshal '%s' failed: %s", val, err)
			return
		}

		maps = append(maps, item)
		return
	})
	return maps
}

// UpdateMaps update all maps to etcd
func (s *storage) UpdateMaps(maps pb.SyncMapping) {
	srcMaps := s.GetMaps()
	mappings := make(pb.SyncMapping, 0, len(maps))
	for _, val := range maps {
		key := mappingsKey + "/" + val.ClusterName + "/" + val.OrgInstanceID
		data, err := proto.Marshal(val)
		if err != nil {
			log.Errorf(err, "Proto marshal failed: %s", err)
			continue
		}
		_, err = s.client.Put(context.Background(), key, util.BytesToStringWithNoCopy(data))
		if err != nil {
			log.Errorf(err, "Save mapping to etcd failed: %s", err)
			continue
		}
		mappings = append(mappings, val)
	}
	s.cleanExpired(srcMaps, mappings)
}

// GetMaps Get maps from storage
func (s *storage) GetMaps() (mapping pb.SyncMapping) {
	maps := make(pb.SyncMapping, 0, 10)
	s.getPrefixKey(mappingsKey, func(key, val []byte) (next bool) {
		next = true
		item := &pb.MappingEntry{}
		if err := proto.Unmarshal(val, item); err != nil {
			log.Errorf(err, "Proto unmarshal '%s' failed: %s", val, err)
			return
		}
		maps = append(maps, item)
		return
	})
	return maps
}
<|endoftext|>