<commit_before>// The colors package provide a simple way to bring colorful charcaters to terminal interface.
//
// This example will output the text with a Blue background and a Black foreground
//      colors.Println("@{bK}Example Text")
//
// This one will output the text with a red foreground
//      colors.Println("@rExample Text")
//
// This one will escape the @
//      colors.Println("@@")
//
// Full color syntax code
//      @{rgbcmykwRGBCMYKW}  foreground/background color
//      @{|}  Reset format style
//      @{!./_} Bold / Dim / Italic / underline
//      @{^&} Blink / Fast blink
//      @{?} Reverse the foreground and background color
//      @{-} Hide the text
// Note some of the functions are not widely supported, like "Fast blink" and "Italic".
package colors

import (
	"bytes"
	"fmt"
	"log"
)

// Escape character for color syntax
const escapeChar = '@'

// Short for reset to default style
var resetChar = fmt.Sprintf("%c|", escapeChar)

// Mapping from character to concrete escape code.
var codeMap = map[int]int{
	'|': 0,
	'!': 1,
	'.': 2,
	'/': 3,
	'_': 4,
	'^': 5,
	'&': 6,
	'?': 7,
	'-': 8,

	'k': 30,
	'r': 31,
	'g': 32,
	'y': 33,
	'b': 34,
	'm': 35,
	'c': 36,
	'w': 37,
	'd': 39,

	'K': 40,
	'R': 41,
	'G': 42,
	'Y': 43,
	'B': 44,
	'M': 45,
	'C': 46,
	'W': 47,
	'D': 49,
}

// Compile color syntax string like "rG" to escape code.
func colorMap(x string) string {
	attr := 0
	fg := 39
	bg := 49

	for _, key := range x {
		c, ok := codeMap[key]
		switch {
		case !ok:
			log.Fatalf("Wrong color syntax: %c", key)
		case 0 <= c && c <= 8:
			attr = c
		case 30 <= c && c <= 37:
			fg = c
		case 40 <= c && c <= 47:
			bg = c
		}
	}
	return fmt.Sprintf("\033[%d;%d;%dm", attr, fg, bg)
}

// Handle state after meeting one '@'
func compileColorSyntax(input, output *bytes.Buffer) {
	i, _, err := input.ReadRune()
	if err != nil {
		// EOF got
		log.Fatal("Parse failed on color syntax")
	}

	switch i {
	default:
		output.WriteString(colorMap(string(i)))
	case '{':
		color := bytes.NewBufferString("")
		for {
			i, _, err := input.ReadRune()
			if err != nil {
				log.Fatal("Parse failed on color syntax")
			}
			if i == '}' {
				break
			}
			color.WriteRune(i)
		}
		output.WriteString(colorMap(color.String()))
	case escapeChar:
		output.WriteRune(escapeChar)
	}
}

// Compile the string and replace color syntax with concrete escape code.
func compile(x string) string {
	if x == "" {
		return ""
	}

	input := bytes.NewBufferString(x)
	output := bytes.NewBufferString("")

	for {
		i, _, err := input.ReadRune()
		if err != nil {
			break
		}
		switch i {
		default:
			output.WriteRune(i)
		case escapeChar:
			compileColorSyntax(input, output)
		}
	}
	return output.String()
}

// Compile multiple values, only do compiling on string type.
func compileValues(a *[]interface{}) {
	for i, x := range *a {
		if str, ok := x.(string); ok {
			(*a)[i] = compile(str)
		}
	}
}

// Similar to fmt.Print, will reset the color at the end.
func Print(a ...interface{}) (int, error) {
	a = append(a, resetChar)
	compileValues(&a)
	return fmt.Print(a...)
}

// Similar to fmt.Println, will reset the color at the end.
func Println(a ...interface{}) (int, error) {
	a = append(a, resetChar)
	compileValues(&a)
	return fmt.Println(a...)
}

// Similar to fmt.Printf, will reset the color at the end.
func Printf(format string, a ...interface{}) (int, error) {
	format += resetChar
	format = compile(format)
	return fmt.Printf(format, a...)
}
<commit_msg>Add Sprintf & Sprint<commit_after>// The colors package provide a simple way to bring colorful charcaters to terminal interface.
//
// This example will output the text with a Blue background and a Black foreground
//      colors.Println("@{bK}Example Text")
//
// This one will output the text with a red foreground
//      colors.Println("@rExample Text")
//
// This one will escape the @
//      colors.Println("@@")
//
// Full color syntax code
//      @{rgbcmykwRGBCMYKW}  foreground/background color
//      @{|}  Reset format style
//      @{!./_} Bold / Dim / Italic / underline
//      @{^&} Blink / Fast blink
//      @{?} Reverse the foreground and background color
//      @{-} Hide the text
// Note some of the functions are not widely supported, like "Fast blink" and "Italic".
package colors

import (
	"bytes"
	"fmt"
	"log"
)

// Escape character for color syntax
const escapeChar = '@'

// Short for reset to default style
var resetChar = fmt.Sprintf("%c|", escapeChar)

// Mapping from character to concrete escape code.
var codeMap = map[int]int{
	'|': 0,
	'!': 1,
	'.': 2,
	'/': 3,
	'_': 4,
	'^': 5,
	'&': 6,
	'?': 7,
	'-': 8,

	'k': 30,
	'r': 31,
	'g': 32,
	'y': 33,
	'b': 34,
	'm': 35,
	'c': 36,
	'w': 37,
	'd': 39,

	'K': 40,
	'R': 41,
	'G': 42,
	'Y': 43,
	'B': 44,
	'M': 45,
	'C': 46,
	'W': 47,
	'D': 49,
}

// Compile color syntax string like "rG" to escape code.
func colorMap(x string) string {
	attr := 0
	fg := 39
	bg := 49

	for _, key := range x {
		c, ok := codeMap[key]
		switch {
		case !ok:
			log.Fatalf("Wrong color syntax: %c", key)
		case 0 <= c && c <= 8:
			attr = c
		case 30 <= c && c <= 37:
			fg = c
		case 40 <= c && c <= 47:
			bg = c
		}
	}
	return fmt.Sprintf("\033[%d;%d;%dm", attr, fg, bg)
}

// Handle state after meeting one '@'
func compileColorSyntax(input, output *bytes.Buffer) {
	i, _, err := input.ReadRune()
	if err != nil {
		// EOF got
		log.Fatal("Parse failed on color syntax")
	}

	switch i {
	default:
		output.WriteString(colorMap(string(i)))
	case '{':
		color := bytes.NewBufferString("")
		for {
			i, _, err := input.ReadRune()
			if err != nil {
				log.Fatal("Parse failed on color syntax")
			}
			if i == '}' {
				break
			}
			color.WriteRune(i)
		}
		output.WriteString(colorMap(color.String()))
	case escapeChar:
		output.WriteRune(escapeChar)
	}
}

// Compile the string and replace color syntax with concrete escape code.
func compile(x string) string {
	if x == "" {
		return ""
	}

	input := bytes.NewBufferString(x)
	output := bytes.NewBufferString("")

	for {
		i, _, err := input.ReadRune()
		if err != nil {
			break
		}
		switch i {
		default:
			output.WriteRune(i)
		case escapeChar:
			compileColorSyntax(input, output)
		}
	}
	return output.String()
}

// Compile multiple values, only do compiling on string type.
func compileValues(a *[]interface{}) {
	for i, x := range *a {
		if str, ok := x.(string); ok {
			(*a)[i] = compile(str)
		}
	}
}

// Similar to fmt.Print, will reset the color at the end.
func Print(a ...interface{}) (int, error) {
	a = append(a, resetChar)
	compileValues(&a)
	return fmt.Print(a...)
}

// Similar to fmt.Println, will reset the color at the end.
func Println(a ...interface{}) (int, error) {
	a = append(a, resetChar)
	compileValues(&a)
	return fmt.Println(a...)
}

// Similar to fmt.Printf, will reset the color at the end.
func Printf(format string, a ...interface{}) (int, error) {
	format += resetChar
	format = compile(format)
	return fmt.Printf(format, a...)
}

// Similar to fmt.Sprint, will reset the color at the end.
func Sprint(a ...interface{}) string {
        a = append(a, resetChar)
        compileValues(&a)
        return fmt.Sprint(a...)
}

// Similar to fmt.Sprintf, will reset the color at the end.
func Sprintf(format string, a ...interface{}) string {
        format += resetChar
        format = compile(format)
        return fmt.Sprintf(format, a...)
}
<|endoftext|>


<commit_before>package report

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/pajlada/pajbot2/pkg"
	"github.com/pajlada/pajbot2/pkg/pubsub"
	"github.com/pajlada/pajbot2/pkg/users"
)

type ReportUser struct {
	ID   string
	Name string
	Type string `json:",omitempty"`
}

type Report struct {
	ID       uint32
	Channel  ReportUser
	Reporter ReportUser
	Target   ReportUser
	Reason   string `json:",omitempty"`
	Logs     []string
	Time     time.Time
}

type Holder struct {
	db        *sql.DB
	pubSub    *pubsub.PubSub
	userStore pkg.UserStore

	reportsMutex *sync.Mutex
	reports      map[uint32]Report
}

var _ pubsub.Connection = &Holder{}
var _ pubsub.SubscriptionHandler = &Holder{}

func New(db *sql.DB, pubSub *pubsub.PubSub, userStore pkg.UserStore) (*Holder, error) {
	h := &Holder{
		db:        db,
		pubSub:    pubSub,
		userStore: userStore,

		reportsMutex: &sync.Mutex{},
		reports:      make(map[uint32]Report),
	}

	err := h.Load()
	if err != nil {
		return nil, err
	}

	pubSub.Subscribe(h, "HandleReport", nil)
	pubSub.Subscribe(h, "TimeoutEvent", nil)
	pubSub.Subscribe(h, "BanEvent", nil)
	pubSub.HandleSubscribe(h, "ReportReceived")

	return h, nil
}

func (h *Holder) Load() error {
	rows, err := h.db.Query("SELECT `id`, `channel_id`, `channel_name`, `channel_type`, `reporter_id`, `reporter_name`, `target_id`, `target_name`, `reason`, `logs`, `time` FROM `Report`")
	if err != nil {
		return err
	}
	defer rows.Close()

	h.reportsMutex.Lock()
	defer h.reportsMutex.Unlock()

	for rows.Next() {
		var report Report
		var logsString string

		if err := rows.Scan(&report.ID, &report.Channel.ID, &report.Channel.Name, &report.Channel.Type, &report.Reporter.ID, &report.Reporter.Name, &report.Target.ID, &report.Target.Name, &report.Reason, &logsString, &report.Time); err != nil {
			return err
		}

		report.Logs = strings.Split(logsString, "\n")

		h.reports[report.ID] = report
	}

	return nil
}

type handleReportMessage struct {
	Action    string
	ChannelID string
	ReportID  uint32
	Duration  *uint32
}

func (h *Holder) Register(report Report) (*Report, bool, error) {
	const queryF = `
	INSERT INTO Report
		(channel_id, channel_name, channel_type,
		reporter_id, reporter_name, target_id, target_name, reason, logs, time)
	VALUES
		(?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	`

	h.reportsMutex.Lock()
	defer h.reportsMutex.Unlock()

	// Don't accept reports for users that have already been reported
	for _, oldReport := range h.reports {
		if oldReport.Channel.ID == report.Channel.ID && oldReport.Target.ID == report.Target.ID {
			fmt.Println("Report already registered for this target in this channel")
			return &oldReport, false, nil
		}
	}

	res, err := h.db.Exec(queryF, report.Channel.ID, report.Channel.Name, report.Channel.Type, report.Reporter.ID, report.Reporter.Name, report.Target.ID, report.Target.Name, report.Reason, strings.Join(report.Logs, "\n"), report.Time)
	if err != nil {
		fmt.Printf("Error inserting report %v into SQL: %s\n", report, err)
		return nil, false, err
	}

	id, err := res.LastInsertId()
	if err != nil {
		fmt.Printf("Error getting last insert id: %s\n", err)
		return nil, false, err
	}

	report.ID = uint32(id)

	h.pubSub.Publish("ReportReceived", report, pkg.PubSubAdminAuth())

	h.reports[report.ID] = report

	return &report, true, nil
}

func (h *Holder) Update(report Report) error {
	if report.ID == 0 {
		return errors.New("Missing report ID in Update")
	}

	const queryF = `UPDATE Report SET time=?, logs=? WHERE id=?`
	_, err := h.db.Exec(queryF, report.Time, strings.Join(report.Logs, "\n"), report.ID)
	if err != nil {
		return err
	}

	h.reportsMutex.Lock()
	defer h.reportsMutex.Unlock()
	h.reports[report.ID] = report

	// TODO: Send some "ReportUpdated" message

	return nil
}

type reportHandled struct {
	ReportID uint32
	Handler  ReportUser
	Action   string
}

func (h *Holder) handleReport(action handleReportMessage, auth *pkg.PubSubAuthorization) error {
	h.reportsMutex.Lock()
	defer h.reportsMutex.Unlock()

	if auth == nil {
		fmt.Println("Missing auth in HandleReport")
		return nil
	}

	report, ok := h.reports[action.ReportID]
	if !ok {
		fmt.Printf("No report found with ID %d\n", action.ReportID)
		// No report found with this ID
		return nil
	}

	// Remove report from SQL and our local map
	err := h.dismissReport(report.ID)
	if err != nil {
		fmt.Println("Error dismissing report", err)
		return err
	}

	// TODO: Insert into new table: HandledReport

	msg := reportHandled{
		ReportID: report.ID,
		Handler: ReportUser{
			ID:   auth.TwitchUserID,
			Name: h.userStore.GetName(auth.TwitchUserID),
		},
		Action: action.Action,
	}

	h.pubSub.Publish("ReportHandled", &msg, pkg.PubSubAdminAuth())

	switch action.Action {
	case "ban":
		h.pubSub.Publish("Ban", &pkg.PubSubBan{
			Channel: report.Channel.Name,
			Target:  report.Target.Name,
			Reason:  report.Reason,
		}, pkg.PubSubAdminAuth())

	case "timeout":
		var duration uint32
		duration = 600
		if action.Duration != nil {
			duration = *action.Duration
		}
		h.pubSub.Publish("Timeout", &pkg.PubSubTimeout{
			Channel:  report.Channel.Name,
			Target:   report.Target.Name,
			Duration: duration,
			Reason:   report.Reason,
		}, pkg.PubSubAdminAuth())

	case "undo":
		h.pubSub.Publish("Untimeout", &pkg.PubSubUntimeout{
			Channel: report.Channel.Name,
			Target:  report.Target.Name,
		}, pkg.PubSubAdminAuth())
	default:
		fmt.Println("Unhandled action", action.Action)
	}

	return nil
}

// dismissReport assumes that reportsMutex has already been locked
func (h *Holder) dismissReport(reportID uint32) error {
	// Delete from SQL
	const queryF = "DELETE FROM Report WHERE `id`=?"

	_, err := h.db.Exec(queryF, reportID)
	if err != nil {
		return err
	}

	// Delete from our internal storage
	delete(h.reports, reportID)

	return nil
}

func (h *Holder) handleBanEvent(banEvent pkg.PubSubBanEvent) error {
	h.reportsMutex.Lock()
	defer h.reportsMutex.Unlock()

	for reportID, report := range h.reports {
		if report.Channel.ID == banEvent.Channel.ID && report.Target.ID == banEvent.Target.ID {
			// Found matching report
			h.dismissReport(reportID)

			break
		}
	}

	return nil
}

func (h *Holder) MessageReceived(topic string, data []byte, auth *pkg.PubSubAuthorization) error {
	switch topic {
	case "HandleReport":
		var msg handleReportMessage
		err := json.Unmarshal(data, &msg)
		if err != nil {
			fmt.Println("Error unmarshalling:", err)
			return nil
		}

		fmt.Printf("Handle report: %#v\n", msg)

		return h.handleReport(msg, auth)

	case "BanEvent":
		var msg pkg.PubSubBanEvent
		err := json.Unmarshal(data, &msg)
		if err != nil {
			fmt.Println("Error unmarshalling:", err)
			return nil
		}

		return h.handleBanEvent(msg)
	}

	return nil
}

func (h *Holder) ConnectionSubscribed(connection pubsub.Connection, topic string, auth *pkg.PubSubAuthorization) (error, bool) {
	switch topic {
	case "ReportReceived":
		if auth == nil {
			return nil, false
		}

		// Verify authorization
		const queryF = `
SELECT twitch_username FROM User
	WHERE twitch_userid=? AND twitch_nonce=? LIMIT 1;
`

		rows, err := h.db.Query(queryF, auth.TwitchUserID, auth.Nonce)
		if err != nil {
			fmt.Println(err)
			return err, true
		}
		defer rows.Close()

		if !rows.Next() {
			return nil, false
		}

		hasPermission, err := users.HasGlobalPermission(auth.TwitchUserID, pkg.PermissionModeration)
		if err != nil {
			fmt.Println(err)
			return err, false
		}

		if !hasPermission {
			return nil, false
		}

		fmt.Println("Send reports to new connection")

		for _, report := range h.reports {
			bytes, err := json.Marshal(report)
			if err != nil {
				fmt.Println(err)
				return err, true
			}
			connection.MessageReceived(topic, bytes, pkg.PubSubAdminAuth())
		}
	}

	return nil, true
}
<commit_msg>report reason in report timeouts/bans<commit_after>package report

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/pajlada/pajbot2/pkg"
	"github.com/pajlada/pajbot2/pkg/pubsub"
	"github.com/pajlada/pajbot2/pkg/users"
)

type ReportUser struct {
	ID   string
	Name string
	Type string `json:",omitempty"`
}

type Report struct {
	ID       uint32
	Channel  ReportUser
	Reporter ReportUser
	Target   ReportUser
	Reason   string `json:",omitempty"`
	Logs     []string
	Time     time.Time
}

type Holder struct {
	db        *sql.DB
	pubSub    *pubsub.PubSub
	userStore pkg.UserStore

	reportsMutex *sync.Mutex
	reports      map[uint32]Report
}

var _ pubsub.Connection = &Holder{}
var _ pubsub.SubscriptionHandler = &Holder{}

func New(db *sql.DB, pubSub *pubsub.PubSub, userStore pkg.UserStore) (*Holder, error) {
	h := &Holder{
		db:        db,
		pubSub:    pubSub,
		userStore: userStore,

		reportsMutex: &sync.Mutex{},
		reports:      make(map[uint32]Report),
	}

	err := h.Load()
	if err != nil {
		return nil, err
	}

	pubSub.Subscribe(h, "HandleReport", nil)
	pubSub.Subscribe(h, "TimeoutEvent", nil)
	pubSub.Subscribe(h, "BanEvent", nil)
	pubSub.HandleSubscribe(h, "ReportReceived")

	return h, nil
}

func (h *Holder) Load() error {
	rows, err := h.db.Query("SELECT `id`, `channel_id`, `channel_name`, `channel_type`, `reporter_id`, `reporter_name`, `target_id`, `target_name`, `reason`, `logs`, `time` FROM `Report`")
	if err != nil {
		return err
	}
	defer rows.Close()

	h.reportsMutex.Lock()
	defer h.reportsMutex.Unlock()

	for rows.Next() {
		var report Report
		var logsString string

		if err := rows.Scan(&report.ID, &report.Channel.ID, &report.Channel.Name, &report.Channel.Type, &report.Reporter.ID, &report.Reporter.Name, &report.Target.ID, &report.Target.Name, &report.Reason, &logsString, &report.Time); err != nil {
			return err
		}

		report.Logs = strings.Split(logsString, "\n")

		h.reports[report.ID] = report
	}

	return nil
}

type handleReportMessage struct {
	Action    string
	ChannelID string
	ReportID  uint32
	Duration  *uint32
}

func (h *Holder) Register(report Report) (*Report, bool, error) {
	const queryF = `
	INSERT INTO Report
		(channel_id, channel_name, channel_type,
		reporter_id, reporter_name, target_id, target_name, reason, logs, time)
	VALUES
		(?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	`

	h.reportsMutex.Lock()
	defer h.reportsMutex.Unlock()

	// Don't accept reports for users that have already been reported
	for _, oldReport := range h.reports {
		if oldReport.Channel.ID == report.Channel.ID && oldReport.Target.ID == report.Target.ID {
			fmt.Println("Report already registered for this target in this channel")
			return &oldReport, false, nil
		}
	}

	res, err := h.db.Exec(queryF, report.Channel.ID, report.Channel.Name, report.Channel.Type, report.Reporter.ID, report.Reporter.Name, report.Target.ID, report.Target.Name, report.Reason, strings.Join(report.Logs, "\n"), report.Time)
	if err != nil {
		fmt.Printf("Error inserting report %v into SQL: %s\n", report, err)
		return nil, false, err
	}

	id, err := res.LastInsertId()
	if err != nil {
		fmt.Printf("Error getting last insert id: %s\n", err)
		return nil, false, err
	}

	report.ID = uint32(id)

	h.pubSub.Publish("ReportReceived", report, pkg.PubSubAdminAuth())

	h.reports[report.ID] = report

	return &report, true, nil
}

func (h *Holder) Update(report Report) error {
	if report.ID == 0 {
		return errors.New("Missing report ID in Update")
	}

	const queryF = `UPDATE Report SET time=?, logs=? WHERE id=?`
	_, err := h.db.Exec(queryF, report.Time, strings.Join(report.Logs, "\n"), report.ID)
	if err != nil {
		return err
	}

	h.reportsMutex.Lock()
	defer h.reportsMutex.Unlock()
	h.reports[report.ID] = report

	// TODO: Send some "ReportUpdated" message

	return nil
}

type reportHandled struct {
	ReportID uint32
	Handler  ReportUser
	Action   string
}

func (h *Holder) handleReport(action handleReportMessage, auth *pkg.PubSubAuthorization) error {
	h.reportsMutex.Lock()
	defer h.reportsMutex.Unlock()

	if auth == nil {
		fmt.Println("Missing auth in HandleReport")
		return nil
	}

	report, ok := h.reports[action.ReportID]
	if !ok {
		fmt.Printf("No report found with ID %d\n", action.ReportID)
		// No report found with this ID
		return nil
	}

	// Remove report from SQL and our local map
	err := h.dismissReport(report.ID)
	if err != nil {
		fmt.Println("Error dismissing report", err)
		return err
	}

	// TODO: Insert into new table: HandledReport

	msg := reportHandled{
		ReportID: report.ID,
		Handler: ReportUser{
			ID:   auth.TwitchUserID,
			Name: h.userStore.GetName(auth.TwitchUserID),
		},
		Action: action.Action,
	}

	h.pubSub.Publish("ReportHandled", &msg, pkg.PubSubAdminAuth())

	switch action.Action {
	case "ban":
		h.pubSub.Publish("Ban", &pkg.PubSubBan{
			Channel: report.Channel.Name,
			Target:  report.Target.Name,
			// Reason:  report.Reason,
		}, pkg.PubSubAdminAuth())

	case "timeout":
		var duration uint32
		duration = 600
		if action.Duration != nil {
			duration = *action.Duration
		}
		h.pubSub.Publish("Timeout", &pkg.PubSubTimeout{
			Channel:  report.Channel.Name,
			Target:   report.Target.Name,
			Duration: duration,
			// Reason:   report.Reason,
		}, pkg.PubSubAdminAuth())

	case "undo":
		h.pubSub.Publish("Untimeout", &pkg.PubSubUntimeout{
			Channel: report.Channel.Name,
			Target:  report.Target.Name,
		}, pkg.PubSubAdminAuth())
	default:
		fmt.Println("Unhandled action", action.Action)
	}

	return nil
}

// dismissReport assumes that reportsMutex has already been locked
func (h *Holder) dismissReport(reportID uint32) error {
	// Delete from SQL
	const queryF = "DELETE FROM Report WHERE `id`=?"

	_, err := h.db.Exec(queryF, reportID)
	if err != nil {
		return err
	}

	// Delete from our internal storage
	delete(h.reports, reportID)

	return nil
}

func (h *Holder) handleBanEvent(banEvent pkg.PubSubBanEvent) error {
	h.reportsMutex.Lock()
	defer h.reportsMutex.Unlock()

	for reportID, report := range h.reports {
		if report.Channel.ID == banEvent.Channel.ID && report.Target.ID == banEvent.Target.ID {
			// Found matching report
			h.dismissReport(reportID)

			break
		}
	}

	return nil
}

func (h *Holder) MessageReceived(topic string, data []byte, auth *pkg.PubSubAuthorization) error {
	switch topic {
	case "HandleReport":
		var msg handleReportMessage
		err := json.Unmarshal(data, &msg)
		if err != nil {
			fmt.Println("Error unmarshalling:", err)
			return nil
		}

		fmt.Printf("Handle report: %#v\n", msg)

		return h.handleReport(msg, auth)

	case "BanEvent":
		var msg pkg.PubSubBanEvent
		err := json.Unmarshal(data, &msg)
		if err != nil {
			fmt.Println("Error unmarshalling:", err)
			return nil
		}

		return h.handleBanEvent(msg)
	}

	return nil
}

func (h *Holder) ConnectionSubscribed(connection pubsub.Connection, topic string, auth *pkg.PubSubAuthorization) (error, bool) {
	switch topic {
	case "ReportReceived":
		if auth == nil {
			return nil, false
		}

		// Verify authorization
		const queryF = `
SELECT twitch_username FROM User
	WHERE twitch_userid=? AND twitch_nonce=? LIMIT 1;
`

		rows, err := h.db.Query(queryF, auth.TwitchUserID, auth.Nonce)
		if err != nil {
			fmt.Println(err)
			return err, true
		}
		defer rows.Close()

		if !rows.Next() {
			return nil, false
		}

		hasPermission, err := users.HasGlobalPermission(auth.TwitchUserID, pkg.PermissionModeration)
		if err != nil {
			fmt.Println(err)
			return err, false
		}

		if !hasPermission {
			return nil, false
		}

		fmt.Println("Send reports to new connection")

		for _, report := range h.reports {
			bytes, err := json.Marshal(report)
			if err != nil {
				fmt.Println(err)
				return err, true
			}
			connection.MessageReceived(topic, bytes, pkg.PubSubAdminAuth())
		}
	}

	return nil, true
}
<|endoftext|>


<commit_before>package server

import (
	jnet "junta/net"
	"junta/paxos"
	"junta/proto"
	"junta/store"
	"junta/util"
	"net"
	"os"
	"reflect"
	"strconv"
	"strings"
	"time"
)

const packetSize = 3000

const lease = 3e9 // ns == 3s

var ErrBadPrefix = os.NewError("bad prefix in path")

type conn struct {
	*proto.Conn
	c   net.Conn
	s   *Server
	cal bool
}

type Manager interface {
	paxos.Proposer
	PutFrom(string, paxos.Msg)
	Alpha() int
}

type Server struct {
	Addr         string
	St           *store.Store
	Mg           Manager
	Self, Prefix string
}

func (sv *Server) ListenAndServeUdp(outs chan paxos.Packet) os.Error {
	logger := util.NewLogger("udp server %s", sv.Addr)

	logger.Println("binding")
	u, err := net.ListenPacket("udp", sv.Addr)
	if err != nil {
		logger.Println(err)
		return err
	}
	defer u.Close()
	logger.Println("listening")

	err = sv.ServeUdp(u, outs)
	if err != nil {
		logger.Printf("%s: %s", u, err)
	}
	return err
}

func (sv *Server) ServeUdp(u jnet.Conn, outs chan paxos.Packet) os.Error {
	r := jnet.Ackify(u, outs)

	for p := range r {
		sv.Mg.PutFrom(p.Addr, p.Msg)
	}

	panic("unreachable")
}

var clg = util.NewLogger("cal")

func (s *Server) Serve(l net.Listener, cal chan int) os.Error {
	var ok bool
	for {
		if !ok {
			_, ok = <-cal
		}
		clg.Println(ok)

		rw, e := l.Accept()
		if e != nil {
			return e
		}
		c := &conn{proto.NewConn(rw), rw, s, ok}
		go c.serve()
	}

	panic("unreachable")
}

func (sv *Server) leader() string {
	parts, cas := sv.St.Get("/junta/leader")
	if cas == store.Dir && cas == store.Missing {
		return ""
	}
	return parts[0]
}

func (sv *Server) addrFor(id string) string {
	parts, cas := sv.St.Get("/junta/members/" + id)
	if cas == store.Dir && cas == store.Missing {
		return ""
	}
	return parts[0]
}

// Checks that path begins with the proper prefix and returns the short path
// without the prefix.
func (sv *Server) checkPath(path string) (string, os.Error) {
	logger := util.NewLogger("checkPath")
	if !strings.HasPrefix(path, sv.Prefix+"/") {
		logger.Printf("prefix %q not in %q", sv.Prefix+"/", path)
		return "", ErrBadPrefix
	}
	return path[len(sv.Prefix):], nil
}

// Repeatedly propose nop values until a successful read from `done`.
func (sv *Server) AdvanceUntil(done chan int) {
	for _, ok := <-done; !ok; _, ok = <-done {
		sv.Mg.Propose(store.Nop)
	}
}

func (c *conn) redirect(rid uint) {
	leader := c.s.leader()
	addr := c.s.addrFor(leader)
	if addr == "" {
		c.SendError(rid, "unknown address for leader")
	} else {
		c.SendRedirect(rid, addr)
	}
}

func get(s *Server, data interface{}) (interface{}, os.Error) {
	r := data.(*proto.ReqGet)
	shortPath, err := s.checkPath(r.Path)
	if err != nil {
		return nil, err
	}

	v, cas := s.St.Get(shortPath)
	return []interface{}{v, cas}, nil
}

func sget(s *Server, data interface{}) (interface{}, os.Error) {
	r := data.(*proto.ReqGet)
	shortPath, err := s.checkPath(r.Path)
	if err != nil {
		return "", err
	}

	body, err := store.GetString(s.St.SyncPath(shortPath), shortPath), nil
	if err != nil {
		return nil, err
	}

	return []interface{}{body}, nil
}

func set(s *Server, data interface{}) (interface{}, os.Error) {
	r := data.(*proto.ReqSet)

	shortPath, err := s.checkPath(r.Path)
	if err != nil {
		return nil, err
	}

	seqn, _, err := paxos.Set(s.Mg, shortPath, r.Body, r.Cas)
	if err != nil {
		return nil, err
	}
	return []interface{}{strconv.Uitoa64(seqn)}, nil
}

func del(s *Server, data interface{}) (interface{}, os.Error) {
	r := data.(*proto.ReqDel)

	shortPath, err := s.checkPath(r.Path)
	if err != nil {
		return 0, err
	}

	_, err = paxos.Del(s.Mg, shortPath, r.Cas)
	if err != nil {
		return nil, err
	}

	return []interface{}{"true"}, nil
}

func nop(s *Server, data interface{}) (interface{}, os.Error) {
	s.Mg.Propose(store.Nop)
	return []interface{}{"true"}, nil
}

func join(s *Server, data interface{}) (interface{}, os.Error) {
	r := data.(*proto.ReqJoin)
	key := "/junta/members/" + r.Who
	seqn, _, err := paxos.Set(s.Mg, key, r.Addr, store.Missing)
	if err != nil {
		return nil, err
	}

	done := make(chan int)
	go s.AdvanceUntil(done)
	s.St.Sync(seqn + uint64(s.Mg.Alpha()))
	close(done)
	seqn, snap := s.St.Snapshot()
	return []interface{}{strconv.Uitoa64(seqn), snap}, nil
}

func checkin(s *Server, data interface{}) (interface{}, os.Error) {
	r := data.(*proto.ReqCheckin)
	t := time.Nanoseconds() + lease
	_, cas, err := paxos.Set(s.Mg, "/session/"+r.Sid, strconv.Itoa64(t), r.Cas)
	if err != nil {
		return nil, err
	}
	return []interface{}{strconv.Itoa64(t), cas}, nil
}

func indirect(x interface{}) interface{} {
	return reflect.Indirect(reflect.NewValue(x)).Interface()
}

type op struct {
	p interface{}
	f func(*Server, interface{}) (interface{}, os.Error)

	redirect bool
}

var ops = map[string]op{
	"get":{p:new(*proto.ReqGet), f:get},
	"sget":{p:new(*proto.ReqGet), f:sget},
	"set":{p:new(*proto.ReqSet), f:set, redirect:true},
	"del":{p:new(*proto.ReqDel), f:del, redirect:true},
	"nop":{p:new(*[]interface{}), f:nop, redirect:true},
	"join":{p:new(*proto.ReqJoin), f:join, redirect:true},
	"checkin":{p:new(*proto.ReqCheckin), f:checkin, redirect:true},
}

func (c *conn) serve() {
	logger := util.NewLogger("%v", c.c.RemoteAddr())
	logger.Println("accepted connection")
	for {
		rid, verb, data, err := c.ReadRequest()
		if err != nil {
			if err == os.EOF {
				logger.Println("connection closed by peer")
			} else {
				logger.Println(err)
			}
			return
		}

		rlogger := util.NewLogger("%v - req [%d]", c.c.RemoteAddr(), rid)

		if o, ok := ops[verb]; ok {
			rlogger.Printf("%s %v", verb, data)

			err := proto.Fit(data, o.p)
			if err != nil {
				c.SendError(rid, err.String())
				continue
			}

			if o.redirect && !c.cal {
				c.redirect(rid)
				continue
			}

			res, err := o.f(c.s, indirect(o.p))
			if err != nil {
				c.SendError(rid, err.String())
			} else {
				c.SendResponse(rid, res)
			}
			continue
		}

		rlogger.Printf("unknown command <%s>", verb)
		c.SendError(rid, proto.InvalidCommand+" "+verb)
	}
}
<commit_msg>refactor<commit_after>package server

import (
	jnet "junta/net"
	"junta/paxos"
	"junta/proto"
	"junta/store"
	"junta/util"
	"net"
	"os"
	"reflect"
	"strconv"
	"strings"
	"time"
)

const packetSize = 3000

const lease = 3e9 // ns == 3s

var ErrBadPrefix = os.NewError("bad prefix in path")

type conn struct {
	*proto.Conn
	c   net.Conn
	s   *Server
	cal bool
}

type Manager interface {
	paxos.Proposer
	PutFrom(string, paxos.Msg)
	Alpha() int
}

type Server struct {
	Addr         string
	St           *store.Store
	Mg           Manager
	Self, Prefix string
}

func (sv *Server) ListenAndServeUdp(outs chan paxos.Packet) os.Error {
	logger := util.NewLogger("udp server %s", sv.Addr)

	logger.Println("binding")
	u, err := net.ListenPacket("udp", sv.Addr)
	if err != nil {
		logger.Println(err)
		return err
	}
	defer u.Close()
	logger.Println("listening")

	err = sv.ServeUdp(u, outs)
	if err != nil {
		logger.Printf("%s: %s", u, err)
	}
	return err
}

func (sv *Server) ServeUdp(u jnet.Conn, outs chan paxos.Packet) os.Error {
	r := jnet.Ackify(u, outs)

	for p := range r {
		sv.Mg.PutFrom(p.Addr, p.Msg)
	}

	panic("unreachable")
}

var clg = util.NewLogger("cal")

func (s *Server) Serve(l net.Listener, cal chan int) os.Error {
	var ok bool
	for {
		if !ok {
			_, ok = <-cal
		}
		clg.Println(ok)

		rw, e := l.Accept()
		if e != nil {
			return e
		}
		c := &conn{proto.NewConn(rw), rw, s, ok}
		go c.serve()
	}

	panic("unreachable")
}

func (sv *Server) leader() string {
	parts, cas := sv.St.Get("/junta/leader")
	if cas == store.Dir && cas == store.Missing {
		return ""
	}
	return parts[0]
}

func (sv *Server) addrFor(id string) string {
	parts, cas := sv.St.Get("/junta/members/" + id)
	if cas == store.Dir && cas == store.Missing {
		return ""
	}
	return parts[0]
}

// Checks that path begins with the proper prefix and returns the short path
// without the prefix.
func (sv *Server) checkPath(path string) (string, os.Error) {
	logger := util.NewLogger("checkPath")
	if !strings.HasPrefix(path, sv.Prefix+"/") {
		logger.Printf("prefix %q not in %q", sv.Prefix+"/", path)
		return "", ErrBadPrefix
	}
	return path[len(sv.Prefix):], nil
}

// Repeatedly propose nop values until a successful read from `done`.
func (sv *Server) AdvanceUntil(done chan int) {
	for _, ok := <-done; !ok; _, ok = <-done {
		sv.Mg.Propose(store.Nop)
	}
}

func (c *conn) redirect(rid uint) {
	leader := c.s.leader()
	addr := c.s.addrFor(leader)
	if addr == "" {
		c.SendError(rid, "unknown address for leader")
	} else {
		c.SendRedirect(rid, addr)
	}
}

func get(s *Server, data interface{}) (interface{}, os.Error) {
	r := data.(*proto.ReqGet)
	shortPath, err := s.checkPath(r.Path)
	if err != nil {
		return nil, err
	}

	v, cas := s.St.Get(shortPath)
	return []interface{}{v, cas}, nil
}

func sget(s *Server, data interface{}) (interface{}, os.Error) {
	r := data.(*proto.ReqGet)
	shortPath, err := s.checkPath(r.Path)
	if err != nil {
		return "", err
	}

	body, err := store.GetString(s.St.SyncPath(shortPath), shortPath), nil
	if err != nil {
		return nil, err
	}

	return []interface{}{body}, nil
}

func set(s *Server, data interface{}) (interface{}, os.Error) {
	r := data.(*proto.ReqSet)

	shortPath, err := s.checkPath(r.Path)
	if err != nil {
		return nil, err
	}

	seqn, _, err := paxos.Set(s.Mg, shortPath, r.Body, r.Cas)
	if err != nil {
		return nil, err
	}
	return []interface{}{strconv.Uitoa64(seqn)}, nil
}

func del(s *Server, data interface{}) (interface{}, os.Error) {
	r := data.(*proto.ReqDel)

	shortPath, err := s.checkPath(r.Path)
	if err != nil {
		return 0, err
	}

	_, err = paxos.Del(s.Mg, shortPath, r.Cas)
	if err != nil {
		return nil, err
	}

	return []interface{}{"true"}, nil
}

func nop(s *Server, data interface{}) (interface{}, os.Error) {
	s.Mg.Propose(store.Nop)
	return []interface{}{"true"}, nil
}

func join(s *Server, data interface{}) (interface{}, os.Error) {
	r := data.(*proto.ReqJoin)
	key := "/junta/members/" + r.Who
	seqn, _, err := paxos.Set(s.Mg, key, r.Addr, store.Missing)
	if err != nil {
		return nil, err
	}

	done := make(chan int)
	go s.AdvanceUntil(done)
	s.St.Sync(seqn + uint64(s.Mg.Alpha()))
	close(done)
	seqn, snap := s.St.Snapshot()
	return []interface{}{strconv.Uitoa64(seqn), snap}, nil
}

func checkin(s *Server, data interface{}) (interface{}, os.Error) {
	r := data.(*proto.ReqCheckin)
	t := time.Nanoseconds() + lease
	_, cas, err := paxos.Set(s.Mg, "/session/"+r.Sid, strconv.Itoa64(t), r.Cas)
	if err != nil {
		return nil, err
	}
	return []interface{}{strconv.Itoa64(t), cas}, nil
}

func indirect(x interface{}) interface{} {
	return reflect.Indirect(reflect.NewValue(x)).Interface()
}

type handler func(*Server, interface{}) (interface{}, os.Error)

type op struct {
	p interface{}
	f handler

	redirect bool
}

var ops = map[string]op{
	"get":{p:new(*proto.ReqGet), f:get},
	"sget":{p:new(*proto.ReqGet), f:sget},
	"set":{p:new(*proto.ReqSet), f:set, redirect:true},
	"del":{p:new(*proto.ReqDel), f:del, redirect:true},
	"nop":{p:new(*[]interface{}), f:nop, redirect:true},
	"join":{p:new(*proto.ReqJoin), f:join, redirect:true},
	"checkin":{p:new(*proto.ReqCheckin), f:checkin, redirect:true},
}

func (c *conn) handle(rid uint, f handler, data interface{}) {
	res, err := f(c.s, data)
	if err != nil {
		c.SendError(rid, err.String())
	} else {
		c.SendResponse(rid, res)
	}
}

func (c *conn) serve() {
	logger := util.NewLogger("%v", c.c.RemoteAddr())
	logger.Println("accepted connection")
	for {
		rid, verb, data, err := c.ReadRequest()
		if err != nil {
			if err == os.EOF {
				logger.Println("connection closed by peer")
			} else {
				logger.Println(err)
			}
			return
		}

		rlogger := util.NewLogger("%v - req [%d]", c.c.RemoteAddr(), rid)

		if o, ok := ops[verb]; ok {
			rlogger.Printf("%s %v", verb, data)

			err := proto.Fit(data, o.p)
			if err != nil {
				c.SendError(rid, err.String())
				continue
			}

			if o.redirect && !c.cal {
				c.redirect(rid)
				continue
			}

			c.handle(rid, o.f, indirect(o.p))
			continue
		}

		rlogger.Printf("unknown command <%s>", verb)
		c.SendError(rid, proto.InvalidCommand+" "+verb)
	}
}
<|endoftext|>


<commit_before>package server

import (
	"context"
	"encoding/json"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"time"

	"github.com/containous/traefik/v2/pkg/api"
	"github.com/containous/traefik/v2/pkg/config/dynamic"
	"github.com/containous/traefik/v2/pkg/config/runtime"
	"github.com/containous/traefik/v2/pkg/config/static"
	"github.com/containous/traefik/v2/pkg/log"
	"github.com/containous/traefik/v2/pkg/metrics"
	"github.com/containous/traefik/v2/pkg/middlewares/accesslog"
	"github.com/containous/traefik/v2/pkg/middlewares/requestdecorator"
	"github.com/containous/traefik/v2/pkg/provider"
	"github.com/containous/traefik/v2/pkg/safe"
	"github.com/containous/traefik/v2/pkg/tls"
	"github.com/containous/traefik/v2/pkg/tracing"
	"github.com/containous/traefik/v2/pkg/tracing/jaeger"
	"github.com/containous/traefik/v2/pkg/types"
)

// Server is the reverse-proxy/load-balancer engine
type Server struct {
	entryPointsTCP             TCPEntryPoints
	configurationChan          chan dynamic.Message
	configurationValidatedChan chan dynamic.Message
	signals                    chan os.Signal
	stopChan                   chan bool
	currentConfigurations      safe.Safe
	providerConfigUpdateMap    map[string]chan dynamic.Message
	accessLoggerMiddleware     *accesslog.Handler
	tracer                     *tracing.Tracing
	routinesPool               *safe.Pool
	defaultRoundTripper        http.RoundTripper
	metricsRegistry            metrics.Registry
	provider                   provider.Provider
	configurationListeners     []func(dynamic.Configuration)
	requestDecorator           *requestdecorator.RequestDecorator
	providersThrottleDuration  time.Duration
	tlsManager                 *tls.Manager
	api                        func(configuration *runtime.Configuration) http.Handler
	restHandler                http.Handler
}

// RouteAppenderFactory the route appender factory interface
type RouteAppenderFactory interface {
	NewAppender(ctx context.Context, runtimeConfiguration *runtime.Configuration) types.RouteAppender
}

func setupTracing(conf *static.Tracing) tracing.Backend {
	var backend tracing.Backend

	if conf.Jaeger != nil {
		backend = conf.Jaeger
	}

	if conf.Zipkin != nil {
		if backend != nil {
			log.WithoutContext().Error("Multiple tracing backend are not supported: cannot create Zipkin backend.")
		} else {
			backend = conf.Zipkin
		}
	}

	if conf.Datadog != nil {
		if backend != nil {
			log.WithoutContext().Error("Multiple tracing backend are not supported: cannot create Datadog backend.")
		} else {
			backend = conf.Datadog
		}
	}

	if conf.Instana != nil {
		if backend != nil {
			log.WithoutContext().Error("Multiple tracing backend are not supported: cannot create Instana backend.")
		} else {
			backend = conf.Instana
		}
	}

	if conf.Haystack != nil {
		if backend != nil {
			log.WithoutContext().Error("Multiple tracing backend are not supported: cannot create Haystack backend.")
		} else {
			backend = conf.Haystack
		}
	}

	if backend == nil {
		log.WithoutContext().Debug("Could not initialize tracing, use Jaeger by default")
		backend := &jaeger.Config{}
		backend.SetDefaults()
	}

	return backend
}

// NewServer returns an initialized Server.
func NewServer(staticConfiguration static.Configuration, provider provider.Provider, entryPoints TCPEntryPoints, tlsManager *tls.Manager) *Server {
	server := &Server{}

	if staticConfiguration.API != nil {
		server.api = api.NewBuilder(staticConfiguration)
	}

	if staticConfiguration.Providers != nil && staticConfiguration.Providers.Rest != nil {
		server.restHandler = staticConfiguration.Providers.Rest.Handler()
	}

	server.provider = provider
	server.entryPointsTCP = entryPoints
	server.configurationChan = make(chan dynamic.Message, 100)
	server.configurationValidatedChan = make(chan dynamic.Message, 100)
	server.signals = make(chan os.Signal, 1)
	server.stopChan = make(chan bool, 1)
	server.configureSignals()
	currentConfigurations := make(dynamic.Configurations)
	server.currentConfigurations.Set(currentConfigurations)
	server.providerConfigUpdateMap = make(map[string]chan dynamic.Message)
	server.tlsManager = tlsManager

	if staticConfiguration.Providers != nil {
		server.providersThrottleDuration = time.Duration(staticConfiguration.Providers.ProvidersThrottleDuration)
	}

	transport, err := createHTTPTransport(staticConfiguration.ServersTransport)
	if err != nil {
		log.WithoutContext().Errorf("Could not configure HTTP Transport, fallbacking on default transport: %v", err)
		server.defaultRoundTripper = http.DefaultTransport
	} else {
		server.defaultRoundTripper = transport
	}

	server.routinesPool = safe.NewPool(context.Background())

	if staticConfiguration.Tracing != nil {
		tracingBackend := setupTracing(staticConfiguration.Tracing)
		if tracingBackend != nil {
			server.tracer, err = tracing.NewTracing(staticConfiguration.Tracing.ServiceName, staticConfiguration.Tracing.SpanNameLimit, tracingBackend)
			if err != nil {
				log.WithoutContext().Warnf("Unable to create tracer: %v", err)
			}
		}
	}

	server.requestDecorator = requestdecorator.New(staticConfiguration.HostResolver)

	server.metricsRegistry = registerMetricClients(staticConfiguration.Metrics)

	if staticConfiguration.AccessLog != nil {
		var err error
		server.accessLoggerMiddleware, err = accesslog.NewHandler(staticConfiguration.AccessLog)
		if err != nil {
			log.WithoutContext().Warnf("Unable to create access logger : %v", err)
		}
	}
	return server
}

// Start starts the server and Stop/Close it when context is Done
func (s *Server) Start(ctx context.Context) {
	go func() {
		defer s.Close()
		<-ctx.Done()
		logger := log.FromContext(ctx)
		logger.Info("I have to go...")
		logger.Info("Stopping server gracefully")
		s.Stop()
	}()

	s.startTCPServers()
	s.routinesPool.Go(func(stop chan bool) {
		s.listenProviders(stop)
	})
	s.routinesPool.Go(func(stop chan bool) {
		s.listenConfigurations(stop)
	})
	s.startProvider()
	s.routinesPool.Go(func(stop chan bool) {
		s.listenSignals(stop)
	})
}

// Wait blocks until server is shutted down.
func (s *Server) Wait() {
	<-s.stopChan
}

// Stop stops the server
func (s *Server) Stop() {
	defer log.WithoutContext().Info("Server stopped")

	var wg sync.WaitGroup
	for epn, ep := range s.entryPointsTCP {
		wg.Add(1)
		go func(entryPointName string, entryPoint *TCPEntryPoint) {
			ctx := log.With(context.Background(), log.Str(log.EntryPointName, entryPointName))
			defer wg.Done()

			entryPoint.Shutdown(ctx)

			log.FromContext(ctx).Debugf("Entry point %s closed", entryPointName)
		}(epn, ep)
	}
	wg.Wait()
	s.stopChan <- true
}

// Close destroys the server
func (s *Server) Close() {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	go func(ctx context.Context) {
		<-ctx.Done()
		if ctx.Err() == context.Canceled {
			return
		} else if ctx.Err() == context.DeadlineExceeded {
			panic("Timeout while stopping traefik, killing instance ✝")
		}
	}(ctx)

	stopMetricsClients()
	s.routinesPool.Cleanup()
	close(s.configurationChan)
	close(s.configurationValidatedChan)
	signal.Stop(s.signals)
	close(s.signals)
	close(s.stopChan)

	if s.accessLoggerMiddleware != nil {
		if err := s.accessLoggerMiddleware.Close(); err != nil {
			log.WithoutContext().Errorf("Could not close the access log file: %s", err)
		}
	}

	if s.tracer != nil {
		s.tracer.Close()
	}

	cancel()
}

func (s *Server) startTCPServers() {
	// Use an empty configuration in order to initialize the default handlers with internal routes
	routers := s.loadConfigurationTCP(dynamic.Configurations{})
	for entryPointName, router := range routers {
		s.entryPointsTCP[entryPointName].switchRouter(router)
	}

	for entryPointName, serverEntryPoint := range s.entryPointsTCP {
		ctx := log.With(context.Background(), log.Str(log.EntryPointName, entryPointName))
		go serverEntryPoint.startTCP(ctx)
	}
}

func (s *Server) listenProviders(stop chan bool) {
	for {
		select {
		case <-stop:
			return
		case configMsg, ok := <-s.configurationChan:
			if !ok {
				return
			}
			if configMsg.Configuration != nil {
				s.preLoadConfiguration(configMsg)
			} else {
				log.WithoutContext().WithField(log.ProviderName, configMsg.ProviderName).
					Debug("Received nil configuration from provider, skipping.")
			}
		}
	}
}

// AddListener adds a new listener function used when new configuration is provided
func (s *Server) AddListener(listener func(dynamic.Configuration)) {
	if s.configurationListeners == nil {
		s.configurationListeners = make([]func(dynamic.Configuration), 0)
	}
	s.configurationListeners = append(s.configurationListeners, listener)
}

func (s *Server) startProvider() {
	logger := log.WithoutContext()

	jsonConf, err := json.Marshal(s.provider)
	if err != nil {
		logger.Debugf("Unable to marshal provider configuration %T: %v", s.provider, err)
	}

	logger.Infof("Starting provider %T %s", s.provider, jsonConf)
	currentProvider := s.provider

	safe.Go(func() {
		err := currentProvider.Provide(s.configurationChan, s.routinesPool)
		if err != nil {
			logger.Errorf("Error starting provider %T: %s", s.provider, err)
		}
	})
}

func registerMetricClients(metricsConfig *types.Metrics) metrics.Registry {
	if metricsConfig == nil {
		return metrics.NewVoidRegistry()
	}

	var registries []metrics.Registry

	if metricsConfig.Prometheus != nil {
		ctx := log.With(context.Background(), log.Str(log.MetricsProviderName, "prometheus"))
		prometheusRegister := metrics.RegisterPrometheus(ctx, metricsConfig.Prometheus)
		if prometheusRegister != nil {
			registries = append(registries, prometheusRegister)
			log.FromContext(ctx).Debug("Configured Prometheus metrics")
		}
	}

	if metricsConfig.Datadog != nil {
		ctx := log.With(context.Background(), log.Str(log.MetricsProviderName, "datadog"))
		registries = append(registries, metrics.RegisterDatadog(ctx, metricsConfig.Datadog))
		log.FromContext(ctx).Debugf("Configured Datadog metrics: pushing to %s once every %s",
			metricsConfig.Datadog.Address, metricsConfig.Datadog.PushInterval)
	}

	if metricsConfig.StatsD != nil {
		ctx := log.With(context.Background(), log.Str(log.MetricsProviderName, "statsd"))
		registries = append(registries, metrics.RegisterStatsd(ctx, metricsConfig.StatsD))
		log.FromContext(ctx).Debugf("Configured StatsD metrics: pushing to %s once every %s",
			metricsConfig.StatsD.Address, metricsConfig.StatsD.PushInterval)
	}

	if metricsConfig.InfluxDB != nil {
		ctx := log.With(context.Background(), log.Str(log.MetricsProviderName, "influxdb"))
		registries = append(registries, metrics.RegisterInfluxDB(ctx, metricsConfig.InfluxDB))
		log.FromContext(ctx).Debugf("Configured InfluxDB metrics: pushing to %s once every %s",
			metricsConfig.InfluxDB.Address, metricsConfig.InfluxDB.PushInterval)
	}

	return metrics.NewMultiRegistry(registries)
}

func stopMetricsClients() {
	metrics.StopDatadog()
	metrics.StopStatsd()
	metrics.StopInfluxDB()
}
<commit_msg>fix: default tracing backend.<commit_after>package server

import (
	"context"
	"encoding/json"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"time"

	"github.com/containous/traefik/v2/pkg/api"
	"github.com/containous/traefik/v2/pkg/config/dynamic"
	"github.com/containous/traefik/v2/pkg/config/runtime"
	"github.com/containous/traefik/v2/pkg/config/static"
	"github.com/containous/traefik/v2/pkg/log"
	"github.com/containous/traefik/v2/pkg/metrics"
	"github.com/containous/traefik/v2/pkg/middlewares/accesslog"
	"github.com/containous/traefik/v2/pkg/middlewares/requestdecorator"
	"github.com/containous/traefik/v2/pkg/provider"
	"github.com/containous/traefik/v2/pkg/safe"
	"github.com/containous/traefik/v2/pkg/tls"
	"github.com/containous/traefik/v2/pkg/tracing"
	"github.com/containous/traefik/v2/pkg/tracing/jaeger"
	"github.com/containous/traefik/v2/pkg/types"
)

// Server is the reverse-proxy/load-balancer engine
type Server struct {
	entryPointsTCP             TCPEntryPoints
	configurationChan          chan dynamic.Message
	configurationValidatedChan chan dynamic.Message
	signals                    chan os.Signal
	stopChan                   chan bool
	currentConfigurations      safe.Safe
	providerConfigUpdateMap    map[string]chan dynamic.Message
	accessLoggerMiddleware     *accesslog.Handler
	tracer                     *tracing.Tracing
	routinesPool               *safe.Pool
	defaultRoundTripper        http.RoundTripper
	metricsRegistry            metrics.Registry
	provider                   provider.Provider
	configurationListeners     []func(dynamic.Configuration)
	requestDecorator           *requestdecorator.RequestDecorator
	providersThrottleDuration  time.Duration
	tlsManager                 *tls.Manager
	api                        func(configuration *runtime.Configuration) http.Handler
	restHandler                http.Handler
}

// RouteAppenderFactory the route appender factory interface
type RouteAppenderFactory interface {
	NewAppender(ctx context.Context, runtimeConfiguration *runtime.Configuration) types.RouteAppender
}

func setupTracing(conf *static.Tracing) tracing.Backend {
	var backend tracing.Backend

	if conf.Jaeger != nil {
		backend = conf.Jaeger
	}

	if conf.Zipkin != nil {
		if backend != nil {
			log.WithoutContext().Error("Multiple tracing backend are not supported: cannot create Zipkin backend.")
		} else {
			backend = conf.Zipkin
		}
	}

	if conf.Datadog != nil {
		if backend != nil {
			log.WithoutContext().Error("Multiple tracing backend are not supported: cannot create Datadog backend.")
		} else {
			backend = conf.Datadog
		}
	}

	if conf.Instana != nil {
		if backend != nil {
			log.WithoutContext().Error("Multiple tracing backend are not supported: cannot create Instana backend.")
		} else {
			backend = conf.Instana
		}
	}

	if conf.Haystack != nil {
		if backend != nil {
			log.WithoutContext().Error("Multiple tracing backend are not supported: cannot create Haystack backend.")
		} else {
			backend = conf.Haystack
		}
	}

	if backend == nil {
		log.WithoutContext().Debug("Could not initialize tracing, use Jaeger by default")
		bcd := &jaeger.Config{}
		bcd.SetDefaults()
		backend = bcd
	}

	return backend
}

// NewServer returns an initialized Server.
func NewServer(staticConfiguration static.Configuration, provider provider.Provider, entryPoints TCPEntryPoints, tlsManager *tls.Manager) *Server {
	server := &Server{}

	if staticConfiguration.API != nil {
		server.api = api.NewBuilder(staticConfiguration)
	}

	if staticConfiguration.Providers != nil && staticConfiguration.Providers.Rest != nil {
		server.restHandler = staticConfiguration.Providers.Rest.Handler()
	}

	server.provider = provider
	server.entryPointsTCP = entryPoints
	server.configurationChan = make(chan dynamic.Message, 100)
	server.configurationValidatedChan = make(chan dynamic.Message, 100)
	server.signals = make(chan os.Signal, 1)
	server.stopChan = make(chan bool, 1)
	server.configureSignals()
	currentConfigurations := make(dynamic.Configurations)
	server.currentConfigurations.Set(currentConfigurations)
	server.providerConfigUpdateMap = make(map[string]chan dynamic.Message)
	server.tlsManager = tlsManager

	if staticConfiguration.Providers != nil {
		server.providersThrottleDuration = time.Duration(staticConfiguration.Providers.ProvidersThrottleDuration)
	}

	transport, err := createHTTPTransport(staticConfiguration.ServersTransport)
	if err != nil {
		log.WithoutContext().Errorf("Could not configure HTTP Transport, fallbacking on default transport: %v", err)
		server.defaultRoundTripper = http.DefaultTransport
	} else {
		server.defaultRoundTripper = transport
	}

	server.routinesPool = safe.NewPool(context.Background())

	if staticConfiguration.Tracing != nil {
		tracingBackend := setupTracing(staticConfiguration.Tracing)
		if tracingBackend != nil {
			server.tracer, err = tracing.NewTracing(staticConfiguration.Tracing.ServiceName, staticConfiguration.Tracing.SpanNameLimit, tracingBackend)
			if err != nil {
				log.WithoutContext().Warnf("Unable to create tracer: %v", err)
			}
		}
	}

	server.requestDecorator = requestdecorator.New(staticConfiguration.HostResolver)

	server.metricsRegistry = registerMetricClients(staticConfiguration.Metrics)

	if staticConfiguration.AccessLog != nil {
		var err error
		server.accessLoggerMiddleware, err = accesslog.NewHandler(staticConfiguration.AccessLog)
		if err != nil {
			log.WithoutContext().Warnf("Unable to create access logger : %v", err)
		}
	}
	return server
}

// Start starts the server and Stop/Close it when context is Done
func (s *Server) Start(ctx context.Context) {
	go func() {
		defer s.Close()
		<-ctx.Done()
		logger := log.FromContext(ctx)
		logger.Info("I have to go...")
		logger.Info("Stopping server gracefully")
		s.Stop()
	}()

	s.startTCPServers()
	s.routinesPool.Go(func(stop chan bool) {
		s.listenProviders(stop)
	})
	s.routinesPool.Go(func(stop chan bool) {
		s.listenConfigurations(stop)
	})
	s.startProvider()
	s.routinesPool.Go(func(stop chan bool) {
		s.listenSignals(stop)
	})
}

// Wait blocks until server is shutted down.
func (s *Server) Wait() {
	<-s.stopChan
}

// Stop stops the server
func (s *Server) Stop() {
	defer log.WithoutContext().Info("Server stopped")

	var wg sync.WaitGroup
	for epn, ep := range s.entryPointsTCP {
		wg.Add(1)
		go func(entryPointName string, entryPoint *TCPEntryPoint) {
			ctx := log.With(context.Background(), log.Str(log.EntryPointName, entryPointName))
			defer wg.Done()

			entryPoint.Shutdown(ctx)

			log.FromContext(ctx).Debugf("Entry point %s closed", entryPointName)
		}(epn, ep)
	}
	wg.Wait()
	s.stopChan <- true
}

// Close destroys the server
func (s *Server) Close() {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	go func(ctx context.Context) {
		<-ctx.Done()
		if ctx.Err() == context.Canceled {
			return
		} else if ctx.Err() == context.DeadlineExceeded {
			panic("Timeout while stopping traefik, killing instance ✝")
		}
	}(ctx)

	stopMetricsClients()
	s.routinesPool.Cleanup()
	close(s.configurationChan)
	close(s.configurationValidatedChan)
	signal.Stop(s.signals)
	close(s.signals)
	close(s.stopChan)

	if s.accessLoggerMiddleware != nil {
		if err := s.accessLoggerMiddleware.Close(); err != nil {
			log.WithoutContext().Errorf("Could not close the access log file: %s", err)
		}
	}

	if s.tracer != nil {
		s.tracer.Close()
	}

	cancel()
}

func (s *Server) startTCPServers() {
	// Use an empty configuration in order to initialize the default handlers with internal routes
	routers := s.loadConfigurationTCP(dynamic.Configurations{})
	for entryPointName, router := range routers {
		s.entryPointsTCP[entryPointName].switchRouter(router)
	}

	for entryPointName, serverEntryPoint := range s.entryPointsTCP {
		ctx := log.With(context.Background(), log.Str(log.EntryPointName, entryPointName))
		go serverEntryPoint.startTCP(ctx)
	}
}

func (s *Server) listenProviders(stop chan bool) {
	for {
		select {
		case <-stop:
			return
		case configMsg, ok := <-s.configurationChan:
			if !ok {
				return
			}
			if configMsg.Configuration != nil {
				s.preLoadConfiguration(configMsg)
			} else {
				log.WithoutContext().WithField(log.ProviderName, configMsg.ProviderName).
					Debug("Received nil configuration from provider, skipping.")
			}
		}
	}
}

// AddListener adds a new listener function used when new configuration is provided
func (s *Server) AddListener(listener func(dynamic.Configuration)) {
	if s.configurationListeners == nil {
		s.configurationListeners = make([]func(dynamic.Configuration), 0)
	}
	s.configurationListeners = append(s.configurationListeners, listener)
}

func (s *Server) startProvider() {
	logger := log.WithoutContext()

	jsonConf, err := json.Marshal(s.provider)
	if err != nil {
		logger.Debugf("Unable to marshal provider configuration %T: %v", s.provider, err)
	}

	logger.Infof("Starting provider %T %s", s.provider, jsonConf)
	currentProvider := s.provider

	safe.Go(func() {
		err := currentProvider.Provide(s.configurationChan, s.routinesPool)
		if err != nil {
			logger.Errorf("Error starting provider %T: %s", s.provider, err)
		}
	})
}

func registerMetricClients(metricsConfig *types.Metrics) metrics.Registry {
	if metricsConfig == nil {
		return metrics.NewVoidRegistry()
	}

	var registries []metrics.Registry

	if metricsConfig.Prometheus != nil {
		ctx := log.With(context.Background(), log.Str(log.MetricsProviderName, "prometheus"))
		prometheusRegister := metrics.RegisterPrometheus(ctx, metricsConfig.Prometheus)
		if prometheusRegister != nil {
			registries = append(registries, prometheusRegister)
			log.FromContext(ctx).Debug("Configured Prometheus metrics")
		}
	}

	if metricsConfig.Datadog != nil {
		ctx := log.With(context.Background(), log.Str(log.MetricsProviderName, "datadog"))
		registries = append(registries, metrics.RegisterDatadog(ctx, metricsConfig.Datadog))
		log.FromContext(ctx).Debugf("Configured Datadog metrics: pushing to %s once every %s",
			metricsConfig.Datadog.Address, metricsConfig.Datadog.PushInterval)
	}

	if metricsConfig.StatsD != nil {
		ctx := log.With(context.Background(), log.Str(log.MetricsProviderName, "statsd"))
		registries = append(registries, metrics.RegisterStatsd(ctx, metricsConfig.StatsD))
		log.FromContext(ctx).Debugf("Configured StatsD metrics: pushing to %s once every %s",
			metricsConfig.StatsD.Address, metricsConfig.StatsD.PushInterval)
	}

	if metricsConfig.InfluxDB != nil {
		ctx := log.With(context.Background(), log.Str(log.MetricsProviderName, "influxdb"))
		registries = append(registries, metrics.RegisterInfluxDB(ctx, metricsConfig.InfluxDB))
		log.FromContext(ctx).Debugf("Configured InfluxDB metrics: pushing to %s once every %s",
			metricsConfig.InfluxDB.Address, metricsConfig.InfluxDB.PushInterval)
	}

	return metrics.NewMultiRegistry(registries)
}

func stopMetricsClients() {
	metrics.StopDatadog()
	metrics.StopStatsd()
	metrics.StopInfluxDB()
}
<|endoftext|>


<commit_before>package sharing

import (
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/crypto"
	"github.com/cozy/cozy-stack/pkg/instance"
	"github.com/cozy/cozy-stack/pkg/vfs"
)

// MakeXorKey generates a key for transforming the file identifiers
func MakeXorKey() []byte {
	random := crypto.GenerateRandomBytes(8)
	result := make([]byte, 2*len(random))
	for i, val := range random {
		result[2*i] = val & 0xf
		result[2*i+1] = val >> 4
	}
	return result
}

// XorID transforms the identifier of a file to a new identifier, in a
// reversible way: it makes a XOR on the hexadecimal characters
func XorID(id string, key []byte) string {
	l := len(key)
	buf := []byte(id)
	for i, c := range buf {
		switch {
		case '0' <= c && c <= '9':
			c = (c - '0') ^ key[i%l]
		case 'a' <= c && c <= 'f':
			c = (c - 'a' + 10) ^ key[i%l]
		case 'A' <= c && c <= 'F':
			c = (c - 'A' + 10) ^ key[i%l]
		default:
			continue
		}
		if c < 10 {
			buf[i] = c + '0'
		} else {
			buf[i] = (c - 10) + 'a'
		}
	}
	return string(buf)
}

// EnsureSharedWithMeDir returns the shared-with-me directory, and create it if
// it doesn't exist
func EnsureSharedWithMeDir(inst *instance.Instance) (*vfs.DirDoc, error) {
	fs := inst.VFS()
	dir, _, err := fs.DirOrFileByID(consts.SharedWithMeDirID)
	if err != nil {
		return nil, err
	}

	if dir == nil {
		name := inst.Translate("Tree Shared with me")
		dir, err = vfs.NewDirDocWithPath(name, consts.SharedWithMeDirID, "/", nil)
		if err != nil {
			return nil, err
		}
		if err = fs.CreateDir(dir); err != nil {
			return nil, err
		}
		return dir, nil
	}

	if dir.RestorePath != "" {
		_, err = vfs.RestoreDir(fs, dir)
		if err != nil {
			return nil, err
		}
		children, err := fs.DirBatch(dir, &couchdb.SkipCursor{})
		if err != nil {
			return nil, err
		}
		for _, child := range children {
			d, f := child.Refine()
			if d != nil {
				_, err = vfs.TrashDir(fs, d)
			} else {
				_, err = vfs.TrashFile(fs, f)
			}
			if err != nil {
				return nil, err
			}
		}
	}

	return dir, nil
}

// CreateDirForSharing creates the directory where files for this sharing will
// be put. This directory will be initially inside the Shared with me folder.
func (s *Sharing) CreateDirForSharing(inst *instance.Instance, rule *Rule) error {
	parent, err := EnsureSharedWithMeDir(inst)
	if err != nil {
		return err
	}
	fs := inst.VFS()
	dir, err := vfs.NewDirDoc(fs, rule.Title, parent.DocID, []string{"from-sharing-" + s.SID})
	dir.AddReferencedBy(couchdb.DocReference{
		ID:   s.SID,
		Type: consts.Sharings,
	})
	if rule.Selector == "" || rule.Selector == "id" || rule.Selector == "_id" {
		dir.DocID = rule.Values[0]
	}
	return fs.CreateDir(dir)
}
<commit_msg>Fix the creation of Shared with me dir<commit_after>package sharing

import (
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/crypto"
	"github.com/cozy/cozy-stack/pkg/instance"
	"github.com/cozy/cozy-stack/pkg/vfs"
)

// MakeXorKey generates a key for transforming the file identifiers
func MakeXorKey() []byte {
	random := crypto.GenerateRandomBytes(8)
	result := make([]byte, 2*len(random))
	for i, val := range random {
		result[2*i] = val & 0xf
		result[2*i+1] = val >> 4
	}
	return result
}

// XorID transforms the identifier of a file to a new identifier, in a
// reversible way: it makes a XOR on the hexadecimal characters
func XorID(id string, key []byte) string {
	l := len(key)
	buf := []byte(id)
	for i, c := range buf {
		switch {
		case '0' <= c && c <= '9':
			c = (c - '0') ^ key[i%l]
		case 'a' <= c && c <= 'f':
			c = (c - 'a' + 10) ^ key[i%l]
		case 'A' <= c && c <= 'F':
			c = (c - 'A' + 10) ^ key[i%l]
		default:
			continue
		}
		if c < 10 {
			buf[i] = c + '0'
		} else {
			buf[i] = (c - 10) + 'a'
		}
	}
	return string(buf)
}

// EnsureSharedWithMeDir returns the shared-with-me directory, and create it if
// it doesn't exist
func EnsureSharedWithMeDir(inst *instance.Instance) (*vfs.DirDoc, error) {
	fs := inst.VFS()
	dir, _, err := fs.DirOrFileByID(consts.SharedWithMeDirID)
	if err != nil {
		return nil, err
	}

	if dir == nil {
		name := inst.Translate("Tree Shared with me")
		dir, err = vfs.NewDirDocWithPath(name, consts.RootDirID, "/", nil)
		dir.DocID = consts.SharedWithMeDirID
		if err != nil {
			return nil, err
		}
		if err = fs.CreateDir(dir); err != nil {
			return nil, err
		}
		return dir, nil
	}

	if dir.RestorePath != "" {
		_, err = vfs.RestoreDir(fs, dir)
		if err != nil {
			return nil, err
		}
		children, err := fs.DirBatch(dir, &couchdb.SkipCursor{})
		if err != nil {
			return nil, err
		}
		for _, child := range children {
			d, f := child.Refine()
			if d != nil {
				_, err = vfs.TrashDir(fs, d)
			} else {
				_, err = vfs.TrashFile(fs, f)
			}
			if err != nil {
				return nil, err
			}
		}
	}

	return dir, nil
}

// CreateDirForSharing creates the directory where files for this sharing will
// be put. This directory will be initially inside the Shared with me folder.
func (s *Sharing) CreateDirForSharing(inst *instance.Instance, rule *Rule) error {
	parent, err := EnsureSharedWithMeDir(inst)
	if err != nil {
		return err
	}
	fs := inst.VFS()
	dir, err := vfs.NewDirDocWithParent(rule.Title, parent, []string{"from-sharing-" + s.SID})
	if rule.Selector == "" || rule.Selector == "id" || rule.Selector == "_id" {
		dir.DocID = rule.Values[0]
	}
	dir.AddReferencedBy(couchdb.DocReference{
		ID:   s.SID,
		Type: consts.Sharings,
	})
	return fs.CreateDir(dir)
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google LLC

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 util

import (
	"archive/tar"
	"compress/bzip2"
	"compress/gzip"
	pkgutil "github.com/GoogleCloudPlatform/container-diff/pkg/util"
	"github.com/docker/docker/pkg/archive"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"io"
	"io/ioutil"
	"os"
	"syscall"
)

var hardlinks = make(map[uint64]string)

// AddToTar adds the file i to tar w at path p
func AddToTar(p string, i os.FileInfo, w *tar.Writer) error {
	linkDst := ""
	if i.Mode()&os.ModeSymlink != 0 {
		var err error
		linkDst, err = os.Readlink(p)
		if err != nil {
			return err
		}
	}
	hdr, err := tar.FileInfoHeader(i, linkDst)
	if err != nil {
		return err
	}
	hdr.Name = p

	hardlink, linkDst := checkHardlink(p, i)
	if hardlink {
		hdr.Linkname = linkDst
		hdr.Typeflag = tar.TypeLink
		hdr.Size = 0
	}
	if err := w.WriteHeader(hdr); err != nil {
		return err
	}
	if !(i.Mode().IsRegular()) || hardlink {
		return nil
	}
	r, err := os.Open(p)
	if err != nil {
		return err
	}
	defer r.Close()
	if _, err := io.Copy(w, r); err != nil {
		return err
	}
	return nil
}

// Returns true if path is hardlink, and the link destination
func checkHardlink(p string, i os.FileInfo) (bool, string) {
	hardlink := false
	linkDst := ""
	if sys := i.Sys(); sys != nil {
		if stat, ok := sys.(*syscall.Stat_t); ok {
			nlinks := stat.Nlink
			if nlinks > 1 {
				inode := stat.Ino
				if original, exists := hardlinks[inode]; exists && original != p {
					hardlink = true
					logrus.Debugf("%s inode exists in hardlinks map, linking to %s", p, original)
					linkDst = original
				} else {
					hardlinks[inode] = p
				}
			}
		}
	}
	return hardlink, linkDst
}

// UnpackLocalTarArchive unpacks the tar archive at path to the directory dest
// Returns true if the path was acutally unpacked
func UnpackLocalTarArchive(path, dest string) error {
	// First, we need to check if the path is a local tar archive
	if compressed, compressionLevel := fileIsCompressedTar(path); compressed {
		file, err := os.Open(path)
		if err != nil {
			return err
		}
		defer file.Close()
		if compressionLevel == archive.Gzip {
			return UnpackCompressedTar(path, dest)
		} else if compressionLevel == archive.Bzip2 {
			bzr := bzip2.NewReader(file)
			return pkgutil.UnTar(bzr, dest, nil)
		}
	}
	if fileIsUncompressedTar(path) {
		file, err := os.Open(path)
		if err != nil {
			return err
		}
		defer file.Close()
		return pkgutil.UnTar(file, dest, nil)
	}
	return errors.New("path does not lead to local tar archive")
}

//IsFileLocalTarArchive returns true if the file is a local tar archive
func IsFileLocalTarArchive(src string) bool {
	compressed, _ := fileIsCompressedTar(src)
	uncompressed := fileIsUncompressedTar(src)
	return compressed || uncompressed
}

func fileIsCompressedTar(src string) (bool, archive.Compression) {
	r, err := os.Open(src)
	if err != nil {
		return false, -1
	}
	defer r.Close()
	buf, err := ioutil.ReadAll(r)
	if err != nil {
		return false, -1
	}
	compressionLevel := archive.DetectCompression(buf)
	return (compressionLevel > 0), compressionLevel
}

func fileIsUncompressedTar(src string) bool {
	r, err := os.Open(src)
	defer r.Close()
	if err != nil {
		return false
	}
	fi, err := os.Lstat(src)
	if err != nil {
		return false
	}
	if fi.Size() == 0 {
		return false
	}
	tr := tar.NewReader(r)
	if tr == nil {
		return false
	}
	for {
		_, err := tr.Next()
		if err != nil {
			return false
		}
		return true
	}
}

// UnpackCompressedTar unpacks the compressed tar at path to dir
func UnpackCompressedTar(path, dir string) error {
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()
	gzr, err := gzip.NewReader(file)
	if err != nil {
		return err
	}
	defer gzr.Close()
	return pkgutil.UnTar(gzr, dir, nil)
}
<commit_msg>Update tar_util.go<commit_after>/*
Copyright 2018 Google LLC

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 util

import (
	"archive/tar"
	"compress/bzip2"
	"compress/gzip"
	pkgutil "github.com/GoogleCloudPlatform/container-diff/pkg/util"
	"github.com/docker/docker/pkg/archive"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"io"
	"io/ioutil"
	"os"
	"syscall"
)

var hardlinks = make(map[uint64]string)

// AddToTar adds the file i to tar w at path p
func AddToTar(p string, i os.FileInfo, w *tar.Writer) error {
	linkDst := ""
	if i.Mode()&os.ModeSymlink != 0 {
		var err error
		linkDst, err = os.Readlink(p)
		if err != nil {
			return err
		}
	}
	hdr, err := tar.FileInfoHeader(i, linkDst)
	if err != nil {
		return err
	}
	hdr.Name = p

	hardlink, linkDst := checkHardlink(p, i)
	if hardlink {
		hdr.Linkname = linkDst
		hdr.Typeflag = tar.TypeLink
		hdr.Size = 0
	}
	if err := w.WriteHeader(hdr); err != nil {
		return err
	}
	if !(i.Mode().IsRegular()) || hardlink {
		return nil
	}
	r, err := os.Open(p)
	if err != nil {
		return err
	}
	defer r.Close()
	if _, err := io.Copy(w, r); err != nil {
		return err
	}
	return nil
}

// Returns true if path is hardlink, and the link destination
func checkHardlink(p string, i os.FileInfo) (bool, string) {
	hardlink := false
	linkDst := ""
	if sys := i.Sys(); sys != nil {
		if stat, ok := sys.(*syscall.Stat_t); ok {
			nlinks := stat.Nlink
			if nlinks > 1 {
				inode := stat.Ino
				if original, exists := hardlinks[inode]; exists && original != p {
					hardlink = true
					logrus.Debugf("%s inode exists in hardlinks map, linking to %s", p, original)
					linkDst = original
				} else {
					hardlinks[inode] = p
				}
			}
		}
	}
	return hardlink, linkDst
}

// UnpackLocalTarArchive unpacks the tar archive at path to the directory dest
// Returns true if the path was actually unpacked
func UnpackLocalTarArchive(path, dest string) error {
	// First, we need to check if the path is a local tar archive
	if compressed, compressionLevel := fileIsCompressedTar(path); compressed {
		file, err := os.Open(path)
		if err != nil {
			return err
		}
		defer file.Close()
		if compressionLevel == archive.Gzip {
			return UnpackCompressedTar(path, dest)
		} else if compressionLevel == archive.Bzip2 {
			bzr := bzip2.NewReader(file)
			return pkgutil.UnTar(bzr, dest, nil)
		}
	}
	if fileIsUncompressedTar(path) {
		file, err := os.Open(path)
		if err != nil {
			return err
		}
		defer file.Close()
		return pkgutil.UnTar(file, dest, nil)
	}
	return errors.New("path does not lead to local tar archive")
}

//IsFileLocalTarArchive returns true if the file is a local tar archive
func IsFileLocalTarArchive(src string) bool {
	compressed, _ := fileIsCompressedTar(src)
	uncompressed := fileIsUncompressedTar(src)
	return compressed || uncompressed
}

func fileIsCompressedTar(src string) (bool, archive.Compression) {
	r, err := os.Open(src)
	if err != nil {
		return false, -1
	}
	defer r.Close()
	buf, err := ioutil.ReadAll(r)
	if err != nil {
		return false, -1
	}
	compressionLevel := archive.DetectCompression(buf)
	return (compressionLevel > 0), compressionLevel
}

func fileIsUncompressedTar(src string) bool {
	r, err := os.Open(src)
	defer r.Close()
	if err != nil {
		return false
	}
	fi, err := os.Lstat(src)
	if err != nil {
		return false
	}
	if fi.Size() == 0 {
		return false
	}
	tr := tar.NewReader(r)
	if tr == nil {
		return false
	}
	for {
		_, err := tr.Next()
		if err != nil {
			return false
		}
		return true
	}
}

// UnpackCompressedTar unpacks the compressed tar at path to dir
func UnpackCompressedTar(path, dir string) error {
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()
	gzr, err := gzip.NewReader(file)
	if err != nil {
		return err
	}
	defer gzr.Close()
	return pkgutil.UnTar(gzr, dir, nil)
}
<|endoftext|>


<commit_before>package main

import (
	"sync"
	"testing"
	"time"
)

func TestPlaySingleFile(t *testing.T) {
	wg := sync.WaitGroup{}
	player = Player{WG: &wg}
	err := player.init()
	if err != nil {
		t.Fatalf(err.Error())
	}
	wg.Add(1)
	start := time.Now()
	player.playSingleFile("test_sounds/beep9.mp3")
	wg.Wait()
	duration := time.Since(start)
	if duration.Seconds() < 1 {
		t.Errorf("Expected to play for at least\n---\n%d\n---\nbut played\n---\n%f\n---\n", 1,
			duration.Seconds())
	}
	player.clear()
}
<commit_msg>add test for init<commit_after>package main

import (
	"sync"
	"testing"
	"time"
)

func TestInit(t *testing.T) {
	player = Player{}
	err := player.init()
	if err != nil {
		t.Fatalf(err.Error())
	}
	player.clear()
}

func TestPlaySingleFile(t *testing.T) {
	wg := sync.WaitGroup{}
	player = Player{WG: &wg}
	err := player.init()
	if err != nil {
		t.Fatalf(err.Error())
	}
	wg.Add(1)
	start := time.Now()
	player.playSingleFile("test_sounds/beep9.mp3")
	wg.Wait()
	duration := time.Since(start)
	if duration.Seconds() < 1 {
		t.Errorf("Expected to play for at least\n---\n%d\n---\nbut played\n---\n%f\n---\n", 1,
			duration.Seconds())
	}
	player.clear()
}
<|endoftext|>


<commit_before>package base

import "math/rand"

type Candidate interface{}

func ShuffleCandidates(slice []Candidate, rng *rand.Rand) {
	for i := len(slice) - 1; i > 0; i-- {
		j := rng.Intn(i + 1)
		slice[i], slice[j] = slice[j], slice[i]
	}
}
<commit_msg>For commodity, ShuffleCandidates returns the modified slice<commit_after>package base

import "math/rand"

type Candidate interface{}

// ShuffleCandidates shuffles a slice of candidates.
//
// The original slice is modified, though, for commodity, it is returned.
func ShuffleCandidates(slice []Candidate, rng *rand.Rand) []Candidate {
	for i := len(slice) - 1; i > 0; i-- {
		j := rng.Intn(i + 1)
		slice[i], slice[j] = slice[j], slice[i]
	}
	return slice
}
<|endoftext|>


<commit_before>package base

import (
	"container/list"
	"errors"
	"sync"
)

// An LRU cache of document revision bodies, together with their channel access.
type LRUCache struct {
	cache      map[string]*list.Element // Fast lookup of list element by key
	lruList    *list.List               // List ordered by most recent access (Front is newest)
	capacity   int                      // Max number of entries to cache
	lruLock    sync.Mutex               // For thread-safety
	loaderFunc LRUCacheLoaderFunc
}

type LRUCacheLoaderFunc func(key string) (value interface{}, err error)

// The cache payload data. Stored as the Value of a list Element.
type lruCacheValue struct {
	key          string
	value        interface{}
	err          error      // Error from loaderFunc if it failed
	lruValueLock sync.Mutex // Synchronizes access to this struct
}

// Creates an LRU cache with the given capacity and an optional loader function.
func NewLRUCache(capacity int) (*LRUCache, error) {

	if capacity <= 0 {
		return nil, errors.New("LRU cache capacity must be non-zero")
	}

	return &LRUCache{
		cache:    map[string]*list.Element{},
		lruList:  list.New(),
		capacity: capacity}, nil
}

// Looks up an entry from the cache.
func (lc *LRUCache) Get(key string) (result interface{}, found bool) {
	lc.lruLock.Lock()
	defer lc.lruLock.Unlock()
	if elem, ok := lc.cache[key]; ok {
		lc.lruList.MoveToFront(elem)
		return elem.Value.(*lruCacheValue).value, true
	}
	return result, false
}

// Adds an entry to the cache.
func (lc *LRUCache) Put(key string, value interface{}) {

	// If already present, move to front
	if elem := lc.cache[key]; elem != nil {
		lc.lruList.MoveToFront(elem)
		value = elem.Value.(*lruCacheValue)
		return
	}

	// Not found - add as new
	cacheValue := &lruCacheValue{
		key:   key,
		value: value,
	}
	lc.cache[key] = lc.lruList.PushFront(cacheValue)

	// Purge oldest if over capacity
	for len(lc.cache) > lc.capacity {
		lc.purgeOldest_()
	}
}

func (lc *LRUCache) purgeOldest_() {
	value := lc.lruList.Remove(lc.lruList.Back()).(*lruCacheValue)
	delete(lc.cache, value.key)
}

func (lc *LRUCache) Count() int {
	return len(lc.cache)
}
<commit_msg>Fix attempt for SG Accel #220 (#3744)<commit_after>package base

import (
	"container/list"
	"errors"
	"sync"
)

// An LRU cache of document revision bodies, together with their channel access.
type LRUCache struct {
	cache      map[string]*list.Element // Fast lookup of list element by key
	lruList    *list.List               // List ordered by most recent access (Front is newest)
	capacity   int                      // Max number of entries to cache
	lruLock    sync.Mutex               // For thread-safety
	loaderFunc LRUCacheLoaderFunc
}

type LRUCacheLoaderFunc func(key string) (value interface{}, err error)

// The cache payload data. Stored as the Value of a list Element.
type lruCacheValue struct {
	key          string
	value        interface{}
	err          error      // Error from loaderFunc if it failed
	lruValueLock sync.Mutex // Synchronizes access to this struct
}

// Creates an LRU cache with the given capacity and an optional loader function.
func NewLRUCache(capacity int) (*LRUCache, error) {

	if capacity <= 0 {
		return nil, errors.New("LRU cache capacity must be non-zero")
	}

	return &LRUCache{
		cache:    map[string]*list.Element{},
		lruList:  list.New(),
		capacity: capacity}, nil
}

// Looks up an entry from the cache.
func (lc *LRUCache) Get(key string) (result interface{}, found bool) {
	lc.lruLock.Lock()
	defer lc.lruLock.Unlock()
	if elem, ok := lc.cache[key]; ok {
		lc.lruList.MoveToFront(elem)
		return elem.Value.(*lruCacheValue).value, true
	}
	return result, false
}

// Adds an entry to the cache if it's the first time seen, otherwise just updates the
// position in the LRU cache but ignores the new value, since the entries in the cache
// are treated as immutable.
func (lc *LRUCache) Put(key string, value interface{}) {

	// If already present, move to front
	if elem := lc.cache[key]; elem != nil {
		lc.lruList.MoveToFront(elem)
		value = elem.Value.(*lruCacheValue)  // "return" the value in the cache by updating the value param
		return
	}

	// Not found - add as new
	cacheValue := &lruCacheValue{
		key:   key,
		value: value,
	}
	lc.cache[key] = lc.lruList.PushFront(cacheValue)

	// Purge oldest if over capacity
	for len(lc.cache) > lc.capacity {
		lc.purgeOldest_()
	}
}

func (lc *LRUCache) purgeOldest_() {
	value := lc.lruList.Remove(lc.lruList.Back()).(*lruCacheValue)
	delete(lc.cache, value.key)
}

func (lc *LRUCache) Count() int {
	return len(lc.cache)
}
<|endoftext|>


<commit_before>//
// Copyright (c) 2017 Jon Carlson.  All rights reserved.
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.
//

package main

import (
	"bytes"
	"database/sql"
	"fmt"
	"github.com/joncrlsn/misc"
	"github.com/joncrlsn/pgutil"
	"sort"
	"strconv"
	"strings"
	"text/template"
)

var (
	columnSqlTemplate = initColumnSqlTemplate()
)

// Initializes the Sql template
func initColumnSqlTemplate() *template.Template {
	sql := `
SELECT table_schema
    , {{if eq $.DbSchema "*" }}table_schema || '.' || {{end}}table_name || '.' || column_name  AS compare_name
	, table_name
    , column_name
    , data_type
    , is_nullable
    , column_default
    , character_maximum_length
    , is_identity
    , identity_generation
FROM information_schema.columns
WHERE is_updatable = 'YES'
{{if eq $.DbSchema "*" }}
AND table_schema NOT LIKE 'pg_%'
AND table_schema <> 'information_schema'
{{else}}
AND table_schema = '{{$.DbSchema}}'
{{end}}
ORDER BY compare_name ASC;
`
	t := template.New("ColumnSqlTmpl")
	template.Must(t.Parse(sql))
	return t
}

// ==================================
// Column Rows definition
// ==================================

// ColumnRows is a sortable slice of string maps
type ColumnRows []map[string]string

func (slice ColumnRows) Len() int {
	return len(slice)
}

func (slice ColumnRows) Less(i, j int) bool {
	return slice[i]["compare_name"] < slice[j]["compare_name"]
}

func (slice ColumnRows) Swap(i, j int) {
	slice[i], slice[j] = slice[j], slice[i]
}

// ==================================
// ColumnSchema definition
// (implements Schema -- defined in pgdiff.go)
// ==================================

// ColumnSchema holds a slice of rows from one of the databases as well as
// a reference to the current row of data we're viewing.
type ColumnSchema struct {
	rows   ColumnRows
	rowNum int
	done   bool
}

// get returns the value from the current row for the given key
func (c *ColumnSchema) get(key string) string {
	if c.rowNum >= len(c.rows) {
		return ""
	}
	return c.rows[c.rowNum][key]
}

// NextRow increments the rowNum and tells you whether or not there are more
func (c *ColumnSchema) NextRow() bool {
	if c.rowNum >= len(c.rows)-1 {
		c.done = true
	}
	c.rowNum = c.rowNum + 1
	return !c.done
}

// Compare tells you, in one pass, whether or not the first row matches, is less than, or greater than the second row
func (c *ColumnSchema) Compare(obj interface{}) int {
	c2, ok := obj.(*ColumnSchema)
	if !ok {
		fmt.Println("Error!!!, Compare needs a ColumnSchema instance", c2)
	}

	val := misc.CompareStrings(c.get("compare_name"), c2.get("compare_name"))
	return val
}

// Add prints SQL to add the column
func (c *ColumnSchema) Add() {

	schema := dbInfo2.DbSchema
	if schema == "*" {
		schema = c.get("table_schema")
	}

	// Knowing the version of db2 would eliminate the need for this warning
	if c.get("is_identity") == "YES" {
		fmt.Println("-- WARNING: identity columns are not supported in PostgreSQL versions < 10.")
		fmt.Println("-- Attempting to create identity columns in earlier versions will probably result in errors.")
	}

	if c.get("data_type") == "character varying" {
		maxLength, valid := getMaxLength(c.get("character_maximum_length"))
		if !valid {
			fmt.Printf("ALTER TABLE %s.%s ADD COLUMN %s character varying", schema, c.get("table_name"), c.get("column_name"))
		} else {
			fmt.Printf("ALTER TABLE %s.%s ADD COLUMN %s character varying(%s)", schema, c.get("table_name"), c.get("column_name"), maxLength)
		}
	} else {
		if c.get("data_type") == "ARRAY" {
			fmt.Println("-- Note that adding of array data types are not yet generated properly.")
		}
		fmt.Printf("ALTER TABLE %s.%s ADD COLUMN %s %s", schema, c.get("table_name"), c.get("column_name"), c.get("data_type"))
	}

	if c.get("is_nullable") == "NO" {
		fmt.Printf(" NOT NULL")
	}
	if c.get("column_default") != "null" {
		fmt.Printf(" DEFAULT %s", c.get("column_default"))
	}
	// NOTE: there are more identity column sequence options according to the PostgreSQL
	// CREATE TABLE docs, but these do not appear to be available as of version 10.1
	if c.get("is_identity") == "YES" {
		fmt.Printf(" GENERATED %s AS IDENTITY", c.get("identity_generation"))
	}
	fmt.Printf(";\n")
}

// Drop prints SQL to drop the column
func (c *ColumnSchema) Drop() {
	// if dropping column
	fmt.Printf("ALTER TABLE %s.%s DROP COLUMN IF EXISTS %s;\n", c.get("table_schema"), c.get("table_name"), c.get("column_name"))
}

// Change handles the case where the table and column match, but the details do not
func (c *ColumnSchema) Change(obj interface{}) {
	c2, ok := obj.(*ColumnSchema)
	if !ok {
		fmt.Println("Error!!!, ColumnSchema.Change(obj) needs a ColumnSchema instance", c2)
	}

	// Detect column type change (mostly varchar length, or number size increase)
	// (integer to/from bigint is OK)
	if c.get("data_type") == c2.get("data_type") {
		if c.get("data_type") == "character varying" {
			max1, max1Valid := getMaxLength(c.get("character_maximum_length"))
			max2, max2Valid := getMaxLength(c2.get("character_maximum_length"))
			if !max1Valid && !max2Valid {
				// Leave them alone, they both have undefined max lengths
			} else if (max1Valid || !max2Valid) && (max1 != c2.get("character_maximum_length")) {
				//if !max1Valid {
				//    fmt.Println("-- WARNING: varchar column has no maximum length.  Setting to 1024, which may result in data loss.")
				//}
				max1Int, err1 := strconv.Atoi(max1)
				check("converting string to int", err1)
				max2Int, err2 := strconv.Atoi(max2)
				check("converting string to int", err2)
				if max1Int < max2Int {
					fmt.Println("-- WARNING: The next statement will shorten a character varying column, which may result in data loss.")
				}
				fmt.Printf("-- max1Valid: %v  max2Valid: %v \n", max1Valid, max2Valid)
				fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s TYPE character varying(%s);\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), max1)
			}
		}
	}

	// Code and test a column change from integer to bigint
	if c.get("data_type") != c2.get("data_type") {
		fmt.Printf("-- WARNING: This type change may not work well: (%s to %s).\n", c2.get("data_type"), c.get("data_type"))
		if strings.HasPrefix(c.get("data_type"), "character") {
			max1, max1Valid := getMaxLength(c.get("character_maximum_length"))
			if !max1Valid {
				fmt.Println("-- WARNING: varchar column has no maximum length.  Setting to 1024")
			}
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s TYPE %s(%s);\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), c.get("data_type"), max1)
		} else {
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s TYPE %s;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), c.get("data_type"))
		}
	}

	// Detect column default change (or added, dropped)
	if c.get("column_default") == "null" {
		if c.get("column_default") != "null" {
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s DROP DEFAULT;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"))
		}
	} else if c.get("column_default") != c2.get("column_default") {
		fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s SET DEFAULT %s;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), c.get("column_default"))
	}

	// Detect identity column change
	// Save result to variable instead of printing because order for adding/removing
	// is_nullable affects identity columns
	var identitySql string
	if c.get("is_identity") != c2.get("is_identity") {
		// Knowing the version of db2 would eliminate the need for this warning
		fmt.Println("-- WARNING: identity columns are not supported in PostgreSQL versions < 10.")
		fmt.Println("-- Attempting to create identity columns in earlier versions will probably result in errors.")
		if c.get("is_identity") == "YES" {
			identitySql = fmt.Sprintf("ALTER TABLE \"%s\".\"%s\" ALTER COLUMN \"%s\" ADD GENERATED %s AS IDENTITY;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), c.get("identity_generation"))
		} else {
			identitySql = fmt.Sprintf("ALTER TABLE \"%s\".\"%s\" ALTER COLUMN \"%s\" DROP IDENTITY;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"))
		}
	}

	// Detect not-null and nullable change
	if c.get("is_nullable") != c2.get("is_nullable") {
		if c.get("is_nullable") == "YES" {
			if identitySql != "" {
				fmt.Printf(identitySql)
			}
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s DROP NOT NULL;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"))
		} else {
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s SET NOT NULL;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"))
			if identitySql != "" {
				fmt.Printf(identitySql)
			}
		}
	} else {
		if identitySql != "" {
			fmt.Printf(identitySql)
		}
	}
}

// ==================================
// Standalone Functions
// ==================================

// compareColumns outputs SQL to make the columns match between two databases or schemas
func compareColumns(conn1 *sql.DB, conn2 *sql.DB) {

	buf1 := new(bytes.Buffer)
	columnSqlTemplate.Execute(buf1, dbInfo1)

	buf2 := new(bytes.Buffer)
	columnSqlTemplate.Execute(buf2, dbInfo2)

	rowChan1, _ := pgutil.QueryStrings(conn1, buf1.String())
	rowChan2, _ := pgutil.QueryStrings(conn2, buf2.String())

	//rows1 := make([]map[string]string, 500)
	rows1 := make(ColumnRows, 0)
	for row := range rowChan1 {
		rows1 = append(rows1, row)
	}
	sort.Sort(rows1)

	//rows2 := make([]map[string]string, 500)
	rows2 := make(ColumnRows, 0)
	for row := range rowChan2 {
		rows2 = append(rows2, row)
	}
	sort.Sort(&rows2)

	// We have to explicitly type this as Schema here for some unknown reason
	var schema1 Schema = &ColumnSchema{rows: rows1, rowNum: -1}
	var schema2 Schema = &ColumnSchema{rows: rows2, rowNum: -1}

	// Compare the columns
	doDiff(schema1, schema2)
}

// getMaxLength returns the maximum length and whether or not it is valid
func getMaxLength(maxLength string) (string, bool) {

	if maxLength == "null" {
		// default to 1024
		return "1024", false
	}
	return maxLength, true
}
<commit_msg>Modify the initial template SQL to include array data type and array dimensions Add support for array type to Add()<commit_after>//
// Copyright (c) 2017 Jon Carlson.  All rights reserved.
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.
//

package main

import (
	"bytes"
	"database/sql"
	"fmt"
	"github.com/joncrlsn/misc"
	"github.com/joncrlsn/pgutil"
	"sort"
	"strconv"
	"strings"
	"text/template"
)

var (
	columnSqlTemplate = initColumnSqlTemplate()
)

// Initializes the Sql template
func initColumnSqlTemplate() *template.Template {
	sql := `
SELECT table_schema
    , {{if eq $.DbSchema "*" }}table_schema || '.' || {{end}}table_name || '.' || column_name  AS compare_name
	, table_name
    , column_name
    , data_type
    , is_nullable
    , column_default
    , character_maximum_length
    , is_identity
    , identity_generation
	, substring(udt_name from 2) AS array_type
	, attndims AS array_dimensions
FROM information_schema.columns
JOIN (SELECT attname
		, attndims
		, relname
		, nspname
	FROM pg_attribute
	JOIN pg_class
	ON attrelid = pg_class.oid
	JOIN pg_namespace
	ON pg_class.relnamespace = pg_namespace.oid) s
ON (table_name = s.relname
	AND column_name = s.attname
	AND table_schema = s.nspname)
WHERE is_updatable = 'YES'
{{if eq $.DbSchema "*" }}
AND table_schema NOT LIKE 'pg_%'
AND table_schema <> 'information_schema'
{{else}}
AND table_schema = '{{$.DbSchema}}'
{{end}}
ORDER BY compare_name ASC;
`
	t := template.New("ColumnSqlTmpl")
	template.Must(t.Parse(sql))
	return t
}

// ==================================
// Column Rows definition
// ==================================

// ColumnRows is a sortable slice of string maps
type ColumnRows []map[string]string

func (slice ColumnRows) Len() int {
	return len(slice)
}

func (slice ColumnRows) Less(i, j int) bool {
	return slice[i]["compare_name"] < slice[j]["compare_name"]
}

func (slice ColumnRows) Swap(i, j int) {
	slice[i], slice[j] = slice[j], slice[i]
}

// ==================================
// ColumnSchema definition
// (implements Schema -- defined in pgdiff.go)
// ==================================

// ColumnSchema holds a slice of rows from one of the databases as well as
// a reference to the current row of data we're viewing.
type ColumnSchema struct {
	rows   ColumnRows
	rowNum int
	done   bool
}

// get returns the value from the current row for the given key
func (c *ColumnSchema) get(key string) string {
	if c.rowNum >= len(c.rows) {
		return ""
	}
	return c.rows[c.rowNum][key]
}

// NextRow increments the rowNum and tells you whether or not there are more
func (c *ColumnSchema) NextRow() bool {
	if c.rowNum >= len(c.rows)-1 {
		c.done = true
	}
	c.rowNum = c.rowNum + 1
	return !c.done
}

// Compare tells you, in one pass, whether or not the first row matches, is less than, or greater than the second row
func (c *ColumnSchema) Compare(obj interface{}) int {
	c2, ok := obj.(*ColumnSchema)
	if !ok {
		fmt.Println("Error!!!, Compare needs a ColumnSchema instance", c2)
	}

	val := misc.CompareStrings(c.get("compare_name"), c2.get("compare_name"))
	return val
}

// Add prints SQL to add the column
func (c *ColumnSchema) Add() {

	schema := dbInfo2.DbSchema
	if schema == "*" {
		schema = c.get("table_schema")
	}

	// Knowing the version of db2 would eliminate the need for this warning
	if c.get("is_identity") == "YES" {
		fmt.Println("-- WARNING: identity columns are not supported in PostgreSQL versions < 10.")
		fmt.Println("-- Attempting to create identity columns in earlier versions will probably result in errors.")
	}

	if c.get("data_type") == "character varying" {
		maxLength, valid := getMaxLength(c.get("character_maximum_length"))
		if !valid {
			fmt.Printf("ALTER TABLE %s.%s ADD COLUMN %s character varying", schema, c.get("table_name"), c.get("column_name"))
		} else {
			fmt.Printf("ALTER TABLE %s.%s ADD COLUMN %s character varying(%s)", schema, c.get("table_name"), c.get("column_name"), maxLength)
		}
	} else {
		dataType := c.get("data_type")
		//if c.get("data_type") == "ARRAY" {
			//fmt.Println("-- Note that adding of array data types are not yet generated properly.")
		//}
		if dataType == "ARRAY" {
			dimensions, err := strconv.Atoi(c.get("array_dimensions"))
			check("converting string to int", err)
			dataType = getArrayDefinition(c.get("array_type"), dimensions)
		}
		//fmt.Printf("ALTER TABLE %s.%s ADD COLUMN %s %s", schema, c.get("table_name"), c.get("column_name"), c.get("data_type"))
		fmt.Printf("ALTER TABLE %s.%s ADD COLUMN %s %s", schema, c.get("table_name"), c.get("column_name"), dataType)
	}

	if c.get("is_nullable") == "NO" {
		fmt.Printf(" NOT NULL")
	}
	if c.get("column_default") != "null" {
		fmt.Printf(" DEFAULT %s", c.get("column_default"))
	}
	// NOTE: there are more identity column sequence options according to the PostgreSQL
	// CREATE TABLE docs, but these do not appear to be available as of version 10.1
	if c.get("is_identity") == "YES" {
		fmt.Printf(" GENERATED %s AS IDENTITY", c.get("identity_generation"))
	}
	fmt.Printf(";\n")
}

// Drop prints SQL to drop the column
func (c *ColumnSchema) Drop() {
	// if dropping column
	fmt.Printf("ALTER TABLE %s.%s DROP COLUMN IF EXISTS %s;\n", c.get("table_schema"), c.get("table_name"), c.get("column_name"))
}

// Change handles the case where the table and column match, but the details do not
func (c *ColumnSchema) Change(obj interface{}) {
	c2, ok := obj.(*ColumnSchema)
	if !ok {
		fmt.Println("Error!!!, ColumnSchema.Change(obj) needs a ColumnSchema instance", c2)
	}

	// Detect column type change (mostly varchar length, or number size increase)
	// (integer to/from bigint is OK)
	if c.get("data_type") == c2.get("data_type") {
		if c.get("data_type") == "character varying" {
			max1, max1Valid := getMaxLength(c.get("character_maximum_length"))
			max2, max2Valid := getMaxLength(c2.get("character_maximum_length"))
			if !max1Valid && !max2Valid {
				// Leave them alone, they both have undefined max lengths
			} else if (max1Valid || !max2Valid) && (max1 != c2.get("character_maximum_length")) {
				//if !max1Valid {
				//    fmt.Println("-- WARNING: varchar column has no maximum length.  Setting to 1024, which may result in data loss.")
				//}
				max1Int, err1 := strconv.Atoi(max1)
				check("converting string to int", err1)
				max2Int, err2 := strconv.Atoi(max2)
				check("converting string to int", err2)
				if max1Int < max2Int {
					fmt.Println("-- WARNING: The next statement will shorten a character varying column, which may result in data loss.")
				}
				fmt.Printf("-- max1Valid: %v  max2Valid: %v \n", max1Valid, max2Valid)
				fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s TYPE character varying(%s);\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), max1)
			}
		}
	}

	// Code and test a column change from integer to bigint
	if c.get("data_type") != c2.get("data_type") {
		fmt.Printf("-- WARNING: This type change may not work well: (%s to %s).\n", c2.get("data_type"), c.get("data_type"))
		if strings.HasPrefix(c.get("data_type"), "character") {
			max1, max1Valid := getMaxLength(c.get("character_maximum_length"))
			if !max1Valid {
				fmt.Println("-- WARNING: varchar column has no maximum length.  Setting to 1024")
			}
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s TYPE %s(%s);\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), c.get("data_type"), max1)
		} else {
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s TYPE %s;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), c.get("data_type"))
		}
	}

	// Detect column default change (or added, dropped)
	if c.get("column_default") == "null" {
		if c.get("column_default") != "null" {
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s DROP DEFAULT;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"))
		}
	} else if c.get("column_default") != c2.get("column_default") {
		fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s SET DEFAULT %s;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), c.get("column_default"))
	}

	// Detect identity column change
	// Save result to variable instead of printing because order for adding/removing
	// is_nullable affects identity columns
	var identitySql string
	if c.get("is_identity") != c2.get("is_identity") {
		// Knowing the version of db2 would eliminate the need for this warning
		fmt.Println("-- WARNING: identity columns are not supported in PostgreSQL versions < 10.")
		fmt.Println("-- Attempting to create identity columns in earlier versions will probably result in errors.")
		if c.get("is_identity") == "YES" {
			identitySql = fmt.Sprintf("ALTER TABLE \"%s\".\"%s\" ALTER COLUMN \"%s\" ADD GENERATED %s AS IDENTITY;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), c.get("identity_generation"))
		} else {
			identitySql = fmt.Sprintf("ALTER TABLE \"%s\".\"%s\" ALTER COLUMN \"%s\" DROP IDENTITY;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"))
		}
	}

	// Detect not-null and nullable change
	if c.get("is_nullable") != c2.get("is_nullable") {
		if c.get("is_nullable") == "YES" {
			if identitySql != "" {
				fmt.Printf(identitySql)
			}
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s DROP NOT NULL;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"))
		} else {
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s SET NOT NULL;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"))
			if identitySql != "" {
				fmt.Printf(identitySql)
			}
		}
	} else {
		if identitySql != "" {
			fmt.Printf(identitySql)
		}
	}
}

// ==================================
// Standalone Functions
// ==================================

// compareColumns outputs SQL to make the columns match between two databases or schemas
func compareColumns(conn1 *sql.DB, conn2 *sql.DB) {

	buf1 := new(bytes.Buffer)
	columnSqlTemplate.Execute(buf1, dbInfo1)

	buf2 := new(bytes.Buffer)
	columnSqlTemplate.Execute(buf2, dbInfo2)

	rowChan1, _ := pgutil.QueryStrings(conn1, buf1.String())
	rowChan2, _ := pgutil.QueryStrings(conn2, buf2.String())

	//rows1 := make([]map[string]string, 500)
	rows1 := make(ColumnRows, 0)
	for row := range rowChan1 {
		rows1 = append(rows1, row)
	}
	sort.Sort(rows1)

	//rows2 := make([]map[string]string, 500)
	rows2 := make(ColumnRows, 0)
	for row := range rowChan2 {
		rows2 = append(rows2, row)
	}
	sort.Sort(&rows2)

	// We have to explicitly type this as Schema here for some unknown reason
	var schema1 Schema = &ColumnSchema{rows: rows1, rowNum: -1}
	var schema2 Schema = &ColumnSchema{rows: rows2, rowNum: -1}

	// Compare the columns
	doDiff(schema1, schema2)
}

// getMaxLength returns the maximum length and whether or not it is valid
func getMaxLength(maxLength string) (string, bool) {

	if maxLength == "null" {
		// default to 1024
		return "1024", false
	}
	return maxLength, true
}

func getArrayDefinition(arrayType string, dimensions int) string {
	return arrayType + strings.Repeat("[]", dimensions)
}
<|endoftext|>