<commit_before>package main

import (
	"bufio"
	"compress/zlib"
	"encoding/binary"
	"fmt"
	"io"
	"os"
)

type PrtGenerator struct {
	enclosedsChan chan *EnclosedChunkJob
	completeChan  chan bool

	outFile *os.File
	w       *bufio.Writer
	zw      io.WriteCloser

	particleCount int64
	total         int
	boundary      *BoundaryLocator
}

func (o *PrtGenerator) Start(outFilename string, total int, maxProcs int, boundary *BoundaryLocator) {
	o.enclosedsChan = make(chan *EnclosedChunkJob, maxProcs*2)
	o.completeChan = make(chan bool)
	o.total = total
	o.boundary = boundary

	maxProcs = 1
	for i := 0; i < maxProcs; i++ {
		go o.chunkProcessor()
	}

	var openErr os.Error

	o.outFile, openErr = os.Open(outFilename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
	if openErr != nil {
		fmt.Fprintln(os.Stderr, openErr) // TODO: return openErr
		return
	}

	o.w = bufio.NewWriter(o.outFile)
	WriteHeader(o.w, -1, []ChannelDefinition{{"Position", 4, 3, 0}, {"BlockID", 1, 1, 12}})

	var zErr os.Error
	o.zw, zErr = zlib.NewWriterLevel(o.w, zlib.NoCompression)
	if zErr != nil {
		fmt.Fprintln(os.Stderr, zErr) // TODO: return zErr
		return
	}
}

func (o *PrtGenerator) chunkProcessor() {
	var chunkCount = 0
	for {
		var job = <-o.enclosedsChan

		var e = job.enclosed

		for i := 0; i < len(e.blocks); i += 128 {
			var x, z = (i / 128) / 16, (i / 128) % 16

			var column = BlockColumn(e.blocks[i : i+128])
			for y, blockId := range column {
				if y < yMin {
					continue
				}

				switch {
				case o.boundary.IsBoundary(blockId, e.Get(x, y-1, z)):
					fallthrough
				case o.boundary.IsBoundary(blockId, e.Get(x, y+1, z)):
					fallthrough
				case o.boundary.IsBoundary(blockId, e.Get(x-1, y, z)):
					fallthrough
				case o.boundary.IsBoundary(blockId, e.Get(x+1, y, z)):
					fallthrough
				case o.boundary.IsBoundary(blockId, e.Get(x, y, z-1)):
					fallthrough
				case o.boundary.IsBoundary(blockId, e.Get(x, y, z+1)):
					o.particleCount++
					var (
						xa = -(x + e.xPos*16)
						ya = y - 64
						za = -(z + e.zPos*16)
					)
					binary.Write(o.zw, binary.LittleEndian, float32(xa*2))
					binary.Write(o.zw, binary.LittleEndian, float32(za*2))
					binary.Write(o.zw, binary.LittleEndian, float32(ya*2))
					binary.Write(o.zw, binary.LittleEndian, int32(blockId))
				}
			}
		}

		chunkCount++
		fmt.Printf("%4v/%-4v (%3v,%3v) Particles: %d\n", chunkCount, o.total, job.enclosed.xPos, job.enclosed.zPos, o.particleCount)

		if job.last {
			o.completeChan <- true
			break
		}
	}
}

func (o *PrtGenerator) Close() {
	o.zw.Close()
	o.w.Flush()
	UpdateParticleCount(o.outFile, o.particleCount)
	o.outFile.Close()
}

func (o *PrtGenerator) GetEnclosedJobsChan() chan *EnclosedChunkJob {
	return o.enclosedsChan
}

func (o *PrtGenerator) GetCompleteChan() chan bool {
	return o.completeChan
}

type ChannelDefinition struct {
	Name                    string // max of 31 characters and must meet the regex [a-zA-Z_][0-9a-zA-Z_]*
	DataType, Arity, Offset int32
}

// http://software.primefocusworld.com/software/support/krakatoa/prt_file_format.php
func WriteHeader(w io.Writer, particleCount int64, channels []ChannelDefinition) {
	// Header (56 bytes)
	var magic = []byte{192, 'P', 'R', 'T', '\r', '\n', 26, '\n'}
	w.Write(magic)

	var headerLength = uint32(56)
	binary.Write(w, binary.LittleEndian, headerLength)

	var signature = make([]byte, 32)
	copy(signature, []byte("Extensible Particle Format"))
	w.Write(signature)

	var version = uint32(1)
	binary.Write(w, binary.LittleEndian, version)

	binary.Write(w, binary.LittleEndian, particleCount)

	// Reserved bytes (4 bytes)
	var reserved = int32(4)
	binary.Write(w, binary.LittleEndian, reserved)

	// Channel definition header (8 bytes)
	binary.Write(w, binary.LittleEndian, int32(len(channels)))

	var channelDefinitionSize = int32(44)
	binary.Write(w, binary.LittleEndian, channelDefinitionSize)

	for _, channel := range channels {
		var nameBytes = make([]byte, 32)
		copy(nameBytes, []byte(channel.Name))
		w.Write(nameBytes)

		binary.Write(w, binary.LittleEndian, channel.DataType)
		binary.Write(w, binary.LittleEndian, channel.Arity)
		binary.Write(w, binary.LittleEndian, channel.Offset)
	}
}

func UpdateParticleCount(file *os.File, particleCount int64) os.Error {
	var storedOffset, err = file.Seek(0, 1)
	if err != nil {
		return err
	}
	_, err = file.Seek(0x30, 0)
	if err != nil {
		return err
	}
	err = binary.Write(file, binary.LittleEndian, particleCount)
	if err != nil {
		return err
	}
	_, err = file.Seek(storedOffset, 0)
	return err
}
<commit_msg>Invert only the x-axis not the z-axis for prt dumps<commit_after>package main

import (
	"bufio"
	"compress/zlib"
	"encoding/binary"
	"fmt"
	"io"
	"os"
)

type PrtGenerator struct {
	enclosedsChan chan *EnclosedChunkJob
	completeChan  chan bool

	outFile *os.File
	w       *bufio.Writer
	zw      io.WriteCloser

	particleCount int64
	total         int
	boundary      *BoundaryLocator
}

func (o *PrtGenerator) Start(outFilename string, total int, maxProcs int, boundary *BoundaryLocator) {
	o.enclosedsChan = make(chan *EnclosedChunkJob, maxProcs*2)
	o.completeChan = make(chan bool)
	o.total = total
	o.boundary = boundary

	maxProcs = 1
	for i := 0; i < maxProcs; i++ {
		go o.chunkProcessor()
	}

	var openErr os.Error

	o.outFile, openErr = os.Open(outFilename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
	if openErr != nil {
		fmt.Fprintln(os.Stderr, openErr) // TODO: return openErr
		return
	}

	o.w = bufio.NewWriter(o.outFile)
	WriteHeader(o.w, -1, []ChannelDefinition{{"Position", 4, 3, 0}, {"BlockID", 1, 1, 12}})

	var zErr os.Error
	o.zw, zErr = zlib.NewWriterLevel(o.w, zlib.NoCompression)
	if zErr != nil {
		fmt.Fprintln(os.Stderr, zErr) // TODO: return zErr
		return
	}
}

func (o *PrtGenerator) chunkProcessor() {
	var chunkCount = 0
	for {
		var job = <-o.enclosedsChan

		var e = job.enclosed

		for i := 0; i < len(e.blocks); i += 128 {
			var x, z = (i / 128) / 16, (i / 128) % 16

			var column = BlockColumn(e.blocks[i : i+128])
			for y, blockId := range column {
				if y < yMin {
					continue
				}

				switch {
				case o.boundary.IsBoundary(blockId, e.Get(x, y-1, z)):
					fallthrough
				case o.boundary.IsBoundary(blockId, e.Get(x, y+1, z)):
					fallthrough
				case o.boundary.IsBoundary(blockId, e.Get(x-1, y, z)):
					fallthrough
				case o.boundary.IsBoundary(blockId, e.Get(x+1, y, z)):
					fallthrough
				case o.boundary.IsBoundary(blockId, e.Get(x, y, z-1)):
					fallthrough
				case o.boundary.IsBoundary(blockId, e.Get(x, y, z+1)):
					o.particleCount++
					var (
						xa = -(x + e.xPos*16)
						ya = y - 64
						za = z + e.zPos*16
					)
					binary.Write(o.zw, binary.LittleEndian, float32(xa*2))
					binary.Write(o.zw, binary.LittleEndian, float32(za*2))
					binary.Write(o.zw, binary.LittleEndian, float32(ya*2))
					binary.Write(o.zw, binary.LittleEndian, int32(blockId))
				}
			}
		}

		chunkCount++
		fmt.Printf("%4v/%-4v (%3v,%3v) Particles: %d\n", chunkCount, o.total, job.enclosed.xPos, job.enclosed.zPos, o.particleCount)

		if job.last {
			o.completeChan <- true
			break
		}
	}
}

func (o *PrtGenerator) Close() {
	o.zw.Close()
	o.w.Flush()
	UpdateParticleCount(o.outFile, o.particleCount)
	o.outFile.Close()
}

func (o *PrtGenerator) GetEnclosedJobsChan() chan *EnclosedChunkJob {
	return o.enclosedsChan
}

func (o *PrtGenerator) GetCompleteChan() chan bool {
	return o.completeChan
}

type ChannelDefinition struct {
	Name                    string // max of 31 characters and must meet the regex [a-zA-Z_][0-9a-zA-Z_]*
	DataType, Arity, Offset int32
}

// http://software.primefocusworld.com/software/support/krakatoa/prt_file_format.php
func WriteHeader(w io.Writer, particleCount int64, channels []ChannelDefinition) {
	// Header (56 bytes)
	var magic = []byte{192, 'P', 'R', 'T', '\r', '\n', 26, '\n'}
	w.Write(magic)

	var headerLength = uint32(56)
	binary.Write(w, binary.LittleEndian, headerLength)

	var signature = make([]byte, 32)
	copy(signature, []byte("Extensible Particle Format"))
	w.Write(signature)

	var version = uint32(1)
	binary.Write(w, binary.LittleEndian, version)

	binary.Write(w, binary.LittleEndian, particleCount)

	// Reserved bytes (4 bytes)
	var reserved = int32(4)
	binary.Write(w, binary.LittleEndian, reserved)

	// Channel definition header (8 bytes)
	binary.Write(w, binary.LittleEndian, int32(len(channels)))

	var channelDefinitionSize = int32(44)
	binary.Write(w, binary.LittleEndian, channelDefinitionSize)

	for _, channel := range channels {
		var nameBytes = make([]byte, 32)
		copy(nameBytes, []byte(channel.Name))
		w.Write(nameBytes)

		binary.Write(w, binary.LittleEndian, channel.DataType)
		binary.Write(w, binary.LittleEndian, channel.Arity)
		binary.Write(w, binary.LittleEndian, channel.Offset)
	}
}

func UpdateParticleCount(file *os.File, particleCount int64) os.Error {
	var storedOffset, err = file.Seek(0, 1)
	if err != nil {
		return err
	}
	_, err = file.Seek(0x30, 0)
	if err != nil {
		return err
	}
	err = binary.Write(file, binary.LittleEndian, particleCount)
	if err != nil {
		return err
	}
	_, err = file.Seek(storedOffset, 0)
	return err
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"

	"golang.org/x/crypto/ssh"

	"github.com/vallard/psh/nr"
	"github.com/vallard/psh/server"
)

func errorFunction(errMessage string) {
	fmt.Printf("%s\n", errMessage)
	os.Exit(1)
}

// prints the usage of the tool and exits with the value past in.
func useage(rc int) {
	fmt.Printf("psh <noderange> command args args ...\n")
	os.Exit(rc)
}

func main() {
	if len(os.Args) < 2 {
		useage(1)
	}

	// the first argument is the range or group of nodes
	nodes, err := nr.GetNodeRange(os.Args[1])
	if err != nil {
		errorFunction(err.Error())
	}

	// open the file for ssh stuff.
	cmd := getCommand(os.Args[2:])
	fmt.Printf("Running command: %s\n", cmd)

	ch := make(chan string)
	for _, server := range nodes {
		fmt.Println(server.Host)
		go runcmd(server, cmd, ch)
	}

	for range nodes {
		fmt.Println(<-ch)
	}
}

func runcmd(server server.Server, cmd string, ch chan<- string) {
	sshConfig := &ssh.ClientConfig{
		User: server.User,
		Auth: []ssh.AuthMethod{
			PublicKeyFile(server.Key),
		},
	}

	// create connection
	s := getServerAndPort(server)
	connection, err := ssh.Dial("tcp", s, sshConfig)
	if err != nil {
		ch <- fmt.Sprint(err)
		return
	}

	// create session
	session, err := connection.NewSession()
	if err != nil {
		ch <- fmt.Sprint(err)
		return
	}
	log.Printf("Running command on %s", server.IP)
	session.Run(cmd)
}

func getServerAndPort(server server.Server) string {
	if strings.Contains(server.IP, ":") {
		return server.IP
	}
	s := []string{server.IP, "22"}
	return strings.Join(s, ":")
}

// get the key file. cred: http://blog.ralch.com/tutorial/golang-ssh-connection/
func PublicKeyFile(file string) ssh.AuthMethod {
	file = nr.ExpandShell(file)
	buffer, err := ioutil.ReadFile(file)
	if err != nil {
		return nil
	}

	key, err := ssh.ParsePrivateKey(buffer)
	if err != nil {
		return nil
	}
	return ssh.PublicKeys(key)
}

// get the command from the command line
func getCommand(arr []string) string {
	return strings.Join(arr, " ")
}
<commit_msg>working mvp<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"golang.org/x/crypto/ssh"

	"github.com/vallard/psh/nr"
	"github.com/vallard/psh/server"
)

func errorFunction(errMessage string) {
	fmt.Printf("%s\n", errMessage)
	os.Exit(1)
}

// prints the usage of the tool and exits with the value past in.
func useage(rc int) {
	fmt.Printf("psh <noderange> command args args ...\n")
	os.Exit(rc)
}

func main() {
	if len(os.Args) < 2 {
		useage(1)
	}

	// the first argument is the range or group of nodes
	nodes, err := nr.GetNodeRange(os.Args[1])
	if err != nil {
		errorFunction(err.Error())
	}

	// open the file for ssh stuff.
	cmd := getCommand(os.Args[2:])
	//fmt.Printf("Running command: %s\n", cmd)

	ch := make(chan string)
	for _, server := range nodes {
		//fmt.Println(server.Host)
		go runcmd(server, cmd, ch)
	}

	for range nodes {
		fmt.Println(<-ch)
	}
}

func runcmd(server server.Server, cmd string, ch chan<- string) {
	sshConfig := &ssh.ClientConfig{
		User: server.User,
		Auth: []ssh.AuthMethod{
			PublicKeyFile(server.Key),
		},
	}

	// create connection
	s := getServerAndPort(server)
	connection, err := ssh.Dial("tcp", s, sshConfig)
	if err != nil {
		ch <- fmt.Sprint(err)
		return
	}

	// create session
	session, err := connection.NewSession()
	if err != nil {
		ch <- fmt.Sprint(err)
		return
	}
	defer session.Close()
	// create terminal
	/*
		modes := ssh.TerminalModes{
			ssh.ECHO:          0,
			ssh.TTY_OP_ISPEED: 14400,
			ssh.TTY_OP_OSPEED: 14400,
		}
			if err = session.RequestPty("xterm", 80, 40, modes); err != nil {
				session.Close()
				ch <- fmt.Sprintf("request for pseudo terminal failed: %s", err)
				return
			}
	*/
	var stdoutBuff bytes.Buffer
	//stdout, err := session.StdoutPipe()
	session.Stdout = &stdoutBuff

	// copy pipe stuff
	/*stdin, err := session.StdinPipe()
	if err != nil {
		//ch <- fmt.Errorf("Unable to setup stdin for session: %v", err)
		ch <- fmt.Sprintf("Unable to setup stdin for session: %v", err)
		return
	}
	go io.Copy(stdin, os.Stdin)

	stdout, err := session.StdoutPipe()
	if err != nil {
		//ch <- fmt.Errorf("Unable to setup stdout for session: %v", err)
		ch <- fmt.Sprintf("Unable to setup stdout for session: %v", err)
		return
	}
	go io.Copy(os.Stdout, stdout)

	stderr, err := session.StderrPipe()
	if err != nil {
		//ch <- fmt.Errorf("Unable to setup stderr for session: %v", err)
		ch <- fmt.Sprintf("Unable to setup stderr for session: %v", err)
		return
	}
	go io.Copy(os.Stderr, stderr)
	*/
	//log.Printf("Running command on %s", server.IP)
	err = session.Run(cmd)
	if err != nil {
		ch <- fmt.Sprintf("Unable to setup stderr for session: %v", err)
	} else {
		ch <- server.Host + ": " + stdoutBuff.String()
	}
	return
}

func getServerAndPort(server server.Server) string {
	if strings.Contains(server.IP, ":") {
		return server.IP
	}
	s := []string{server.IP, "22"}
	return strings.Join(s, ":")
}

// get the key file. cred: http://blog.ralch.com/tutorial/golang-ssh-connection/
func PublicKeyFile(file string) ssh.AuthMethod {
	file = nr.ExpandShell(file)
	buffer, err := ioutil.ReadFile(file)
	if err != nil {
		return nil
	}

	key, err := ssh.ParsePrivateKey(buffer)
	if err != nil {
		return nil
	}
	return ssh.PublicKeys(key)
}

// get the command from the command line
func getCommand(arr []string) string {
	return strings.Join(arr, " ")
}
<|endoftext|>


<commit_before>package isolated

import (
	"code.cloudfoundry.org/cli/integration/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("table alignment", func() {
	BeforeEach(func() {
		helpers.LoginCF()
	})

	When("output is in English", func() {
		BeforeEach(func() {
			helpers.LoginCF()
			helpers.TargetOrgAndSpace(ReadOnlyOrg, ReadOnlySpace)
		})

		// Developer note: The spacing in this test is significant and explicit. Do
		// not replace with a regex.
		It("aligns the table correctly", func() {
			username, _ := helpers.GetCredentials()
			session := helpers.CF("target")
			Eventually(session).Should(Say("api endpoint:   %s", apiURL))
			Eventually(session).Should(Say(`api version:    [\d.]+`))
			Eventually(session).Should(Say("user:           %s", username))
			Eventually(session).Should(Say("org:            %s", ReadOnlyOrg))
			Eventually(session).Should(Say("space:          %s", ReadOnlySpace))
			Eventually(session).Should(Exit(0))
		})
	})

	When("output is in language with multibyte characters", func() {
		BeforeEach(func() {
			helpers.LoginCF()
			helpers.TargetOrgAndSpace(ReadOnlyOrg, ReadOnlySpace)
		})

		// Developer note: The spacing in this test is significant and explicit. Do
		// not replace with a regex.
		// TODO: add these strings to the i18n/resources
		// See: https://www.pivotaltracker.com/story/show/151737497
		XIt("aligns the table correctly", func() {
			username, _ := helpers.GetCredentials()
			session := helpers.CFWithEnv(map[string]string{"LANG": "ja-JP.utf8"}, "target")
			Eventually(session).Should(Say("API エンドポイント:   %s", apiURL))
			Eventually(session).Should(Say(`api version:          [\d.]+`))
			Eventually(session).Should(Say("ユーザー:             %s", username))
			Eventually(session).Should(Say("組織:                 %s", ReadOnlyOrg))
			Eventually(session).Should(Say("スペース:             %s", ReadOnlySpace))
			Eventually(session).Should(Exit(0))
		})
	})
})
<commit_msg>Remove pending from multi-byte table alignment example<commit_after>package isolated

import (
	"code.cloudfoundry.org/cli/integration/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("table alignment", func() {
	BeforeEach(func() {
		helpers.LoginCF()
	})

	When("output is in English", func() {
		BeforeEach(func() {
			helpers.LoginCF()
			helpers.TargetOrgAndSpace(ReadOnlyOrg, ReadOnlySpace)
		})

		// Developer note: The spacing in this test is significant and explicit. Do
		// not replace with a regex.
		It("aligns the table correctly", func() {
			username, _ := helpers.GetCredentials()
			session := helpers.CF("target")
			Eventually(session).Should(Say("api endpoint:   %s", apiURL))
			Eventually(session).Should(Say(`api version:    [\d.]+`))
			Eventually(session).Should(Say("user:           %s", username))
			Eventually(session).Should(Say("org:            %s", ReadOnlyOrg))
			Eventually(session).Should(Say("space:          %s", ReadOnlySpace))
			Eventually(session).Should(Exit(0))
		})
	})

	When("output is in language with multibyte characters", func() {
		BeforeEach(func() {
			helpers.LoginCF()
			helpers.TargetOrgAndSpace(ReadOnlyOrg, ReadOnlySpace)
		})

		// Developer note: The spacing in this test is significant and explicit. Do
		// not replace with a regex.
		It("aligns the table correctly", func() {
			username, _ := helpers.GetCredentials()
			session := helpers.CFWithEnv(map[string]string{"LANG": "ja-JP.utf8"}, "target")
			Eventually(session).Should(Say("API エンドポイント   %s", apiURL))
			Eventually(session).Should(Say(`API バージョン:      [\d.]+`))
			Eventually(session).Should(Say("ユーザー:            %s", username))
			Eventually(session).Should(Say("組織:                %s", ReadOnlyOrg))
			Eventually(session).Should(Say("スペース:            %s", ReadOnlySpace))
			Eventually(session).Should(Exit(0))
		})
	})
})
<|endoftext|>


<commit_before>// Implements INSERTS output for exported rows
// e.g.
// INSERT INTO <table> (name, last_name, something);
package formats

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"strconv"
	"strings"
	"time"
)

const (
	insertQuery = "INSERT INTO %s (%s) "
)

// very simple quoting for sql values
func quote(val string) string {
	buf := bytes.NewBufferString("'")
	buf.WriteString(strings.Replace(val, "'", "\\'", -1))
	buf.WriteString("'")
	return buf.String()
}

type InsertsFormat struct {
	DataFormat
	tableName   string
	columns     []string
	multiInsert bool
	writer      *bufio.Writer
}

func NewInsertsFormat(w io.Writer, fileName string, tableName string) (*InsertsFormat, error) {
	return &InsertsFormat{
		writer:      bufio.NewWriter(w),
		tableName:   tableName,
		columns:     make([]string, 0),
		multiInsert: false,
	}, nil
}

func (f *InsertsFormat) WriteHeader(columns []string) error {
	f.columns = columns
	return nil
}

func (f *InsertsFormat) Flush() error { return nil }

func (f *InsertsFormat) WriteRow(values map[string]interface{}) error {
	columnsVal := strings.Join(f.columns, ",")
	queryStr := fmt.Sprintf(insertQuery, f.tableName, columnsVal)
	buf := bytes.NewBufferString(queryStr)
	record := []string{}
	for _, col := range f.columns {
		switch value := (values[col]).(type) {
		case string:
			record = append(record, quote(value))
		case []byte:
			record = append(record, quote(string(value)))
		case int64:
			record = append(record, fmt.Sprintf("%d", value))
		case float64:
			record = append(record, strconv.FormatFloat(value, 'f', -1, 64))
		case time.Time:
			record = append(record, quote(value.Format(time.RFC3339)))
		case bool:
			if value == true {
				record = append(record, "true")
			} else {
				record = append(record, "false")
			}
		}
	}
	buf.WriteString("VALUES (")
	buf.WriteString(strings.Join(record, ","))
	buf.WriteString(");\n")

	_, err := f.writer.Write(buf.Bytes())
	if err != nil {
		return err
	}
	err = f.writer.Flush()
	// defer f.writer.Close()
	return err
}
<commit_msg>Insert nulls for nil values<commit_after>// Implements INSERTS output for exported rows
// e.g.
// INSERT INTO <table> (name, last_name, something);
package formats

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"strconv"
	"strings"
	"time"
)

const (
	insertQuery = "INSERT INTO %s (%s) "
)

// very simple quoting for sql values
func quote(val string) string {
	buf := bytes.NewBufferString("'")
	buf.WriteString(strings.Replace(val, "'", "\\'", -1))
	buf.WriteString("'")
	return buf.String()
}

type InsertsFormat struct {
	DataFormat
	tableName   string
	columns     []string
	multiInsert bool
	writer      *bufio.Writer
}

func NewInsertsFormat(w io.Writer, fileName string, tableName string) (*InsertsFormat, error) {
	return &InsertsFormat{
		writer:      bufio.NewWriter(w),
		tableName:   tableName,
		columns:     make([]string, 0),
		multiInsert: false,
	}, nil
}

func (f *InsertsFormat) WriteHeader(columns []string) error {
	f.columns = columns
	return nil
}

func (f *InsertsFormat) Flush() error { return nil }

func (f *InsertsFormat) WriteRow(values map[string]interface{}) error {
	columnsVal := strings.Join(f.columns, ",")
	queryStr := fmt.Sprintf(insertQuery, f.tableName, columnsVal)
	buf := bytes.NewBufferString(queryStr)
	record := []string{}
	for _, col := range f.columns {
		switch value := (values[col]).(type) {
		case string:
			record = append(record, quote(value))
		case []byte:
			record = append(record, quote(string(value)))
		case int64:
			record = append(record, fmt.Sprintf("%d", value))
		case float64:
			record = append(record, strconv.FormatFloat(value, 'f', -1, 64))
		case time.Time:
			record = append(record, quote(value.Format(time.RFC3339)))
		case bool:
			if value == true {
				record = append(record, "true")
			} else {
				record = append(record, "false")
			}
		case nil:
			record = append(record, "null")
		}
	}
	buf.WriteString("VALUES (")
	buf.WriteString(strings.Join(record, ","))
	buf.WriteString(");\n")

	_, err := f.writer.Write(buf.Bytes())
	if err != nil {
		return err
	}
	err = f.writer.Flush()
	// defer f.writer.Close()
	return err
}
<|endoftext|>


<commit_before>package uic

import (
	"encoding/base64"
	"encoding/json"
	"github.com/Cepave/open-falcon-backend/modules/fe/g"
	"github.com/Cepave/open-falcon-backend/modules/fe/http/base"
	. "github.com/Cepave/open-falcon-backend/modules/fe/model/uic"
	"github.com/Cepave/open-falcon-backend/modules/fe/utils"
	log "github.com/Sirupsen/logrus"
	"github.com/toolkits/str"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"
)

type AuthController struct {
	base.BaseController
}

func (this *AuthController) Logout() {
	u := this.Ctx.Input.GetData("CurrentUser").(*User)
	token := this.Ctx.GetCookie("token")
	if len(token) > 0 {
		url := g.Config().Api.Logout + "/" + token
		log.Println("logout url =", url)
		result := sendHttpGetRequest(url)
		log.Println("logout result =", result)
		this.Ctx.SetCookie("token", "", 0, "/")
		this.Ctx.SetCookie("token", "", 0, "/", g.Config().Http.Cookie)
	}
	RemoveSessionByUid(u.Id)
	this.Ctx.SetCookie("sig", "", 0, "/")
	this.Ctx.SetCookie("sig", "", 0, "/", g.Config().Http.Cookie)
	this.Ctx.SetCookie("name", "", 0, "/")
	this.Ctx.SetCookie("name", "", 0, "/", g.Config().Http.Cookie)
	this.Redirect("/auth/login", 302)
}

func (this *AuthController) LoginGet() {
	appSig := this.GetString("sig", "")
	callback := this.GetString("callback", "")

	cookieSig := this.Ctx.GetCookie("sig")
	if cookieSig == "" {
		this.renderLoginPage(appSig, callback)
		return
	}

	sessionObj := ReadSessionBySig(cookieSig)
	if sessionObj == nil {
		this.renderLoginPage(appSig, callback)
		return
	}

	if int64(sessionObj.Expired) < time.Now().Unix() {
		RemoveSessionByUid(sessionObj.Uid)
		this.renderLoginPage(appSig, callback)
		return
	}

	if appSig != "" && callback != "" {
		this.Redirect(callback, 302)
	} else {
		this.Redirect("/me/info", 302)
	}
}

func (this *AuthController) LoginPost() {
	name := this.GetString("name", "")
	password := this.GetString("password", "")
	if name == "" || password == "" {
		this.ServeErrJson("name or password is blank")
		return
	}

	var u *User

	ldapEnabled := this.MustGetBool("ldap", false)

	if ldapEnabled {
		sucess, err := utils.LdapBind(g.Config().Ldap.Addr,
			g.Config().Ldap.BaseDN,
			g.Config().Ldap.BindDN,
			g.Config().Ldap.BindPasswd,
			g.Config().Ldap.UserField,
			name,
			password)
		if err != nil {
			this.ServeErrJson(err.Error())
			return
		}

		if !sucess {
			this.ServeErrJson("name or password error")
			return
		}

		user_attributes, err := utils.Ldapsearch(g.Config().Ldap.Addr,
			g.Config().Ldap.BaseDN,
			g.Config().Ldap.BindDN,
			g.Config().Ldap.BindPasswd,
			g.Config().Ldap.UserField,
			name,
			g.Config().Ldap.Attributes)
		userSn := ""
		userMail := ""
		userTel := ""
		if err == nil {
			userSn = user_attributes["sn"]
			userMail = user_attributes["mail"]
			userTel = user_attributes["telephoneNumber"]
		}

		arr := strings.Split(name, "@")
		var userName, userEmail string
		if len(arr) == 2 {
			userName = arr[0]
			userEmail = name
		} else {
			userName = name
			userEmail = userMail
		}

		u = ReadUserByName(userName)
		if u == nil {
			// 说明用户不存在
			u = &User{
				Name:   userName,
				Passwd: "",
				Cnname: userSn,
				Phone:  userTel,
				Email:  userEmail,
			}
			_, err = u.Save()
			if err != nil {
				this.ServeErrJson("insert user fail " + err.Error())
				return
			}
		}
	} else {
		u = ReadUserByName(name)
		if u == nil {
			this.ServeErrJson("no such user")
			return
		}

		if u.Passwd != str.Md5Encode(g.Config().Salt+password) {
			this.ServeErrJson("password error")
			return
		}
	}

	appSig := this.GetString("sig", "")
	callback := this.GetString("callback", "")
	if appSig != "" && callback != "" {
		SaveSessionAttrs(u.Id, appSig, int(time.Now().Unix())+3600*24*30)
	} else {
		this.CreateSession(u.Id, 3600*24*30)
	}

	this.ServeDataJson(callback)
}

func (this *AuthController) renderLoginPage(sig, callback string) {
	this.Data["CanRegister"] = g.Config().CanRegister
	this.Data["LdapEnabled"] = g.Config().Ldap.Enabled
	this.Data["Sig"] = sig
	this.Data["Callback"] = callback
	this.Data["Shortcut"] = g.Config().Shortcut
	this.TplName = "auth/login.html"
}

func (this *AuthController) RegisterGet() {
	this.Data["CanRegister"] = g.Config().CanRegister
	this.Data["Shortcut"] = g.Config().Shortcut
	this.TplName = "auth/register.html"
}

func (this *AuthController) RegisterPost() {
	if !g.Config().CanRegister {
		this.ServeErrJson("registration system is not open")
		return
	}

	name := strings.TrimSpace(this.GetString("name", ""))
	password := strings.TrimSpace(this.GetString("password", ""))
	repeatPassword := strings.TrimSpace(this.GetString("repeat_password", ""))
	email := strings.TrimSpace(this.GetString("email", ""))

	if password != repeatPassword {
		this.ServeErrJson("password not equal the repeart one")
		return
	}

	if !utils.IsUsernameValid(name) {
		this.ServeErrJson("name pattern is invalid")
		return
	}

	if ReadUserIdByName(name) > 0 {
		this.ServeErrJson("name is already existent")
		return
	}

	lastId, err := InsertRegisterUser(name, str.Md5Encode(g.Config().Salt+password), email)
	if err != nil {
		this.ServeErrJson("insert user fail " + err.Error())
		return
	}

	this.CreateSession(lastId, 3600*24*30)

	this.ServeOKJson()
}

func (this *AuthController) CreateSession(uid int64, maxAge int) int {
	sig := utils.GenerateUUID()
	user := SelectUserById(uid)
	expired := int(time.Now().Unix()) + maxAge
	SaveSessionAttrs(uid, sig, expired)
	this.Ctx.SetCookie("sig", sig, maxAge, "/")
	this.Ctx.SetCookie("sig", sig, maxAge, "/", g.Config().Http.Cookie)
	this.Ctx.SetCookie("name", user.Name, maxAge, "/")
	this.Ctx.SetCookie("name", user.Name, maxAge, "/", g.Config().Http.Cookie)
	return expired
}

/**
 * @function name:   func (this *AuthController) LoginThirdParty()
 * @description:     This function returns third party login URL.
 * @related issues:  OWL-206
 * @param:           void
 * @return:          void
 * @author:          Don Hsieh
 * @since:           12/17/2015
 * @last modified:   12/17/2015
 * @called by:       beego.Router("/auth/third-party", &AuthController{}, "post:LoginThirdParty")
 *                    in fe/http/uic/uic_routes.go
 */
func (this *AuthController) LoginThirdParty() {
	s := g.Config().Api.Redirect
	s = base64.StdEncoding.EncodeToString([]byte(s))
	strEncoded := url.QueryEscape(s)
	loginUrl := g.Config().Api.Login + "/" + strEncoded
	this.ServeDataJson(loginUrl)
}

/**
 * @function name:   func sendHttpGetRequest(url string) map[string]interface{}
 * @description:     This function sends GET request to given URL.
 * @related issues:  OWL-206, OWL-159
 * @param:           url string
 * @return:          map[string]interface{}
 * @author:          Don Hsieh
 * @since:           12/17/2015
 * @last modified:   12/17/2015
 * @called by:       func (this *AuthController) LoginWithToken()
 *                    in fe/http/uic/auth_controller.go
 */
func sendHttpGetRequest(url string) map[string]interface{} {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		log.Println("Error =", err.Error())
	}

	client := &http.Client{}
	resp, err := client.Do(req)
	defer resp.Body.Close()
	if err != nil {
		log.Println("Error =", err.Error())
		return nil
	} else {
		body, _ := ioutil.ReadAll(resp.Body)
		var nodes = make(map[string]interface{})
		if err := json.Unmarshal(body, &nodes); err != nil {
			log.Println("Error =", err.Error())
			return nil
		}
		return nodes
	}
}

func setUserInfo(nodes map[string]interface{}, userInfo map[string]string) {
	if status, ok := nodes["status"]; ok {
		if int(status.(float64)) == 1 {
			data := nodes["data"].(map[string]interface{})
			access_key := data["access_key"].(string)
			username := data["username"].(string)
			email := data["email"].(string)
			log.Println("access_key =", access_key)
			userInfo["username"] = username
			userInfo["email"] = email
			userInfo["access_key"] = access_key
		}
	}
}

func getUserRole(access_key string) int {
	urlRole := g.Config().Api.Role + "/" + access_key
	nodes := sendHttpGetRequest(urlRole)
	role := -1
	if int(nodes["status"].(float64)) == 1 {
		permission := nodes["data"]
		log.Println("permission =", permission)
		if permission == "admin" {
			role = 0
		} else if permission == "operator" {
			role = 1
		} else if permission == "observer" {
			role = 2
		} else if permission == "deny" {
			role = 3
		}
	}

	// TODO: The role should be able to be changed on BOSS in the future.
	role = 1
	return role
}

/**
 * @function name:   func (this *AuthController) LoginWithToken()
 * @description:     This function logins user with third party token.
 * @related issues:  OWL-247, OWL-206
 * @param:           void
 * @return:          void
 * @author:          Don Hsieh
 * @since:           12/16/2015
 * @last modified:   01/08/2016
 * @called by:       beego.Router("/auth/login/:token", &AuthController{}, "get:LoginWithToken")
 *                    in fe/http/uic/uic_routes.go
 */
func (this *AuthController) LoginWithToken() {
	log.Println("func (this *AuthController) LoginWithToken()")
	token := this.Ctx.Input.Param(":token")
	log.Println("token =", token)
	key := g.Config().Api.Key
	authUrl := g.Config().Api.Access + "/" + token + "/" + key

	nodes := sendHttpGetRequest(authUrl)
	if nodes == nil {
		nodes = sendHttpGetRequest(authUrl)
	}
	log.Println("nodes =", nodes)

	var userInfo = make(map[string]string)
	userInfo["username"] = ""
	userInfo["email"] = ""
	userInfo["access_key"] = ""
	if nodes != nil {
		setUserInfo(nodes, userInfo)
	}
	log.Println("userInfo =", userInfo)

	username := userInfo["username"]
	if len(username) > 0 {
		access_key := userInfo["access_key"]
		user := ReadUserByName(username)
		if user == nil { // create third party user
			InsertRegisterUser(username, "", "")
			user = ReadUserByName(username)
		}
		if len(user.Passwd) == 0 {
			role := getUserRole(access_key)
			if role < 1 {
				role = getUserRole(access_key)
			}
			email := userInfo["email"]
			user.Email = email
			user.Role = role
			user.Update()
		}
		maxAge := 3600 * 24 * 30
		this.Ctx.SetCookie("token", access_key, maxAge, "/")
		this.Ctx.SetCookie("token", access_key, maxAge, "/", g.Config().Http.Cookie)

		appSig := this.GetString("sig", "")
		callback := this.GetString("callback", "")
		if appSig != "" && callback != "" {
			SaveSessionAttrs(user.Id, appSig, int(time.Now().Unix())+3600*24*30)
		} else {
			this.CreateSession(user.Id, 3600*24*30)
		}
		this.Redirect("/me/info", 302)
	} else {
		// not logged in. redirect to login page.
		appSig := this.GetString("sig", "")
		callback := this.GetString("callback", "")
		this.renderLoginPage(appSig, callback)
	}
}
<commit_msg>[OWL-1063][fe] refine func LoginWithToken()<commit_after>package uic

import (
	"encoding/base64"
	"encoding/json"
	"github.com/Cepave/open-falcon-backend/modules/fe/g"
	"github.com/Cepave/open-falcon-backend/modules/fe/http/base"
	. "github.com/Cepave/open-falcon-backend/modules/fe/model/uic"
	"github.com/Cepave/open-falcon-backend/modules/fe/utils"
	log "github.com/Sirupsen/logrus"
	"github.com/toolkits/str"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"
)

type AuthController struct {
	base.BaseController
}

func (this *AuthController) Logout() {
	u := this.Ctx.Input.GetData("CurrentUser").(*User)
	token := this.Ctx.GetCookie("token")
	if len(token) > 0 {
		url := g.Config().Api.Logout + "/" + token
		log.Println("logout url =", url)
		result := sendHttpGetRequest(url)
		log.Println("logout result =", result)
		this.Ctx.SetCookie("token", "", 0, "/")
		this.Ctx.SetCookie("token", "", 0, "/", g.Config().Http.Cookie)
	}
	RemoveSessionByUid(u.Id)
	this.Ctx.SetCookie("sig", "", 0, "/")
	this.Ctx.SetCookie("sig", "", 0, "/", g.Config().Http.Cookie)
	this.Ctx.SetCookie("name", "", 0, "/")
	this.Ctx.SetCookie("name", "", 0, "/", g.Config().Http.Cookie)
	this.Redirect("/auth/login", 302)
}

func (this *AuthController) LoginGet() {
	appSig := this.GetString("sig", "")
	callback := this.GetString("callback", "")

	cookieSig := this.Ctx.GetCookie("sig")
	if cookieSig == "" {
		this.renderLoginPage(appSig, callback)
		return
	}

	sessionObj := ReadSessionBySig(cookieSig)
	if sessionObj == nil {
		this.renderLoginPage(appSig, callback)
		return
	}

	if int64(sessionObj.Expired) < time.Now().Unix() {
		RemoveSessionByUid(sessionObj.Uid)
		this.renderLoginPage(appSig, callback)
		return
	}

	if appSig != "" && callback != "" {
		this.Redirect(callback, 302)
	} else {
		this.Redirect("/me/info", 302)
	}
}

func (this *AuthController) LoginPost() {
	name := this.GetString("name", "")
	password := this.GetString("password", "")
	if name == "" || password == "" {
		this.ServeErrJson("name or password is blank")
		return
	}

	var u *User

	ldapEnabled := this.MustGetBool("ldap", false)

	if ldapEnabled {
		sucess, err := utils.LdapBind(g.Config().Ldap.Addr,
			g.Config().Ldap.BaseDN,
			g.Config().Ldap.BindDN,
			g.Config().Ldap.BindPasswd,
			g.Config().Ldap.UserField,
			name,
			password)
		if err != nil {
			this.ServeErrJson(err.Error())
			return
		}

		if !sucess {
			this.ServeErrJson("name or password error")
			return
		}

		user_attributes, err := utils.Ldapsearch(g.Config().Ldap.Addr,
			g.Config().Ldap.BaseDN,
			g.Config().Ldap.BindDN,
			g.Config().Ldap.BindPasswd,
			g.Config().Ldap.UserField,
			name,
			g.Config().Ldap.Attributes)
		userSn := ""
		userMail := ""
		userTel := ""
		if err == nil {
			userSn = user_attributes["sn"]
			userMail = user_attributes["mail"]
			userTel = user_attributes["telephoneNumber"]
		}

		arr := strings.Split(name, "@")
		var userName, userEmail string
		if len(arr) == 2 {
			userName = arr[0]
			userEmail = name
		} else {
			userName = name
			userEmail = userMail
		}

		u = ReadUserByName(userName)
		if u == nil {
			// 说明用户不存在
			u = &User{
				Name:   userName,
				Passwd: "",
				Cnname: userSn,
				Phone:  userTel,
				Email:  userEmail,
			}
			_, err = u.Save()
			if err != nil {
				this.ServeErrJson("insert user fail " + err.Error())
				return
			}
		}
	} else {
		u = ReadUserByName(name)
		if u == nil {
			this.ServeErrJson("no such user")
			return
		}

		if u.Passwd != str.Md5Encode(g.Config().Salt+password) {
			this.ServeErrJson("password error")
			return
		}
	}

	appSig := this.GetString("sig", "")
	callback := this.GetString("callback", "")
	if appSig != "" && callback != "" {
		SaveSessionAttrs(u.Id, appSig, int(time.Now().Unix())+3600*24*30)
	} else {
		this.CreateSession(u.Id, 3600*24*30)
	}

	this.ServeDataJson(callback)
}

func (this *AuthController) renderLoginPage(sig, callback string) {
	this.Data["CanRegister"] = g.Config().CanRegister
	this.Data["LdapEnabled"] = g.Config().Ldap.Enabled
	this.Data["Sig"] = sig
	this.Data["Callback"] = callback
	this.Data["Shortcut"] = g.Config().Shortcut
	this.TplName = "auth/login.html"
}

func (this *AuthController) RegisterGet() {
	this.Data["CanRegister"] = g.Config().CanRegister
	this.Data["Shortcut"] = g.Config().Shortcut
	this.TplName = "auth/register.html"
}

func (this *AuthController) RegisterPost() {
	if !g.Config().CanRegister {
		this.ServeErrJson("registration system is not open")
		return
	}

	name := strings.TrimSpace(this.GetString("name", ""))
	password := strings.TrimSpace(this.GetString("password", ""))
	repeatPassword := strings.TrimSpace(this.GetString("repeat_password", ""))
	email := strings.TrimSpace(this.GetString("email", ""))

	if password != repeatPassword {
		this.ServeErrJson("password not equal the repeart one")
		return
	}

	if !utils.IsUsernameValid(name) {
		this.ServeErrJson("name pattern is invalid")
		return
	}

	if ReadUserIdByName(name) > 0 {
		this.ServeErrJson("name is already existent")
		return
	}

	lastId, err := InsertRegisterUser(name, str.Md5Encode(g.Config().Salt+password), email)
	if err != nil {
		this.ServeErrJson("insert user fail " + err.Error())
		return
	}

	this.CreateSession(lastId, 3600*24*30)

	this.ServeOKJson()
}

func (this *AuthController) CreateSession(uid int64, maxAge int) int {
	sig := utils.GenerateUUID()
	user := SelectUserById(uid)
	expired := int(time.Now().Unix()) + maxAge
	SaveSessionAttrs(uid, sig, expired)
	this.Ctx.SetCookie("sig", sig, maxAge, "/")
	this.Ctx.SetCookie("sig", sig, maxAge, "/", g.Config().Http.Cookie)
	this.Ctx.SetCookie("name", user.Name, maxAge, "/")
	this.Ctx.SetCookie("name", user.Name, maxAge, "/", g.Config().Http.Cookie)
	return expired
}

/**
 * @function name:   func (this *AuthController) LoginThirdParty()
 * @description:     This function returns third party login URL.
 * @related issues:  OWL-206
 * @param:           void
 * @return:          void
 * @author:          Don Hsieh
 * @since:           12/17/2015
 * @last modified:   12/17/2015
 * @called by:       beego.Router("/auth/third-party", &AuthController{}, "post:LoginThirdParty")
 *                    in fe/http/uic/uic_routes.go
 */
func (this *AuthController) LoginThirdParty() {
	s := g.Config().Api.Redirect
	s = base64.StdEncoding.EncodeToString([]byte(s))
	strEncoded := url.QueryEscape(s)
	loginUrl := g.Config().Api.Login + "/" + strEncoded
	this.ServeDataJson(loginUrl)
}

/**
 * @function name:   func sendHttpGetRequest(url string) map[string]interface{}
 * @description:     This function sends GET request to given URL.
 * @related issues:  OWL-206, OWL-159
 * @param:           url string
 * @return:          map[string]interface{}
 * @author:          Don Hsieh
 * @since:           12/17/2015
 * @last modified:   12/17/2015
 * @called by:       func (this *AuthController) LoginWithToken()
 *                    in fe/http/uic/auth_controller.go
 */
func sendHttpGetRequest(url string) map[string]interface{} {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		log.Println("Error =", err.Error())
	}

	client := &http.Client{}
	resp, err := client.Do(req)
	defer resp.Body.Close()
	if err != nil {
		log.Println("Error =", err.Error())
		return nil
	} else {
		body, _ := ioutil.ReadAll(resp.Body)
		var nodes = make(map[string]interface{})
		if err := json.Unmarshal(body, &nodes); err != nil {
			log.Println("Error =", err.Error())
			return nil
		}
		return nodes
	}
}

func setUserInfo(nodes map[string]interface{}, userInfo map[string]string) {
	if status, ok := nodes["status"]; ok {
		if int(status.(float64)) == 1 {
			data := nodes["data"].(map[string]interface{})
			access_key := data["access_key"].(string)
			username := data["username"].(string)
			email := data["email"].(string)
			log.Println("access_key =", access_key)
			userInfo["username"] = username
			userInfo["email"] = email
			userInfo["access_key"] = access_key
		}
	}
}

func getUserRole(access_key string) int {
	urlRole := g.Config().Api.Role + "/" + access_key
	nodes := sendHttpGetRequest(urlRole)
	role := -1
	if int(nodes["status"].(float64)) == 1 {
		permission := nodes["data"]
		log.Println("permission =", permission)
		if permission == "admin" {
			role = 0
		} else if permission == "operator" {
			role = 1
		} else if permission == "observer" {
			role = 2
		} else if permission == "deny" {
			role = 3
		}
	}

	// TODO: The role should be able to be changed on BOSS in the future.
	role = 1
	return role
}

/**
 * @function name:   func (this *AuthController) LoginWithToken()
 * @description:     This function logins user with third party token.
 * @related issues:  OWL-247, OWL-206
 * @param:           void
 * @return:          void
 * @author:          Don Hsieh
 * @since:           12/16/2015
 * @last modified:   01/08/2016
 * @called by:       beego.Router("/auth/login/:token", &AuthController{}, "get:LoginWithToken")
 *                    in fe/http/uic/uic_routes.go
 */
func (this *AuthController) LoginWithToken() {
	log.Println("func (this *AuthController) LoginWithToken()")
	token := this.Ctx.Input.Param(":token")
	log.Println("token =", token)
	key := g.Config().Api.Key
	authUrl := g.Config().Api.Access + "/" + token + "/" + key

	nodes := sendHttpGetRequest(authUrl)
	if nodes == nil {
		nodes = sendHttpGetRequest(authUrl)
	}
	log.Println("nodes =", nodes)

	var userInfo = make(map[string]string)
	userInfo["username"] = ""
	userInfo["email"] = ""
	userInfo["access_key"] = ""
	if nodes != nil {
		setUserInfo(nodes, userInfo)
	}
	log.Println("userInfo =", userInfo)

	username := userInfo["username"]
	if len(username) > 0 {
		access_key := userInfo["access_key"]
		user := ReadUserByName(username)
		if user == nil { // create third party user
			InsertRegisterUser(username, "", "")
			user = ReadUserByName(username)
		}
		if len(user.Passwd) == 0 {
			role := getUserRole(access_key)
			if role < 1 {
				role = getUserRole(access_key)
			}
			email := userInfo["email"]
			user.Email = email
			user.Role = role
			user.Update()
		}
		maxAge := 3600 * 24 * 30
		this.Ctx.SetCookie("token", access_key, maxAge, "/")
		this.Ctx.SetCookie("token", access_key, maxAge, "/", g.Config().Http.Cookie)

		appSig := this.GetString("sig", "")
		callback := this.GetString("callback", "")
		if appSig != "" && callback != "" {
			SaveSessionAttrs(user.Id, appSig, int(time.Now().Unix())+3600*24*30)
		} else {
			this.CreateSession(user.Id, 3600*24*30)
		}
		if callback != "" {
			this.Redirect(callback, 302)
		} else {
			this.Redirect("/me/info", 302)
		}
	} else {
		// not logged in. redirect to login page.
		appSig := this.GetString("sig", "")
		callback := this.GetString("callback", "")
		this.renderLoginPage(appSig, callback)
	}
}
<|endoftext|>


<commit_before>package main

// Filter RDB file per spec: https://github.com/sripathikrishnan/redis-rdb-tools/wiki/Redis-RDB-Dump-File-Format

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"io"
	"strconv"
)

const (
	rdbOpDB         = 0xFE
	rdbOpExpirySec  = 0xFD
	rdbOpExpiryMSec = 0xFC
	rdbOpEOF        = 0xFF

	rdbLen6Bit  = 0x0
	rdbLen14bit = 0x1
	rdbLen32Bit = 0x2
	rdbLenEnc   = 0x3

	rdbOpString    = 0x00
	rdbOpList      = 0x01
	rdbOpSet       = 0x02
	rdbOpZset      = 0x03
	rdbOpHash      = 0x04
	rdbOpZipmap    = 0x09
	rdbOpZiplist   = 0x0a
	rdbOpIntset    = 0x0b
	rdbOpSortedSet = 0x0c
	rdbOpHashmap   = 0x0d
)

var (
	rdbSignature = []byte{0x52, 0x45, 0x44, 0x49, 0x53}
)

var (
	// ErrWrongSignature is returned when RDB signature can't be parsed
	ErrWrongSignature = errors.New("rdb: wrong signature")
	// ErrVersionUnsupported is returned when RDB version is too high (can't parse)
	ErrVersionUnsupported = errors.New("rdb: version unsupported")
	// ErrUnsupportedOp is returned when unsupported operation is encountered in RDB
	ErrUnsupportedOp = errors.New("rdb: unsupported opcode")
	// ErrUnsupportedStringEnc is returned when unsupported string encoding is encountered in RDB
	ErrUnsupportedStringEnc = errors.New("rdb: unsupported string encoding")
)

// RDBFilter holds internal state of RDB filter while running
type RDBFilter struct {
	reader         *bufio.Reader
	output         chan<- []byte
	dissector      func(string) bool
	originalLength int64
	length         int64
	hash           uint64
	saved          []byte
	rdbVersion     int
	valueState     state
	shouldKeep     bool
	currentOp      byte
}

type state func(filter *RDBFilter) (nextstate state, err error)

// FilterRDB filters RDB file which is read from reader, sending chunks of data through output channel
// dissector function is applied to keys to check whether item should be kept or skipped
// length is original length of RDB file
func FilterRDB(reader *bufio.Reader, output chan<- []byte, dissector func(string) bool, length int64) (err error) {
	filter := &RDBFilter{
		reader:         reader,
		output:         output,
		dissector:      dissector,
		originalLength: length,
		shouldKeep:     true,
	}

	state := stateMagic

	for state != nil {
		state, err = state(filter)
		if err != nil {
			return
		}
	}

	return nil
}

// Read exactly n bytes
func (filter *RDBFilter) safeRead(n uint32) (result []byte, err error) {
	result = make([]byte, n)
	_, err = io.ReadFull(filter.reader, result)
	return
}

// Accumulate some data that might be either filtered out or passed through
func (filter *RDBFilter) write(data []byte) {
	if !filter.shouldKeep {
		return
	}

	if filter.saved == nil {
		filter.saved = make([]byte, len(data), 4096)
		copy(filter.saved, data)
	} else {
		filter.saved = append(filter.saved, data...)
	}
}

// Discard or keep saved data
func (filter *RDBFilter) keepOrDiscard() {
	if filter.shouldKeep && filter.saved != nil {
		filter.output <- filter.saved
		filter.hash = CRC64Update(filter.hash, filter.saved)
		filter.length += int64(len(filter.saved))
	}
	filter.saved = nil
	filter.shouldKeep = true
}

// Read length encoded prefix
func (filter *RDBFilter) readLength() (length uint32, encoding int8, err error) {
	prefix, err := filter.reader.ReadByte()
	if err != nil {
		return 0, 0, err
	}
	filter.write([]byte{prefix})

	kind := (prefix & 0xC0) >> 6

	switch kind {
	case rdbLen6Bit:
		length = uint32(prefix & 0x3F)
		return length, -1, nil
	case rdbLen14bit:
		data, err := filter.reader.ReadByte()
		if err != nil {
			return 0, 0, err
		}
		filter.write([]byte{data})
		length = ((uint32(prefix) & 0x3F) << 8) | uint32(data)
		return length, -1, nil
	case rdbLen32Bit:
		data, err := filter.safeRead(4)
		if err != nil {
			return 0, 0, err
		}
		filter.write(data)
		length = binary.BigEndian.Uint32(data)
		return length, -1, nil
	case rdbLenEnc:
		encoding = int8(prefix & 0x3F)
		return 0, encoding, nil
	}
	panic("never reached")
}

// read string from RDB, only uncompressed version is supported
func (filter *RDBFilter) readString() (string, error) {
	length, encoding, err := filter.readLength()
	if err != nil {
		return "", err
	}

	if encoding != -1 {
		return "", ErrUnsupportedStringEnc
	}

	data, err := filter.safeRead(length)
	if err != nil {
		return "", err
	}
	filter.write(data)
	return string(data), nil
}

// skip (copy) string from RDB
func (filter *RDBFilter) skipString() error {
	length, encoding, err := filter.readLength()
	if err != nil {
		return err
	}

	switch encoding {
	// length-prefixed string
	case -1:
		data, err := filter.safeRead(length)
		if err != nil {
			return err
		}
		filter.write(data)
	// integer as string
	case 0, 1, 2:
		data, err := filter.safeRead(1 << uint8(encoding))
		if err != nil {
			return err
		}
		filter.write(data)
	// compressed string
	case 3:
		clength, _, err := filter.readLength()
		if err != nil {
			return err
		}
		_, _, err = filter.readLength()
		if err != nil {
			return err
		}
		data, err := filter.safeRead(clength)
		if err != nil {
			return err
		}
		filter.write(data)
	default:
		return ErrUnsupportedStringEnc
	}
	return nil
}

// read RDB magic header
func stateMagic(filter *RDBFilter) (state, error) {
	signature, err := filter.safeRead(5)
	if err != nil {
		return nil, err
	}
	if bytes.Compare(signature, rdbSignature) != 0 {
		return nil, ErrWrongSignature
	}
	filter.write(signature)

	versionRaw, err := filter.safeRead(4)
	if err != nil {
		return nil, err
	}
	version, err := strconv.Atoi(string(versionRaw))
	if err != nil {
		return nil, ErrWrongSignature
	}

	if version > 6 {
		return nil, ErrVersionUnsupported
	}

	filter.rdbVersion = version
	filter.write(versionRaw)
	filter.keepOrDiscard()

	return stateOp, nil
}

// main selector of operations
func stateOp(filter *RDBFilter) (state, error) {
	op, err := filter.reader.ReadByte()
	if err != nil {
		return nil, err
	}
	filter.currentOp = op

	switch op {
	case rdbOpDB:
		filter.keepOrDiscard()
		return stateDB, nil
	case rdbOpExpirySec:
		return stateExpirySec, nil
	case rdbOpExpiryMSec:
		return stateExpiryMSec, nil
	case rdbOpString, rdbOpZipmap, rdbOpZiplist, rdbOpIntset, rdbOpSortedSet, rdbOpHashmap:
		filter.valueState = stateSkipString
		return stateKey, nil
	case rdbOpList, rdbOpSet:
		filter.valueState = stateSkipSetOrList
		return stateKey, nil
	case rdbOpZset:
		filter.valueState = stateSkipZset
		return stateKey, nil
	case rdbOpHash:
		filter.valueState = stateSkipHash
		return stateKey, nil
	case rdbOpEOF:
		filter.keepOrDiscard()
		filter.write([]byte{rdbOpEOF})
		filter.keepOrDiscard()
		if filter.rdbVersion > 4 {
			return stateCRC64, nil
		}
		return statePadding, nil
	default:
		return nil, ErrUnsupportedOp
	}
}

// DB index operation
func stateDB(filter *RDBFilter) (state, error) {
	filter.write([]byte{rdbOpDB})
	_, _, err := filter.readLength()
	if err != nil {
		return nil, err
	}
	filter.keepOrDiscard()

	return stateOp, nil
}

func stateExpirySec(filter *RDBFilter) (state, error) {
	expiry, err := filter.safeRead(4)
	if err != nil {
		return nil, err
	}

	filter.write([]byte{rdbOpExpirySec})
	filter.write(expiry)

	return stateOp, nil
}

func stateExpiryMSec(filter *RDBFilter) (state, error) {
	expiry, err := filter.safeRead(8)
	if err != nil {
		return nil, err
	}

	filter.write([]byte{rdbOpExpiryMSec})
	filter.write(expiry)

	return stateOp, nil
}

// read key
func stateKey(filter *RDBFilter) (state, error) {
	filter.write([]byte{filter.currentOp})
	key, err := filter.readString()
	if err != nil {
		return nil, err
	}

	filter.shouldKeep = filter.dissector(key)

	return filter.valueState, nil
}

// skip over string
func stateSkipString(filter *RDBFilter) (state, error) {
	err := filter.skipString()
	if err != nil {
		return nil, err
	}

	filter.keepOrDiscard()
	return stateOp, nil
}

// skip over set or list
func stateSkipSetOrList(filter *RDBFilter) (state, error) {
	length, _, err := filter.readLength()
	if err != nil {
		return nil, err
	}

	var i uint32

	for i = 0; i < length; i++ {
		// list element
		err = filter.skipString()
		if err != nil {
			return nil, err
		}
	}

	filter.keepOrDiscard()
	return stateOp, nil
}

// skip over hash
func stateSkipHash(filter *RDBFilter) (state, error) {
	length, _, err := filter.readLength()
	if err != nil {
		return nil, err
	}

	var i uint32

	for i = 0; i < length; i++ {
		// key
		err = filter.skipString()
		if err != nil {
			return nil, err
		}

		// value
		err = filter.skipString()
		if err != nil {
			return nil, err
		}
	}

	filter.keepOrDiscard()
	return stateOp, nil
}

// skip over zset
func stateSkipZset(filter *RDBFilter) (state, error) {
	length, _, err := filter.readLength()
	if err != nil {
		return nil, err
	}

	var i uint32

	for i = 0; i < length; i++ {
		err = filter.skipString()
		if err != nil {
			return nil, err
		}

		dlen, err := filter.reader.ReadByte()
		if err != nil {
			return nil, err
		}
		filter.write([]byte{dlen})

		if dlen < 0xFD {
			double, err := filter.safeRead(uint32(dlen))
			if err != nil {
				return nil, err
			}

			filter.write(double)
		}
	}

	filter.keepOrDiscard()
	return stateOp, nil
}

// re-calculate crc64
func stateCRC64(filter *RDBFilter) (state, error) {
	_, err := filter.safeRead(8)
	if err != nil {
		return nil, err
	}

	var buf = make([]byte, 8)

	binary.LittleEndian.PutUint64(buf, filter.hash)
	filter.output <- buf
	filter.length += 8

	return statePadding, nil
}

// pad RDB with 0xFF up to original length
func statePadding(filter *RDBFilter) (state, error) {
	paddingLength := filter.originalLength - filter.length

	for paddingLength > 0 {
		var padding []byte
		if paddingLength > 4096 {
			padding = make([]byte, 4096)
		} else {
			padding = make([]byte, paddingLength)
		}
		for i := range padding {
			padding[i] = 0xFF
		}
		paddingLength -= int64(len(padding))
		filter.output <- padding
	}
	return nil, nil
}
<commit_msg>Reuse block of padding instead of generating new one each time.<commit_after>package main

// Filter RDB file per spec: https://github.com/sripathikrishnan/redis-rdb-tools/wiki/Redis-RDB-Dump-File-Format

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"io"
	"strconv"
)

const (
	rdbOpDB         = 0xFE
	rdbOpExpirySec  = 0xFD
	rdbOpExpiryMSec = 0xFC
	rdbOpEOF        = 0xFF

	rdbLen6Bit  = 0x0
	rdbLen14bit = 0x1
	rdbLen32Bit = 0x2
	rdbLenEnc   = 0x3

	rdbOpString    = 0x00
	rdbOpList      = 0x01
	rdbOpSet       = 0x02
	rdbOpZset      = 0x03
	rdbOpHash      = 0x04
	rdbOpZipmap    = 0x09
	rdbOpZiplist   = 0x0a
	rdbOpIntset    = 0x0b
	rdbOpSortedSet = 0x0c
	rdbOpHashmap   = 0x0d
)

var (
	rdbSignature = []byte{0x52, 0x45, 0x44, 0x49, 0x53}
)

var (
	// ErrWrongSignature is returned when RDB signature can't be parsed
	ErrWrongSignature = errors.New("rdb: wrong signature")
	// ErrVersionUnsupported is returned when RDB version is too high (can't parse)
	ErrVersionUnsupported = errors.New("rdb: version unsupported")
	// ErrUnsupportedOp is returned when unsupported operation is encountered in RDB
	ErrUnsupportedOp = errors.New("rdb: unsupported opcode")
	// ErrUnsupportedStringEnc is returned when unsupported string encoding is encountered in RDB
	ErrUnsupportedStringEnc = errors.New("rdb: unsupported string encoding")
)

// RDBFilter holds internal state of RDB filter while running
type RDBFilter struct {
	reader         *bufio.Reader
	output         chan<- []byte
	dissector      func(string) bool
	originalLength int64
	length         int64
	hash           uint64
	saved          []byte
	rdbVersion     int
	valueState     state
	shouldKeep     bool
	currentOp      byte
}

type state func(filter *RDBFilter) (nextstate state, err error)

// FilterRDB filters RDB file which is read from reader, sending chunks of data through output channel
// dissector function is applied to keys to check whether item should be kept or skipped
// length is original length of RDB file
func FilterRDB(reader *bufio.Reader, output chan<- []byte, dissector func(string) bool, length int64) (err error) {
	filter := &RDBFilter{
		reader:         reader,
		output:         output,
		dissector:      dissector,
		originalLength: length,
		shouldKeep:     true,
	}

	state := stateMagic

	for state != nil {
		state, err = state(filter)
		if err != nil {
			return
		}
	}

	return nil
}

// Read exactly n bytes
func (filter *RDBFilter) safeRead(n uint32) (result []byte, err error) {
	result = make([]byte, n)
	_, err = io.ReadFull(filter.reader, result)
	return
}

// Accumulate some data that might be either filtered out or passed through
func (filter *RDBFilter) write(data []byte) {
	if !filter.shouldKeep {
		return
	}

	if filter.saved == nil {
		filter.saved = make([]byte, len(data), 4096)
		copy(filter.saved, data)
	} else {
		filter.saved = append(filter.saved, data...)
	}
}

// Discard or keep saved data
func (filter *RDBFilter) keepOrDiscard() {
	if filter.shouldKeep && filter.saved != nil {
		filter.output <- filter.saved
		filter.hash = CRC64Update(filter.hash, filter.saved)
		filter.length += int64(len(filter.saved))
	}
	filter.saved = nil
	filter.shouldKeep = true
}

// Read length encoded prefix
func (filter *RDBFilter) readLength() (length uint32, encoding int8, err error) {
	prefix, err := filter.reader.ReadByte()
	if err != nil {
		return 0, 0, err
	}
	filter.write([]byte{prefix})

	kind := (prefix & 0xC0) >> 6

	switch kind {
	case rdbLen6Bit:
		length = uint32(prefix & 0x3F)
		return length, -1, nil
	case rdbLen14bit:
		data, err := filter.reader.ReadByte()
		if err != nil {
			return 0, 0, err
		}
		filter.write([]byte{data})
		length = ((uint32(prefix) & 0x3F) << 8) | uint32(data)
		return length, -1, nil
	case rdbLen32Bit:
		data, err := filter.safeRead(4)
		if err != nil {
			return 0, 0, err
		}
		filter.write(data)
		length = binary.BigEndian.Uint32(data)
		return length, -1, nil
	case rdbLenEnc:
		encoding = int8(prefix & 0x3F)
		return 0, encoding, nil
	}
	panic("never reached")
}

// read string from RDB, only uncompressed version is supported
func (filter *RDBFilter) readString() (string, error) {
	length, encoding, err := filter.readLength()
	if err != nil {
		return "", err
	}

	if encoding != -1 {
		return "", ErrUnsupportedStringEnc
	}

	data, err := filter.safeRead(length)
	if err != nil {
		return "", err
	}
	filter.write(data)
	return string(data), nil
}

// skip (copy) string from RDB
func (filter *RDBFilter) skipString() error {
	length, encoding, err := filter.readLength()
	if err != nil {
		return err
	}

	switch encoding {
	// length-prefixed string
	case -1:
		data, err := filter.safeRead(length)
		if err != nil {
			return err
		}
		filter.write(data)
	// integer as string
	case 0, 1, 2:
		data, err := filter.safeRead(1 << uint8(encoding))
		if err != nil {
			return err
		}
		filter.write(data)
	// compressed string
	case 3:
		clength, _, err := filter.readLength()
		if err != nil {
			return err
		}
		_, _, err = filter.readLength()
		if err != nil {
			return err
		}
		data, err := filter.safeRead(clength)
		if err != nil {
			return err
		}
		filter.write(data)
	default:
		return ErrUnsupportedStringEnc
	}
	return nil
}

// read RDB magic header
func stateMagic(filter *RDBFilter) (state, error) {
	signature, err := filter.safeRead(5)
	if err != nil {
		return nil, err
	}
	if bytes.Compare(signature, rdbSignature) != 0 {
		return nil, ErrWrongSignature
	}
	filter.write(signature)

	versionRaw, err := filter.safeRead(4)
	if err != nil {
		return nil, err
	}
	version, err := strconv.Atoi(string(versionRaw))
	if err != nil {
		return nil, ErrWrongSignature
	}

	if version > 6 {
		return nil, ErrVersionUnsupported
	}

	filter.rdbVersion = version
	filter.write(versionRaw)
	filter.keepOrDiscard()

	return stateOp, nil
}

// main selector of operations
func stateOp(filter *RDBFilter) (state, error) {
	op, err := filter.reader.ReadByte()
	if err != nil {
		return nil, err
	}
	filter.currentOp = op

	switch op {
	case rdbOpDB:
		filter.keepOrDiscard()
		return stateDB, nil
	case rdbOpExpirySec:
		return stateExpirySec, nil
	case rdbOpExpiryMSec:
		return stateExpiryMSec, nil
	case rdbOpString, rdbOpZipmap, rdbOpZiplist, rdbOpIntset, rdbOpSortedSet, rdbOpHashmap:
		filter.valueState = stateSkipString
		return stateKey, nil
	case rdbOpList, rdbOpSet:
		filter.valueState = stateSkipSetOrList
		return stateKey, nil
	case rdbOpZset:
		filter.valueState = stateSkipZset
		return stateKey, nil
	case rdbOpHash:
		filter.valueState = stateSkipHash
		return stateKey, nil
	case rdbOpEOF:
		filter.keepOrDiscard()
		filter.write([]byte{rdbOpEOF})
		filter.keepOrDiscard()
		if filter.rdbVersion > 4 {
			return stateCRC64, nil
		}
		return statePadding, nil
	default:
		return nil, ErrUnsupportedOp
	}
}

// DB index operation
func stateDB(filter *RDBFilter) (state, error) {
	filter.write([]byte{rdbOpDB})
	_, _, err := filter.readLength()
	if err != nil {
		return nil, err
	}
	filter.keepOrDiscard()

	return stateOp, nil
}

func stateExpirySec(filter *RDBFilter) (state, error) {
	expiry, err := filter.safeRead(4)
	if err != nil {
		return nil, err
	}

	filter.write([]byte{rdbOpExpirySec})
	filter.write(expiry)

	return stateOp, nil
}

func stateExpiryMSec(filter *RDBFilter) (state, error) {
	expiry, err := filter.safeRead(8)
	if err != nil {
		return nil, err
	}

	filter.write([]byte{rdbOpExpiryMSec})
	filter.write(expiry)

	return stateOp, nil
}

// read key
func stateKey(filter *RDBFilter) (state, error) {
	filter.write([]byte{filter.currentOp})
	key, err := filter.readString()
	if err != nil {
		return nil, err
	}

	filter.shouldKeep = filter.dissector(key)

	return filter.valueState, nil
}

// skip over string
func stateSkipString(filter *RDBFilter) (state, error) {
	err := filter.skipString()
	if err != nil {
		return nil, err
	}

	filter.keepOrDiscard()
	return stateOp, nil
}

// skip over set or list
func stateSkipSetOrList(filter *RDBFilter) (state, error) {
	length, _, err := filter.readLength()
	if err != nil {
		return nil, err
	}

	var i uint32

	for i = 0; i < length; i++ {
		// list element
		err = filter.skipString()
		if err != nil {
			return nil, err
		}
	}

	filter.keepOrDiscard()
	return stateOp, nil
}

// skip over hash
func stateSkipHash(filter *RDBFilter) (state, error) {
	length, _, err := filter.readLength()
	if err != nil {
		return nil, err
	}

	var i uint32

	for i = 0; i < length; i++ {
		// key
		err = filter.skipString()
		if err != nil {
			return nil, err
		}

		// value
		err = filter.skipString()
		if err != nil {
			return nil, err
		}
	}

	filter.keepOrDiscard()
	return stateOp, nil
}

// skip over zset
func stateSkipZset(filter *RDBFilter) (state, error) {
	length, _, err := filter.readLength()
	if err != nil {
		return nil, err
	}

	var i uint32

	for i = 0; i < length; i++ {
		err = filter.skipString()
		if err != nil {
			return nil, err
		}

		dlen, err := filter.reader.ReadByte()
		if err != nil {
			return nil, err
		}
		filter.write([]byte{dlen})

		if dlen < 0xFD {
			double, err := filter.safeRead(uint32(dlen))
			if err != nil {
				return nil, err
			}

			filter.write(double)
		}
	}

	filter.keepOrDiscard()
	return stateOp, nil
}

// re-calculate crc64
func stateCRC64(filter *RDBFilter) (state, error) {
	_, err := filter.safeRead(8)
	if err != nil {
		return nil, err
	}

	var buf = make([]byte, 8)

	binary.LittleEndian.PutUint64(buf, filter.hash)
	filter.output <- buf
	filter.length += 8

	return statePadding, nil
}

// pad RDB with 0xFF up to original length
func statePadding(filter *RDBFilter) (state, error) {
	const paddingSize = 4096

	paddingLength := filter.originalLength - filter.length
	paddingBlock := make([]byte, paddingSize)
	for i := range paddingBlock {
		paddingBlock[i] = 0xFF
	}

	for paddingLength > 0 {
		if paddingLength > paddingSize {
			filter.output <- paddingBlock
			paddingLength -= paddingSize
		} else {
			filter.output <- paddingBlock[:paddingLength]
			break
		}
	}
	return nil, nil
}
<|endoftext|>


<commit_before>package iris

import (
	"net/http"
	"sync"

	"github.com/kataras/go-websocket"
)

// conversionals
const (
	// All is the string which the Emitter use to send a message to all
	All = websocket.All
	// NotMe is the string which the Emitter use to send a message to all except this websocket.Connection
	NotMe = websocket.NotMe
	// Broadcast is the string which the Emitter use to send a message to all except this websocket.Connection, same as 'NotMe'
	Broadcast = websocket.Broadcast
)

// Note I keep this code only to no change the front-end API, we could only use the go-websocket and set our custom upgrader

type (
	// WebsocketServer is the iris websocket server, expose the websocket.Server
	// the below code is a wrapper and bridge between iris-contrib/websocket and kataras/go-websocket
	WebsocketServer struct {
		websocket.Server
		station *Framework
		once    sync.Once
		// Config:
		// if endpoint is not empty then this configuration is used instead of the station's
		// useful when the user/dev wants more than one websocket server inside one iris instance.
		Config WebsocketConfiguration
	}
)

// NewWebsocketServer returns a new empty unitialized websocket server
// it runs on first OnConnection
func NewWebsocketServer(station *Framework) *WebsocketServer {
	return &WebsocketServer{station: station, Server: websocket.New(), Config: DefaultWebsocketConfiguration()}
}

// NewWebsocketServer creates the client side source route and the route path Endpoint with the correct Handler
// receives the websocket configuration and  the iris station
// and returns the websocket server which can be attached to more than one iris station (if needed)
func (ws *WebsocketServer) init() {

	if ws.Config.Endpoint == "" {
		ws.Config = ws.station.Config.Websocket
	}

	c := ws.Config

	if c.Endpoint == "" {
		return
	}

	if c.CheckOrigin == nil {
		c.CheckOrigin = DefaultWebsocketCheckOrigin
	}

	if c.Error == nil {
		c.Error = DefaultWebsocketError
	}
	// set the underline websocket server's configuration
	ws.Server.Set(websocket.Config{
		WriteTimeout:    c.WriteTimeout,
		PongTimeout:     c.PongTimeout,
		PingPeriod:      c.PingPeriod,
		MaxMessageSize:  c.MaxMessageSize,
		BinaryMessages:  c.BinaryMessages,
		ReadBufferSize:  c.ReadBufferSize,
		WriteBufferSize: c.WriteBufferSize,
		Error: func(w http.ResponseWriter, r *http.Request, status int, reason error) {
			ws.station.Context.Run(w, r, func(ctx *Context) {
				c.Error(ctx, status, reason)
			})
		},
		CheckOrigin: c.CheckOrigin,
		IDGenerator: c.IDGenerator,
	})

	// set the routing for client-side source (javascript) (optional)
	clientSideLookupName := "iris-websocket-client-side"
	ws.station.Get(c.Endpoint, ToHandler(ws.Server.Handler()))
	// check if client side already exists
	if ws.station.Routes().Lookup(clientSideLookupName) == nil {
		// serve the client side on domain:port/iris-ws.js
		ws.station.StaticContent("/iris-ws.js", contentJavascript, websocket.ClientSource).ChangeName(clientSideLookupName)
	}
}

// WebsocketConnection is the front-end API that you will use to communicate with the client side
type WebsocketConnection interface {
	websocket.Connection
}

// OnConnection this is the main event you, as developer, will work with each of the websocket connections
func (ws *WebsocketServer) OnConnection(connectionListener func(WebsocketConnection)) {
	ws.once.Do(ws.init)

	ws.Server.OnConnection(func(c websocket.Connection) {
		connectionListener(c)
	})
}
<commit_msg>Set the NewWebsocketServer's configuration to the user's station, these can be changed later. Preparing for writing the tests.<commit_after>package iris

import (
	"net/http"
	"sync"

	"github.com/kataras/go-websocket"
)

// conversionals
const (
	// All is the string which the Emitter use to send a message to all
	All = websocket.All
	// NotMe is the string which the Emitter use to send a message to all except this websocket.Connection
	NotMe = websocket.NotMe
	// Broadcast is the string which the Emitter use to send a message to all except this websocket.Connection, same as 'NotMe'
	Broadcast = websocket.Broadcast
)

// Note I keep this code only to no change the front-end API, we could only use the go-websocket and set our custom upgrader

type (
	// WebsocketServer is the iris websocket server, expose the websocket.Server
	// the below code is a wrapper and bridge between iris-contrib/websocket and kataras/go-websocket
	WebsocketServer struct {
		websocket.Server
		station *Framework
		once    sync.Once
		// Config:
		// if endpoint is not empty then this configuration is used instead of the station's
		// useful when the user/dev wants more than one websocket server inside one iris instance.
		Config WebsocketConfiguration
	}
)

// NewWebsocketServer returns a new empty unitialized websocket server
// it runs on first OnConnection
func NewWebsocketServer(station *Framework) *WebsocketServer {
	return &WebsocketServer{station: station, Server: websocket.New(), Config: station.Config.Websocket}
}

// NewWebsocketServer creates the client side source route and the route path Endpoint with the correct Handler
// receives the websocket configuration and  the iris station
// and returns the websocket server which can be attached to more than one iris station (if needed)
func (ws *WebsocketServer) init() {

	if ws.Config.Endpoint == "" {
		ws.Config = ws.station.Config.Websocket
	}

	c := ws.Config

	if c.Endpoint == "" {
		return
	}

	if c.CheckOrigin == nil {
		c.CheckOrigin = DefaultWebsocketCheckOrigin
	}

	if c.Error == nil {
		c.Error = DefaultWebsocketError
	}
	// set the underline websocket server's configuration
	ws.Server.Set(websocket.Config{
		WriteTimeout:    c.WriteTimeout,
		PongTimeout:     c.PongTimeout,
		PingPeriod:      c.PingPeriod,
		MaxMessageSize:  c.MaxMessageSize,
		BinaryMessages:  c.BinaryMessages,
		ReadBufferSize:  c.ReadBufferSize,
		WriteBufferSize: c.WriteBufferSize,
		Error: func(w http.ResponseWriter, r *http.Request, status int, reason error) {
			ws.station.Context.Run(w, r, func(ctx *Context) {
				c.Error(ctx, status, reason)
			})
		},
		CheckOrigin: c.CheckOrigin,
		IDGenerator: c.IDGenerator,
	})

	// set the routing for client-side source (javascript) (optional)
	clientSideLookupName := "iris-websocket-client-side"
	ws.station.Get(c.Endpoint, ToHandler(ws.Server.Handler()))
	// check if client side already exists
	if ws.station.Routes().Lookup(clientSideLookupName) == nil {
		// serve the client side on domain:port/iris-ws.js
		ws.station.StaticContent("/iris-ws.js", contentJavascript, websocket.ClientSource).ChangeName(clientSideLookupName)
	}
}

// WebsocketConnection is the front-end API that you will use to communicate with the client side
type WebsocketConnection interface {
	websocket.Connection
}

// OnConnection this is the main event you, as developer, will work with each of the websocket connections
func (ws *WebsocketServer) OnConnection(connectionListener func(WebsocketConnection)) {
	ws.once.Do(ws.init)

	ws.Server.OnConnection(func(c websocket.Connection) {
		connectionListener(c)
	})
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"github.com/russross/blackfriday"
	"gopkg.in/yaml.v1"
	"io/ioutil"
	"log"
	"strings"
	"sync"
)

var (
	port = flag.String("port", ":8888",
		"the port to run the release notes aggregator on")
	conf        = flag.String("conf", "./example/repos.yml", "path to the config file")
	confFileDir string
	wg          sync.WaitGroup
)

func main() {
	flag.Parse()
	if len(*conf) == 0 {
		panic("Configuration file is not set.")
	}

	log.Printf("Using configuration file: %s\n", *conf)
	repos := parseConfigFile()

	confString := *conf
	confFileDir = confString[:strings.LastIndex(confString, "/")]

	output := make(chan []byte)

	if len(repos) > 0 {
		getReleaseNotes(repos, output)
		writeToFile(output)
	} else {
		panic("repos is nil")
	}
}

// parseConfigFile parses the configugration yaml file to which the global
// variable points
func parseConfigFile() []string {
	b, err := ioutil.ReadFile(*conf)
	generalErrorHandler(err)

	m := make(map[string][]string)

	err = yaml.Unmarshal(b, &m)
	generalErrorHandler(err)

	if val, ok := m["repos"]; ok {
		return val
	}

	return nil
}

func generalErrorHandler(err error) {
	if err != nil {
		panic(err)
	}
}

func getReleaseNotes(repos []string, output chan []byte) {
	for _, repo := range repos {
		wg.Add(1)
		go func(repo string) {
			var path string

			if repo[0:1] == "/" { // using absolute path
				path = repo
			} else {
				if repo[:1] == "./" { // relative path dot slash
					path = confFileDir + "/" + repo[2:]
				} else { // relative path naked
					path = confFileDir + "/" + repo
				}
			}

			file, err := ioutil.ReadFile(path + "/RELEASE.md")
			generalErrorHandler(err)

			output <- blackfriday.MarkdownBasic(file)
			wg.Done()
		}(repo)
	}

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

func writeToFile(output <-chan []byte) {
	for notes := range output {
		log.Printf("%s\n", notes)
	}
}
<commit_msg>Use templates to write valid html to file for static file server<commit_after>package main

import (
	"flag"
	"github.com/russross/blackfriday"
	"gopkg.in/yaml.v1"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"
)

// ReleaseFile is used to communicate on a channel the name and body
type ReleaseFile struct {
	Name string
	Body template.HTML
}

func (rf ReleaseFile) FormattedName() string {
	if rf.Name[:2] == "./" {
		return rf.Name[2:]
	}
	return rf.Name
}

var (
	port  = flag.String("port", ":8888",
		"the port to run the release notes aggregator on")
	conf        = flag.String("conf", "./example/repos.yml", "path to the config file")
	confFileDir string
	wg          sync.WaitGroup
)

func main() {
	tmpl, err := template.New("default.tmpl.html").ParseFiles("./templates/default.tmpl.html")
	generalErrorHandler(err)

	flag.Parse()
	if len(*conf) == 0 {
		panic("Configuration file is not set.")
	}

	log.Printf("Using configuration file: %s\n", *conf)
	repos := parseConfigFile()

	confString := *conf
	confFileDir = confString[:strings.LastIndex(confString, "/")]

	output := make(chan *ReleaseFile)

	if len(repos) > 0 {
		getReleaseNotes(repos, output)
		writeToFile(output, tmpl)
	} else {
		panic("repos is nil")
	}

	http.HandleFunc("/", RootHandler)
	http.HandleFunc("/static/", StaticHandler)
	http.ListenAndServe(*port, nil)
}

// RootHandler takes care of stuff that comes out of the compiled_html folder
func RootHandler(rw http.ResponseWriter, r *http.Request) {
	http.ServeFile(rw, r, "./compiled_html/"+r.URL.Path[1:])
}

// StaticHandler takes care of the static assets
func StaticHandler(rw http.ResponseWriter, r *http.Request) {
	http.ServeFile(rw, r, r.URL.Path[1:])
}

// parseConfigFile parses the configugration yaml file to which the global
// variable points
func parseConfigFile() []string {
	b, err := ioutil.ReadFile(*conf)
	generalErrorHandler(err)

	m := make(map[string][]string)

	err = yaml.Unmarshal(b, &m)
	generalErrorHandler(err)

	if val, ok := m["repos"]; ok {
		return val
	}

	return nil
}

func generalErrorHandler(err error) {
	if err != nil {
		panic(err)
	}
}

func getReleaseNotes(repos []string, output chan *ReleaseFile) {
	for _, repo := range repos {
		wg.Add(1)
		go func(repo string) {
			var path string

			if repo[0:1] == "/" { // using absolute path
				path = repo
			} else {
				if repo[:1] == "./" { // relative path dot slash
					path = confFileDir + "/" + repo[2:]
				} else { // relative path naked
					path = confFileDir + "/" + repo
				}
			}

			file, err := ioutil.ReadFile(path + "/RELEASE.md")
			generalErrorHandler(err)

			output <- &ReleaseFile{
				Name: repo,
				Body: template.HTML(blackfriday.MarkdownBasic(file)),
			}
			wg.Done()
		}(repo)
	}

	// Start a goroutine to wait for the waitgroup to hit zero, because otherwise
	// we'll get deadlocked goroutines
	go func() {
		wg.Wait()
		close(output)
	}()
}

func writeToFile(output <-chan *ReleaseFile, tmpl *template.Template) {
	for note := range output {
		f, err := os.Create("./compiled_html/"+note.Name+".html")
		generalErrorHandler(err)

		defer f.Close()

		err = tmpl.Execute(f, note)

		generalErrorHandler(err)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2022 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package hetznertasks

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"math/rand"
	"strings"

	"github.com/hetznercloud/hcloud-go/hcloud"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/cloudup/hetzner"
	"k8s.io/kops/upup/pkg/fi/cloudup/terraform"
	"k8s.io/kops/upup/pkg/fi/cloudup/terraformWriter"
)

// +kops:fitask
type ServerGroup struct {
	Name      *string
	Lifecycle fi.Lifecycle
	SSHKeys   []*SSHKey
	Network   *Network

	Count      int
	NeedUpdate []string

	Location string
	Size     string
	Image    string

	EnableIPv4 bool
	EnableIPv6 bool

	UserData fi.Resource

	Labels map[string]string
}

func (v *ServerGroup) Find(c *fi.Context) (*ServerGroup, error) {
	cloud := c.Cloud.(hetzner.HetznerCloud)
	client := cloud.ServerClient()

	labelSelector := []string{
		fmt.Sprintf("%s=%s", hetzner.TagKubernetesClusterName, c.Cluster.Name),
		fmt.Sprintf("%s=%s", hetzner.TagKubernetesInstanceGroup, fi.StringValue(v.Name)),
	}
	listOptions := hcloud.ListOpts{
		PerPage:       50,
		LabelSelector: strings.Join(labelSelector, ","),
	}
	serverListOptions := hcloud.ServerListOpts{ListOpts: listOptions}
	servers, err := client.AllWithOpts(context.TODO(), serverListOptions)
	if err != nil {
		return nil, err
	}
	if len(servers) == 0 {
		return nil, nil
	}

	// Calculate the user-data hash
	userDataBytes, err := fi.ResourceAsBytes(v.UserData)
	if err != nil {
		return nil, err
	}
	userDataHash := safeBytesHash(userDataBytes)

	// Add the expected user-data hash label
	v.Labels[hetzner.TagKubernetesInstanceUserData] = userDataHash

	actual := *v
	actual.Count = len(servers)

	// Find servers that need to be updated
	for _, server := range servers {
		// Ignore servers that are already labeled as needing update
		if _, ok := server.Labels[hetzner.TagKubernetesInstanceNeedsUpdate]; ok {
			continue
		}

		// Check if server matches the expected group template
		if server.Labels[hetzner.TagKubernetesInstanceUserData] != userDataHash {
			actual.NeedUpdate = append(actual.NeedUpdate, server.Name)
			continue
		}
		if server.Datacenter == nil || server.Datacenter.Location == nil || server.Datacenter.Location.Name != v.Location {
			actual.NeedUpdate = append(actual.NeedUpdate, server.Name)
			continue
		}
		if server.ServerType == nil || server.ServerType.Name != v.Size {
			actual.NeedUpdate = append(actual.NeedUpdate, server.Name)
			continue
		}
		if server.Image == nil || server.Image.Name != v.Image {
			actual.NeedUpdate = append(actual.NeedUpdate, server.Name)
			continue
		}
		if (server.PublicNet.IPv4.IP != nil) != v.EnableIPv4 {
			actual.NeedUpdate = append(actual.NeedUpdate, server.Name)
			continue
		}
		if (server.PublicNet.IPv6.IP != nil) != v.EnableIPv6 {
			actual.NeedUpdate = append(actual.NeedUpdate, server.Name)
			continue
		}
	}

	return &actual, nil
}

func (v *ServerGroup) Run(c *fi.Context) error {
	return fi.DefaultDeltaRunMethod(v, c)
}

func (_ *ServerGroup) CheckChanges(a, e, changes *ServerGroup) error {
	if e.Name == nil {
		return fi.RequiredField("Name")
	}
	if e.Location == "" {
		return fi.RequiredField("Location")
	}
	if e.Size == "" {
		return fi.RequiredField("Size")
	}
	if e.Image == "" {
		return fi.RequiredField("Image")
	}
	if e.UserData == nil {
		return fi.RequiredField("UserData")
	}
	return nil
}

func (_ *ServerGroup) RenderHetzner(t *hetzner.HetznerAPITarget, a, e, changes *ServerGroup) error {
	client := t.Cloud.ServerClient()

	if a != nil {
		// Add "kops.k8s.io/needs-update" label to servers needing update
		for _, serverName := range a.NeedUpdate {
			server, _, err := client.GetByName(context.TODO(), serverName)
			if err != nil {
				return err
			}
			if server == nil {
				continue
			}

			server.Labels[hetzner.TagKubernetesInstanceNeedsUpdate] = ""
			_, _, err = client.Update(context.TODO(), server, hcloud.ServerUpdateOpts{
				Name:   server.Name,
				Labels: server.Labels,
			})
			if err != nil {
				return err
			}
		}
	}

	actualCount := 0
	if a != nil {
		actualCount = a.Count
	}
	expectedCount := e.Count

	if actualCount >= expectedCount {
		return nil
	}

	if len(e.SSHKeys) == 0 {
		return fmt.Errorf("failed to find ssh keys for server group %q", fi.StringValue(e.Name))
	}
	if e.Network == nil {
		return fmt.Errorf("failed to find network for server group %q", fi.StringValue(e.Name))
	}

	userData, err := fi.ResourceAsString(e.UserData)
	if err != nil {
		return err
	}
	userDataBytes, err := fi.ResourceAsBytes(e.UserData)
	if err != nil {
		return err
	}
	userDataHash := safeBytesHash(userDataBytes)

	for i := 1; i <= expectedCount-actualCount; i++ {
		// Append a random/unique ID to the node name
		name := fmt.Sprintf("%s-%x", fi.StringValue(e.Name), rand.Int63())

		opts := hcloud.ServerCreateOpts{
			Name:             name,
			StartAfterCreate: fi.Bool(true),
			Networks: []*hcloud.Network{
				{
					ID: fi.IntValue(e.Network.ID),
				},
			},
			Location: &hcloud.Location{
				Name: e.Location,
			},
			ServerType: &hcloud.ServerType{
				Name: e.Size,
			},
			Image: &hcloud.Image{
				Name: e.Image,
			},
			UserData: userData,
			Labels:   e.Labels,
			PublicNet: &hcloud.ServerCreatePublicNet{
				EnableIPv4: e.EnableIPv4,
				EnableIPv6: e.EnableIPv6,
			},
		}

		// Add the SSH keys
		for _, sshkey := range e.SSHKeys {
			opts.SSHKeys = append(opts.SSHKeys, &hcloud.SSHKey{ID: fi.IntValue(sshkey.ID)})
		}

		// Add the user-data hash label
		opts.Labels[hetzner.TagKubernetesInstanceUserData] = userDataHash

		_, _, err = client.Create(context.TODO(), opts)
		if err != nil {
			return err
		}
	}

	return nil
}

func safeBytesHash(data []byte) string {
	// Calculate the SHA256 checksum of the data
	sum256 := sha256.Sum256(data)

	// Replace the unsupported chars with supported ones
	safe256 := base64.StdEncoding.EncodeToString(sum256[:])
	safe256 = strings.ReplaceAll(safe256, "+", "-")
	safe256 = strings.ReplaceAll(safe256, "/", "_")

	// Trim the unsupported "=" padding chars
	safe256 = strings.TrimRight(safe256, "=")

	return fmt.Sprintf("sha256.%s", safe256)
}

type terraformServer struct {
	Count      *int                       `cty:"count"`
	Name       *terraformWriter.Literal   `cty:"name"`
	Location   *string                    `cty:"location"`
	ServerType *string                    `cty:"server_type"`
	Image      *string                    `cty:"image"`
	SSHKeys    []*terraformWriter.Literal `cty:"ssh_keys"`
	Network    []*terraformServerNetwork  `cty:"network"`
	PublicNet  *terraformServerPublicNet  `cty:"public_net"`
	UserData   *terraformWriter.Literal   `cty:"user_data"`
	Labels     map[string]string          `cty:"labels"`
}

type terraformServerNetwork struct {
	ID         *terraformWriter.Literal `cty:"network_id"`
	IP         *string                  `cty:"ip"`
	AliasIPs   *string                  `cty:"alias_ips"`
	MACAddress *string                  `cty:"mac_address"`
}

type terraformServerPublicNet struct {
	EnableIPv4 *bool `cty:"ipv4_enabled"`
	EnableIPv6 *bool `cty:"ipv6_enabled"`
}

func (_ *ServerGroup) RenderTerraform(t *terraform.TerraformTarget, a, e, changes *ServerGroup) error {
	name := terraformWriter.LiteralWithIndex(fi.StringValue(e.Name))
	tf := &terraformServer{
		Count:      fi.Int(e.Count),
		Name:       name,
		Location:   fi.String(e.Location),
		ServerType: fi.String(e.Size),
		Image:      fi.String(e.Image),
		Network: []*terraformServerNetwork{
			{
				ID: e.Network.TerraformLink(),
			},
		},
		PublicNet: &terraformServerPublicNet{
			EnableIPv4: fi.Bool(e.EnableIPv4),
			EnableIPv6: fi.Bool(e.EnableIPv6),
		},
		Labels: e.Labels,
	}

	for _, sshkey := range e.SSHKeys {
		tf.SSHKeys = append(tf.SSHKeys, sshkey.TerraformLink())
	}

	if e.UserData != nil {
		data, err := fi.ResourceAsBytes(e.UserData)
		if err != nil {
			return err
		}
		if data != nil {
			tf.UserData, err = t.AddFileBytes("hcloud_server", fi.StringValue(e.Name), "user_data", data, true)
			if err != nil {
				return err
			}
		}
	}

	return t.RenderResource("hcloud_server", fi.StringValue(e.Name), tf)
}
<commit_msg>Delete the oldest servers when over the desired count for Hetzner<commit_after>/*
Copyright 2022 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package hetznertasks

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"math/rand"
	"strings"

	"github.com/hetznercloud/hcloud-go/hcloud"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/cloudup/hetzner"
	"k8s.io/kops/upup/pkg/fi/cloudup/terraform"
	"k8s.io/kops/upup/pkg/fi/cloudup/terraformWriter"
)

// +kops:fitask
type ServerGroup struct {
	Name      *string
	Lifecycle fi.Lifecycle
	SSHKeys   []*SSHKey
	Network   *Network

	Count      int
	NeedUpdate []string

	Location string
	Size     string
	Image    string

	EnableIPv4 bool
	EnableIPv6 bool

	UserData fi.Resource

	Labels map[string]string
}

func (v *ServerGroup) Find(c *fi.Context) (*ServerGroup, error) {
	cloud := c.Cloud.(hetzner.HetznerCloud)
	client := cloud.ServerClient()

	labelSelector := []string{
		fmt.Sprintf("%s=%s", hetzner.TagKubernetesClusterName, c.Cluster.Name),
		fmt.Sprintf("%s=%s", hetzner.TagKubernetesInstanceGroup, fi.StringValue(v.Name)),
	}
	listOptions := hcloud.ListOpts{
		PerPage:       50,
		LabelSelector: strings.Join(labelSelector, ","),
	}
	serverListOptions := hcloud.ServerListOpts{ListOpts: listOptions}
	servers, err := client.AllWithOpts(context.TODO(), serverListOptions)
	if err != nil {
		return nil, err
	}
	if len(servers) == 0 {
		return nil, nil
	}

	// Calculate the user-data hash
	userDataBytes, err := fi.ResourceAsBytes(v.UserData)
	if err != nil {
		return nil, err
	}
	userDataHash := safeBytesHash(userDataBytes)

	// Add the expected user-data hash label
	v.Labels[hetzner.TagKubernetesInstanceUserData] = userDataHash

	actual := *v
	actual.Count = len(servers)

	// Find servers that need to be updated
	for i, server := range servers {
		// Ignore servers that are already labeled as needing update
		if _, ok := server.Labels[hetzner.TagKubernetesInstanceNeedsUpdate]; ok {
			continue
		}

		// Check if server index is higher than desired count
		if i >= v.Count {
			actual.NeedUpdate = append(actual.NeedUpdate, server.Name)
			continue
		}

		// Check if server matches the expected group template
		if server.Labels[hetzner.TagKubernetesInstanceUserData] != userDataHash {
			actual.NeedUpdate = append(actual.NeedUpdate, server.Name)
			continue
		}
		if server.Datacenter == nil || server.Datacenter.Location == nil || server.Datacenter.Location.Name != v.Location {
			actual.NeedUpdate = append(actual.NeedUpdate, server.Name)
			continue
		}
		if server.ServerType == nil || server.ServerType.Name != v.Size {
			actual.NeedUpdate = append(actual.NeedUpdate, server.Name)
			continue
		}
		if server.Image == nil || server.Image.Name != v.Image {
			actual.NeedUpdate = append(actual.NeedUpdate, server.Name)
			continue
		}
		if (server.PublicNet.IPv4.IP != nil) != v.EnableIPv4 {
			actual.NeedUpdate = append(actual.NeedUpdate, server.Name)
			continue
		}
		if (server.PublicNet.IPv6.IP != nil) != v.EnableIPv6 {
			actual.NeedUpdate = append(actual.NeedUpdate, server.Name)
			continue
		}
	}

	return &actual, nil
}

func (v *ServerGroup) Run(c *fi.Context) error {
	return fi.DefaultDeltaRunMethod(v, c)
}

func (_ *ServerGroup) CheckChanges(a, e, changes *ServerGroup) error {
	if e.Name == nil {
		return fi.RequiredField("Name")
	}
	if e.Location == "" {
		return fi.RequiredField("Location")
	}
	if e.Size == "" {
		return fi.RequiredField("Size")
	}
	if e.Image == "" {
		return fi.RequiredField("Image")
	}
	if e.UserData == nil {
		return fi.RequiredField("UserData")
	}
	return nil
}

func (_ *ServerGroup) RenderHetzner(t *hetzner.HetznerAPITarget, a, e, changes *ServerGroup) error {
	client := t.Cloud.ServerClient()

	if a != nil {
		// Add "kops.k8s.io/needs-update" label to servers needing update
		for _, serverName := range a.NeedUpdate {
			server, _, err := client.GetByName(context.TODO(), serverName)
			if err != nil {
				return err
			}
			if server == nil {
				continue
			}

			server.Labels[hetzner.TagKubernetesInstanceNeedsUpdate] = ""
			_, _, err = client.Update(context.TODO(), server, hcloud.ServerUpdateOpts{
				Name:   server.Name,
				Labels: server.Labels,
			})
			if err != nil {
				return err
			}
		}
	}

	actualCount := 0
	if a != nil {
		actualCount = a.Count
	}
	expectedCount := e.Count

	if actualCount >= expectedCount {
		return nil
	}

	if len(e.SSHKeys) == 0 {
		return fmt.Errorf("failed to find ssh keys for server group %q", fi.StringValue(e.Name))
	}
	if e.Network == nil {
		return fmt.Errorf("failed to find network for server group %q", fi.StringValue(e.Name))
	}

	userData, err := fi.ResourceAsString(e.UserData)
	if err != nil {
		return err
	}
	userDataBytes, err := fi.ResourceAsBytes(e.UserData)
	if err != nil {
		return err
	}
	userDataHash := safeBytesHash(userDataBytes)

	for i := 1; i <= expectedCount-actualCount; i++ {
		// Append a random/unique ID to the node name
		name := fmt.Sprintf("%s-%x", fi.StringValue(e.Name), rand.Int63())

		opts := hcloud.ServerCreateOpts{
			Name:             name,
			StartAfterCreate: fi.Bool(true),
			Networks: []*hcloud.Network{
				{
					ID: fi.IntValue(e.Network.ID),
				},
			},
			Location: &hcloud.Location{
				Name: e.Location,
			},
			ServerType: &hcloud.ServerType{
				Name: e.Size,
			},
			Image: &hcloud.Image{
				Name: e.Image,
			},
			UserData: userData,
			Labels:   e.Labels,
			PublicNet: &hcloud.ServerCreatePublicNet{
				EnableIPv4: e.EnableIPv4,
				EnableIPv6: e.EnableIPv6,
			},
		}

		// Add the SSH keys
		for _, sshkey := range e.SSHKeys {
			opts.SSHKeys = append(opts.SSHKeys, &hcloud.SSHKey{ID: fi.IntValue(sshkey.ID)})
		}

		// Add the user-data hash label
		opts.Labels[hetzner.TagKubernetesInstanceUserData] = userDataHash

		_, _, err = client.Create(context.TODO(), opts)
		if err != nil {
			return err
		}
	}

	return nil
}

func safeBytesHash(data []byte) string {
	// Calculate the SHA256 checksum of the data
	sum256 := sha256.Sum256(data)

	// Replace the unsupported chars with supported ones
	safe256 := base64.StdEncoding.EncodeToString(sum256[:])
	safe256 = strings.ReplaceAll(safe256, "+", "-")
	safe256 = strings.ReplaceAll(safe256, "/", "_")

	// Trim the unsupported "=" padding chars
	safe256 = strings.TrimRight(safe256, "=")

	return fmt.Sprintf("sha256.%s", safe256)
}

type terraformServer struct {
	Count      *int                       `cty:"count"`
	Name       *terraformWriter.Literal   `cty:"name"`
	Location   *string                    `cty:"location"`
	ServerType *string                    `cty:"server_type"`
	Image      *string                    `cty:"image"`
	SSHKeys    []*terraformWriter.Literal `cty:"ssh_keys"`
	Network    []*terraformServerNetwork  `cty:"network"`
	PublicNet  *terraformServerPublicNet  `cty:"public_net"`
	UserData   *terraformWriter.Literal   `cty:"user_data"`
	Labels     map[string]string          `cty:"labels"`
}

type terraformServerNetwork struct {
	ID         *terraformWriter.Literal `cty:"network_id"`
	IP         *string                  `cty:"ip"`
	AliasIPs   *string                  `cty:"alias_ips"`
	MACAddress *string                  `cty:"mac_address"`
}

type terraformServerPublicNet struct {
	EnableIPv4 *bool `cty:"ipv4_enabled"`
	EnableIPv6 *bool `cty:"ipv6_enabled"`
}

func (_ *ServerGroup) RenderTerraform(t *terraform.TerraformTarget, a, e, changes *ServerGroup) error {
	name := terraformWriter.LiteralWithIndex(fi.StringValue(e.Name))
	tf := &terraformServer{
		Count:      fi.Int(e.Count),
		Name:       name,
		Location:   fi.String(e.Location),
		ServerType: fi.String(e.Size),
		Image:      fi.String(e.Image),
		Network: []*terraformServerNetwork{
			{
				ID: e.Network.TerraformLink(),
			},
		},
		PublicNet: &terraformServerPublicNet{
			EnableIPv4: fi.Bool(e.EnableIPv4),
			EnableIPv6: fi.Bool(e.EnableIPv6),
		},
		Labels: e.Labels,
	}

	for _, sshkey := range e.SSHKeys {
		tf.SSHKeys = append(tf.SSHKeys, sshkey.TerraformLink())
	}

	if e.UserData != nil {
		data, err := fi.ResourceAsBytes(e.UserData)
		if err != nil {
			return err
		}
		if data != nil {
			tf.UserData, err = t.AddFileBytes("hcloud_server", fi.StringValue(e.Name), "user_data", data, true)
			if err != nil {
				return err
			}
		}
	}

	return t.RenderResource("hcloud_server", fi.StringValue(e.Name), tf)
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"path"
	"strconv"
	"sync"

	"github.com/MJKWoolnough/minecraft"
)

func (c *Config) Name(_ struct{}, serverName *string) error {
	c.mu.RLock()
	defer c.mu.RUnlock()
	*serverName = c.ServerName
	return nil
}

func (c *Config) List(_ struct{}, list *[]Server) error {
	c.mu.RLock()
	defer c.mu.RUnlock()
	*list = make([]Server, 0, len(c.Servers))
	for _, s := range c.Servers {
		*list = append(*list, s)
	}
	return nil
}

func (c *Config) Save(s Server, _ *struct{}) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	defer c.save()
	ns, ok := c.Servers[s.ID]
	if !ok {
		return ErrNoServer
	}
	s.Path = ns.Path
	s.status = ns.status
	c.Servers[s.ID] = s
	return nil
}

func (c *Config) MapList(_ struct{}, list *[]Map) error {
	c.mu.RLock()
	defer c.mu.RUnlock()
	*list = make([]Map, 0, len(c.Maps))
	for _, m := range c.Maps {
		*list = append(*list, m)
	}
	return nil
}

type DefaultMap struct {
	Mode               int
	Name               string
	GameMode           int
	Seed               int64
	Structures, Cheats bool
}

func (c *Config) CreateDefaultMap(data DefaultMap, _ *struct{}) error {
	d, err := setupMapDir()
	if err != nil {
		return err
	}
	p, err := mincraft.NewFilePath(d)
	if err != nil {
		return err
	}
	l, err := minecraft.NewLevel(p)
	if err != nil {
		return err
	}
	l.GameMode(int32(data.GameMode))
	l.LevelName(data.Name)
	switch data.Mode {
	case 0:
		l.Generator(minecraft.DefaultGenerator)
	case 1:
		l.Generator(minecraft.FlatGenerator)
	case 2:
		l.Generator(minecraft.LargeBiomeGenerator)
	case 3:
		l.Generator(minecraft.AmplifiedGenerator)
	case 4:
		l.Generator(minecraft.CustomGenerator)
	}
	l.Seed(data.Seed)
	l.AllowCommands(data.Cheats)
	l.MapFeatures(data.Structures)
	l.Save()
	return nil
}

var mapDirLock sync.Mutex

func setupMapDir() (string, error) {
	mapDirLock.Lock()
	defer mapDirLock.Unlock()
	num := 0
	for {
		dir := path.Join(config.MapsDir, strconv.Itoa(num))
		_, err := os.Stat(dir)
		if os.IsNotExist(err) {
			err := os.MkdirAll(dir, 0777)
			if err != nil {
				return "", err
			}
			return dir, nil
		}
		num++
	}
}
<commit_msg>Added missing newMap call<commit_after>package main

import (
	"os"
	"path"
	"strconv"
	"sync"

	"github.com/MJKWoolnough/minecraft"
)

func (c *Config) Name(_ struct{}, serverName *string) error {
	c.mu.RLock()
	defer c.mu.RUnlock()
	*serverName = c.ServerName
	return nil
}

func (c *Config) List(_ struct{}, list *[]Server) error {
	c.mu.RLock()
	defer c.mu.RUnlock()
	*list = make([]Server, 0, len(c.Servers))
	for _, s := range c.Servers {
		*list = append(*list, s)
	}
	return nil
}

func (c *Config) Save(s Server, _ *struct{}) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	defer c.save()
	ns, ok := c.Servers[s.ID]
	if !ok {
		return ErrNoServer
	}
	s.Path = ns.Path
	s.status = ns.status
	c.Servers[s.ID] = s
	return nil
}

func (c *Config) MapList(_ struct{}, list *[]Map) error {
	c.mu.RLock()
	defer c.mu.RUnlock()
	*list = make([]Map, 0, len(c.Maps))
	for _, m := range c.Maps {
		*list = append(*list, m)
	}
	return nil
}

type DefaultMap struct {
	Mode               int
	Name               string
	GameMode           int
	Seed               int64
	Structures, Cheats bool
}

func (c *Config) CreateDefaultMap(data DefaultMap, _ *struct{}) error {
	d, err := setupMapDir()
	if err != nil {
		return err
	}
	p, err := mincraft.NewFilePath(d)
	if err != nil {
		return err
	}
	l, err := minecraft.NewLevel(p)
	if err != nil {
		return err
	}
	l.GameMode(int32(data.GameMode))
	l.LevelName(data.Name)
	switch data.Mode {
	case 0:
		l.Generator(minecraft.DefaultGenerator)
	case 1:
		l.Generator(minecraft.FlatGenerator)
	case 2:
		l.Generator(minecraft.LargeBiomeGenerator)
	case 3:
		l.Generator(minecraft.AmplifiedGenerator)
	case 4:
		l.Generator(minecraft.CustomGenerator)
	}
	l.Seed(data.Seed)
	l.AllowCommands(data.Cheats)
	l.MapFeatures(data.Structures)
	l.Save()
	c.newMap(data.Name, d)
	return nil
}

var mapDirLock sync.Mutex

func setupMapDir() (string, error) {
	mapDirLock.Lock()
	defer mapDirLock.Unlock()
	num := 0
	for {
		dir := path.Join(config.MapsDir, strconv.Itoa(num))
		_, err := os.Stat(dir)
		if os.IsNotExist(err) {
			err := os.MkdirAll(dir, 0777)
			if err != nil {
				return "", err
			}
			return dir, nil
		}
		num++
	}
}
<|endoftext|>