<commit_before>package slt

import (
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"syscall"
	"time"

	"github.com/capnm/sysinfo" // TODO: can be thrown using syscall!!

)

func KiraWrap(icsFileName, intTime, randomNumber string, noGPU bool) () {
	var (
		timeLimit string
		err error
		kiraString string
		kiraArgs []string
		kiraWrappedCmd *exec.Cmd
		pathName, icsName, outName, errName, ext string
		icsFile, outFile, errFile *os.File
		regRes map[string]string
		done = make(chan string, 1)
		randomSeed string = ""
		u  *user.User
		host, wd string
	)

	pathName = filepath.Dir(icsFileName)
	icsName = filepath.Base(icsFileName)
	timeLimit = intTime
	randomSeed = randomNumber

	fmt.Println("###################################################")

	// Extract fileNameBody, round and ext
	log.Println("Extract files names")
	ext = filepath.Ext(icsName)
	regRes, err = Reg(icsName)
	if err != nil {
		log.Println("Can't derive standard names from STDOUT => wrap it!!")
		errName = "err-" + icsName + ext
		outName = "out-" + icsName + ext
	} else {
		if regRes["prefix"] != "ics" {
			log.Fatalf("Please specify a STDIN file, found %v prefix", regRes["prefix"])
		}

		// Creating new filenames
		errName = "err-" + regRes["baseName"] + "-run" + regRes["run"] + "-rnd" + regRes["rnd"] + ext
		outName = "out-" + regRes["baseName"] + "-run" + regRes["run"] + "-rnd" + regRes["rnd"] + ext
	}

	if icsName, err = filepath.Abs(filepath.Join(pathName, icsName)); err != nil {
		log.Fatal("Error composing icsName: ", err)
	}
	if errName, err = filepath.Abs(filepath.Join(pathName, errName)); err != nil {
		log.Fatal("Error composing errName: ", err)
	}
	if outName, err = filepath.Abs(filepath.Join(pathName, outName)); err != nil {
		log.Fatal("Error composing outName: ", err)
	}

	if icsFile, err = os.Open(icsName); err != nil {log.Fatal(err)}
	defer icsFile.Close()
	if outFile, err = os.Create(outName); err != nil {log.Fatal(err)}
	defer outFile.Close()
	if errFile, err = os.Create(errName); err != nil {log.Fatal(err)}
	defer errFile.Close()

	errFile.WriteString("\n#==============================\n")
	errFile.WriteString(fmt.Sprintf("\n#   %v Start with kiraWrap.\n", time.Now().Format(time.RFC850)))
	errFile.WriteString("\n#==============================\n")

	outFile.WriteString("\n#==============================\n")
	outFile.WriteString(fmt.Sprintf("\n#   %v Start with kiraWrap.\n", time.Now().Format(time.RFC850)))
	outFile.WriteString("\n#==============================\n")

	kiraArgs =  []string{"-t", timeLimit,// +  // number of timesteps to compute
		"-d", "1",// +  // log output interval
		"-D", "1",// +  // snapshot interval
		// 				"-b", "1",// +  // frequency of full binary output
		"-f", "0",// +  // dynamical friction (0 = no friction, 1 = friction)
		"-n", "10",// +  // terminate if the cluster remains with only 10 particles
		"-e", "0.000",// + // softening
		// 				"-B",// // switch on binary evolution
		//"-s 36543" // random seed
	}

	if noGPU {
		log.Println("Selected the no GPU integration.")
		log.Println("Assuming kira is in $HOME/bin/kira-no-GPU, if not, please copy it there... for sake of simplicity!:P")
		kiraString = filepath.Join(os.Getenv("HOME"), "/bin/", "kira-no-GPU")
	} else if tf {
		log.Println("Selected TF versionn.")
		log.Println("Assuming kira is in $HOME/bin/kiraTF, if not, please copy it there... for sake of simplicity!:P")
		kiraString = filepath.Join(os.Getenv("HOME"), "/bin/", "kiraTF")
// 		kiraArgs = append(kiraArgs, "-F") // F now means dynamical friction
	} else if as {
		log.Println("Selected TF versionn.")
		log.Println("Assuming kira is in $HOME/bin/kira-AS, if not, please copy it there... for sake of simplicity!:P")
		kiraString = filepath.Join(os.Getenv("HOME"), "/bin/", "kira-AS")
	}
	} else {
		log.Println("Assuming kira is in $HOME/bin/kira, if not, please copy it there... for sake of simplicity!:P")
		kiraString = filepath.Join(os.Getenv("HOME"), "/bin/", "kira")
	}

	if !noBinaries {
		log.Println("Binary evolution on, adding flags to kira")
		kiraArgs = append(kiraArgs, "-b", "1", "-B")
	} else {
		kiraArgs = append(kiraArgs, "-S")
		log.Println("Binary evolution off")
	}

	// Add the random seed if specified
	if randomSeed != "" {
		kiraArgs = append(kiraArgs, "-s", randomSeed)
	}

	kiraWrappedCmd = exec.Command(kiraString, kiraArgs...)
	if kiraWrappedCmd.Stdin = icsFile; err != nil {log.Fatal("Error connecting ICs to kira STDIN: ", err)}
	if kiraWrappedCmd.Stdout = outFile; err != nil {log.Fatal("Error connecting ICs to kira STDOUT: ", err)}
	if kiraWrappedCmd.Stderr = errFile; err != nil {log.Fatal("Error connecting ICs to kira STDERR: ", err)}

	if u, err = user.Current(); err != nil {
		log.Fatal("Can't retrieve username: ", err)
	}
	if host, err = os.Hostname(); err != nil {
		log.Fatal("Can't retrieve hostname: ", err)
	}
	if wd, err = os.Getwd(); err != nil {
		log.Fatal("Can't retrieve working dir: ", err)
	}

	log.Println("Run summary:")

	log.Printf("Username: %v (%v)\n", u.Username, u.Name)
	log.Println("Hostname: ", host)
	log.Println("Working dir: ", wd)
	log.Println("LD_LIBRARY_PATH: ", os.Getenv("LD_LIBRARY_PATH"))

	log.Println("Command: ", kiraString, kiraArgs)
	log.Println("STDIN = ", icsName)
	log.Println("STDOUT = ", outName)
	log.Println("STDERR = ", errName)

	log.Println("Ready... steady... Go!")

	if err = kiraWrappedCmd.Start(); err != nil {
		log.Fatal("Error starting kiraWrappedCmd: ", err)
	}

	log.Println("Waiting kira to finish while checking for problems...")
	// Wait for the process to end normally
	go waitProcess(kiraWrappedCmd, done)
	// Check for pp3-stalling situations
	go killTrigger(errName, kiraWrappedCmd, done)

	if err = errors.New(<-done); err.Error() != "" {
		errFile.WriteString("\n"+err.Error()+"\n")
	}

	errFile.WriteString("\n#==============================\n")
	errFile.WriteString(fmt.Sprintf("\n#   %v Done with kiraWrap.\n", time.Now().Format(time.RFC850)))
	errFile.WriteString(fmt.Sprintf("\n#   Username: %v (%V)\n", u.Username, u.Name))
	errFile.WriteString(fmt.Sprintf("\n#   Hostname: %v\n", host))
	errFile.WriteString("\n==============================\n")

	outFile.WriteString("\n#==============================\n")
	outFile.WriteString(fmt.Sprintf("\n#   %v Done with kiraWrap.\n", time.Now().Format(time.RFC850)))
	outFile.WriteString("\n#==============================\n")

	fmt.Print("\x07") // Beep when finish!!:D
}

func killTrigger(errName string, kiraWrappedCmd *exec.Cmd, done chan string) () {
	const toGB = float64(1. / (1024*1024*1024))
	const maxStderrGB = float64(2)
	const minDiskGB = float64(5)
	const maxMemPerCent = float64(98)

	var (
		fileInfo os.FileInfo
		sysInfo *sysinfo.SI
		memAvail float64
		diskAvailGB float64
		wd string
		reason string
		err error
	)

	if wd, err = os.Getwd(); err != nil {
		log.Fatal("Can't retrieve local dir: ", err)
	}

	fsInfo := syscall.Statfs_t{}

	for {
		// Check STDERR file size
		if fileInfo, err = os.Stat(errName); err != nil {
			log.Fatal("Error checking STDERR file size, err")
		}
		// STDERR exceesing aloowed dimension of 2GB
		// probably the simulation is stalling because
		// of pp3 locked on a binary
		if float64(fileInfo.Size()) * toGB > maxStderrGB {
			reason = "probable pp3 stalling"
			break
		}

		// Check memory availability
		sysInfo = sysinfo.Get()
		memAvail = float64(sysInfo.FreeRam) / float64(sysInfo.TotalRam) * 100
		if memAvail > maxMemPerCent {
			reason = fmt.Sprintf("memory less than 95%% on the system: %2.2f",  memAvail)
			break
		}

		// Check HDD space availability
		if err = syscall.Statfs(wd, &fsInfo); err!= nil{
			log.Fatal("Cant't retrieve file system information: ", err)
		}
		if diskAvailGB = float64(fsInfo.Bavail) * float64(fsInfo.Frsize) * toGB; diskAvailGB < minDiskGB {
			reason = " available disk space less than 5 GB on the system"
			break
		}

		// Wait some time
		time.Sleep(time.Duration(2) * time.Minute)
	}
	log.Println("Kill kira because ", reason)
	if err := kiraWrappedCmd.Process.Kill(); err != nil {
		log.Fatal("failed to kill: ", err)
	}
	done <- "killed because " + reason
}

func waitProcess(kiraWrappedCmd *exec.Cmd, done chan string) () {
	err := kiraWrappedCmd.Wait()
	log.Println("Process exited with error ", err)
	if err == nil {
		log.Println("Maybe the simulation is COMPLETE or the STD* are mixed and wrong.")
	}
	done <- ""
}

<commit_msg>fix kiraWrap<commit_after>package slt

import (
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"syscall"
	"time"

	"github.com/capnm/sysinfo" // TODO: can be thrown using syscall!!

)

func KiraWrap(icsFileName, intTime, randomNumber string, noGPU bool) () {
	var (
		timeLimit string
		err error
		kiraString string
		kiraArgs []string
		kiraWrappedCmd *exec.Cmd
		pathName, icsName, outName, errName, ext string
		icsFile, outFile, errFile *os.File
		regRes map[string]string
		done = make(chan string, 1)
		randomSeed string = ""
		u  *user.User
		host, wd string
	)

	pathName = filepath.Dir(icsFileName)
	icsName = filepath.Base(icsFileName)
	timeLimit = intTime
	randomSeed = randomNumber

	fmt.Println("###################################################")

	// Extract fileNameBody, round and ext
	log.Println("Extract files names")
	ext = filepath.Ext(icsName)
	regRes, err = Reg(icsName)
	if err != nil {
		log.Println("Can't derive standard names from STDOUT => wrap it!!")
		errName = "err-" + icsName + ext
		outName = "out-" + icsName + ext
	} else {
		if regRes["prefix"] != "ics" {
			log.Fatalf("Please specify a STDIN file, found %v prefix", regRes["prefix"])
		}

		// Creating new filenames
		errName = "err-" + regRes["baseName"] + "-run" + regRes["run"] + "-rnd" + regRes["rnd"] + ext
		outName = "out-" + regRes["baseName"] + "-run" + regRes["run"] + "-rnd" + regRes["rnd"] + ext
	}

	if icsName, err = filepath.Abs(filepath.Join(pathName, icsName)); err != nil {
		log.Fatal("Error composing icsName: ", err)
	}
	if errName, err = filepath.Abs(filepath.Join(pathName, errName)); err != nil {
		log.Fatal("Error composing errName: ", err)
	}
	if outName, err = filepath.Abs(filepath.Join(pathName, outName)); err != nil {
		log.Fatal("Error composing outName: ", err)
	}

	if icsFile, err = os.Open(icsName); err != nil {log.Fatal(err)}
	defer icsFile.Close()
	if outFile, err = os.Create(outName); err != nil {log.Fatal(err)}
	defer outFile.Close()
	if errFile, err = os.Create(errName); err != nil {log.Fatal(err)}
	defer errFile.Close()

	errFile.WriteString("\n#==============================\n")
	errFile.WriteString(fmt.Sprintf("\n#   %v Start with kiraWrap.\n", time.Now().Format(time.RFC850)))
	errFile.WriteString("\n#==============================\n")

	outFile.WriteString("\n#==============================\n")
	outFile.WriteString(fmt.Sprintf("\n#   %v Start with kiraWrap.\n", time.Now().Format(time.RFC850)))
	outFile.WriteString("\n#==============================\n")

	kiraArgs =  []string{"-t", timeLimit,// +  // number of timesteps to compute
		"-d", "1",// +  // log output interval
		"-D", "1",// +  // snapshot interval
		// 				"-b", "1",// +  // frequency of full binary output
		"-f", "0",// +  // dynamical friction (0 = no friction, 1 = friction)
		"-n", "10",// +  // terminate if the cluster remains with only 10 particles
		"-e", "0.000",// + // softening
		// 				"-B",// // switch on binary evolution
		//"-s 36543" // random seed
	}

	if noGPU {
		log.Println("Selected the no GPU integration.")
		log.Println("Assuming kira is in $HOME/bin/kira-no-GPU, if not, please copy it there... for sake of simplicity!:P")
		kiraString = filepath.Join(os.Getenv("HOME"), "/bin/", "kira-no-GPU")
	} else if tf {
		log.Println("Selected TF versionn.")
		log.Println("Assuming kira is in $HOME/bin/kiraTF, if not, please copy it there... for sake of simplicity!:P")
		kiraString = filepath.Join(os.Getenv("HOME"), "/bin/", "kiraTF")
// 		kiraArgs = append(kiraArgs, "-F") // F now means dynamical friction
	} else if as {
		log.Println("Selected TF versionn.")
		log.Println("Assuming kira is in $HOME/bin/kira-AS, if not, please copy it there... for sake of simplicity!:P")
		kiraString = filepath.Join(os.Getenv("HOME"), "/bin/", "kira-AS")
	} else {
		log.Println("Assuming kira is in $HOME/bin/kira, if not, please copy it there... for sake of simplicity!:P")
		kiraString = filepath.Join(os.Getenv("HOME"), "/bin/", "kira")
	}

	if !noBinaries {
		log.Println("Binary evolution on, adding flags to kira")
		kiraArgs = append(kiraArgs, "-b", "1", "-B")
	} else {
		kiraArgs = append(kiraArgs, "-S")
		log.Println("Binary evolution off")
	}

	// Add the random seed if specified
	if randomSeed != "" {
		kiraArgs = append(kiraArgs, "-s", randomSeed)
	}

	kiraWrappedCmd = exec.Command(kiraString, kiraArgs...)
	if kiraWrappedCmd.Stdin = icsFile; err != nil {log.Fatal("Error connecting ICs to kira STDIN: ", err)}
	if kiraWrappedCmd.Stdout = outFile; err != nil {log.Fatal("Error connecting ICs to kira STDOUT: ", err)}
	if kiraWrappedCmd.Stderr = errFile; err != nil {log.Fatal("Error connecting ICs to kira STDERR: ", err)}

	if u, err = user.Current(); err != nil {
		log.Fatal("Can't retrieve username: ", err)
	}
	if host, err = os.Hostname(); err != nil {
		log.Fatal("Can't retrieve hostname: ", err)
	}
	if wd, err = os.Getwd(); err != nil {
		log.Fatal("Can't retrieve working dir: ", err)
	}

	log.Println("Run summary:")

	log.Printf("Username: %v (%v)\n", u.Username, u.Name)
	log.Println("Hostname: ", host)
	log.Println("Working dir: ", wd)
	log.Println("LD_LIBRARY_PATH: ", os.Getenv("LD_LIBRARY_PATH"))

	log.Println("Command: ", kiraString, kiraArgs)
	log.Println("STDIN = ", icsName)
	log.Println("STDOUT = ", outName)
	log.Println("STDERR = ", errName)

	log.Println("Ready... steady... Go!")

	if err = kiraWrappedCmd.Start(); err != nil {
		log.Fatal("Error starting kiraWrappedCmd: ", err)
	}

	log.Println("Waiting kira to finish while checking for problems...")
	// Wait for the process to end normally
	go waitProcess(kiraWrappedCmd, done)
	// Check for pp3-stalling situations
	go killTrigger(errName, kiraWrappedCmd, done)

	if err = errors.New(<-done); err.Error() != "" {
		errFile.WriteString("\n"+err.Error()+"\n")
	}

	errFile.WriteString("\n#==============================\n")
	errFile.WriteString(fmt.Sprintf("\n#   %v Done with kiraWrap.\n", time.Now().Format(time.RFC850)))
	errFile.WriteString(fmt.Sprintf("\n#   Username: %v (%V)\n", u.Username, u.Name))
	errFile.WriteString(fmt.Sprintf("\n#   Hostname: %v\n", host))
	errFile.WriteString("\n==============================\n")

	outFile.WriteString("\n#==============================\n")
	outFile.WriteString(fmt.Sprintf("\n#   %v Done with kiraWrap.\n", time.Now().Format(time.RFC850)))
	outFile.WriteString("\n#==============================\n")

	fmt.Print("\x07") // Beep when finish!!:D
}

func killTrigger(errName string, kiraWrappedCmd *exec.Cmd, done chan string) () {
	const toGB = float64(1. / (1024*1024*1024))
	const maxStderrGB = float64(2)
	const minDiskGB = float64(5)
	const maxMemPerCent = float64(98)

	var (
		fileInfo os.FileInfo
		sysInfo *sysinfo.SI
		memAvail float64
		diskAvailGB float64
		wd string
		reason string
		err error
	)

	if wd, err = os.Getwd(); err != nil {
		log.Fatal("Can't retrieve local dir: ", err)
	}

	fsInfo := syscall.Statfs_t{}

	for {
		// Check STDERR file size
		if fileInfo, err = os.Stat(errName); err != nil {
			log.Fatal("Error checking STDERR file size, err")
		}
		// STDERR exceesing aloowed dimension of 2GB
		// probably the simulation is stalling because
		// of pp3 locked on a binary
		if float64(fileInfo.Size()) * toGB > maxStderrGB {
			reason = "probable pp3 stalling"
			break
		}

		// Check memory availability
		sysInfo = sysinfo.Get()
		memAvail = float64(sysInfo.FreeRam) / float64(sysInfo.TotalRam) * 100
		if memAvail > maxMemPerCent {
			reason = fmt.Sprintf("memory less than 95%% on the system: %2.2f",  memAvail)
			break
		}

		// Check HDD space availability
		if err = syscall.Statfs(wd, &fsInfo); err!= nil{
			log.Fatal("Cant't retrieve file system information: ", err)
		}
		if diskAvailGB = float64(fsInfo.Bavail) * float64(fsInfo.Frsize) * toGB; diskAvailGB < minDiskGB {
			reason = " available disk space less than 5 GB on the system"
			break
		}

		// Wait some time
		time.Sleep(time.Duration(2) * time.Minute)
	}
	log.Println("Kill kira because ", reason)
	if err := kiraWrappedCmd.Process.Kill(); err != nil {
		log.Fatal("failed to kill: ", err)
	}
	done <- "killed because " + reason
}

func waitProcess(kiraWrappedCmd *exec.Cmd, done chan string) () {
	err := kiraWrappedCmd.Wait()
	log.Println("Process exited with error ", err)
	if err == nil {
		log.Println("Maybe the simulation is COMPLETE or the STD* are mixed and wrong.")
	}
	done <- ""
}

<|endoftext|>


<commit_before>/*
Copyright Suzhou Tongji Fintech Research Institute 2017 All Rights Reserved.
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 sm2

import (
	"crypto/rand"
	"crypto/x509/pkix"
	"encoding/asn1"
	"fmt"
	"io/ioutil"
	"log"
	"math/big"
	"net"
	"os"
	"testing"
	"time"
)

func TestSm2(t *testing.T) {
	priv, err := GenerateKey() // 生成密钥对
	if err != nil {
		log.Fatal(err)
	}
	ok, err := WritePrivateKeytoPem("priv.pem", priv, nil) // 生成密钥文件
	if ok != true {
		log.Fatal(err)
	}
	pubKey, _ := priv.Public().(*PublicKey)
	ok, err = WritePublicKeytoPem("pub.pem", pubKey, nil) // 生成公钥文件
	if ok != true {
		log.Fatal(err)
	}
	msg := []byte("test")
	err = ioutil.WriteFile("ifile", msg, os.FileMode(0644)) // 生成测试文件
	if err != nil {
		log.Fatal(err)
	}
	privKey, err := ReadPrivateKeyFromPem("priv.pem", nil) // 读取密钥
	if err != nil {
		log.Fatal(err)
	}
	pubKey, err = ReadPublicKeyFromPem("pub.pem", nil) // 读取公钥
	if err != nil {
		log.Fatal(err)
	}
	msg, _ = ioutil.ReadFile("ifile")                // 从文件读取数据
	sign, err := privKey.Sign(rand.Reader, msg, nil) // 签名
	if err != nil {
		log.Fatal(err)
	}
	err = ioutil.WriteFile("ofile", sign, os.FileMode(0644))
	if err != nil {
		log.Fatal(err)
	}
	signdata, _ := ioutil.ReadFile("ofile")
	ok = privKey.Verify(msg, signdata) // 密钥验证
	if ok != true {
		fmt.Printf("Verify error\n")
	} else {
		fmt.Printf("Verify ok\n")
	}
	ok = pubKey.Verify(msg, signdata) // 公钥验证
	if ok != true {
		fmt.Printf("Verify error\n")
	} else {
		fmt.Printf("Verify ok\n")
	}
	templateReq := CertificateRequest{
		Subject: pkix.Name{
			CommonName:   "test.example.com",
			Organization: []string{"Test"},
		},
		//		SignatureAlgorithm: ECDSAWithSHA256,
		SignatureAlgorithm: SM2WithSM3,
	}
	_, err = CreateCertificateRequestToPem("req.pem", &templateReq, privKey)
	if err != nil {
		log.Fatal(err)
	}
	req, err := ReadCertificateRequestFromPem("req.pem")
	if err != nil {
		log.Fatal(err)
	}
	err = req.CheckSignature()
	if err != nil {
		log.Fatal(err)
	} else {
		fmt.Printf("CheckSignature ok\n")
	}
	testExtKeyUsage := []ExtKeyUsage{ExtKeyUsageClientAuth, ExtKeyUsageServerAuth}
	testUnknownExtKeyUsage := []asn1.ObjectIdentifier{[]int{1, 2, 3}, []int{2, 59, 1}}
	extraExtensionData := []byte("extra extension")
	commonName := "test.example.com"
	template := Certificate{
		// SerialNumber is negative to ensure that negative
		// values are parsed. This is due to the prevalence of
		// buggy code that produces certificates with negative
		// serial numbers.
		SerialNumber: big.NewInt(-1),
		Subject: pkix.Name{
			CommonName:   commonName,
			Organization: []string{"TEST"},
			Country:      []string{"China"},
			ExtraNames: []pkix.AttributeTypeAndValue{
				{
					Type:  []int{2, 5, 4, 42},
					Value: "Gopher",
				},
				// This should override the Country, above.
				{
					Type:  []int{2, 5, 4, 6},
					Value: "NL",
				},
			},
		},
		NotBefore: time.Unix(1000, 0),
		NotAfter:  time.Unix(100000, 0),

		//		SignatureAlgorithm: ECDSAWithSHA256,
		SignatureAlgorithm: SM2WithSM3,

		SubjectKeyId: []byte{1, 2, 3, 4},
		KeyUsage:     KeyUsageCertSign,

		ExtKeyUsage:        testExtKeyUsage,
		UnknownExtKeyUsage: testUnknownExtKeyUsage,

		BasicConstraintsValid: true,
		IsCA: true,

		OCSPServer:            []string{"http://ocsp.example.com"},
		IssuingCertificateURL: []string{"http://crt.example.com/ca1.crt"},

		DNSNames:       []string{"test.example.com"},
		EmailAddresses: []string{"gopher@golang.org"},
		IPAddresses:    []net.IP{net.IPv4(127, 0, 0, 1).To4(), net.ParseIP("2001:4860:0:2001::68")},

		PolicyIdentifiers:   []asn1.ObjectIdentifier{[]int{1, 2, 3}},
		PermittedDNSDomains: []string{".example.com", "example.com"},

		CRLDistributionPoints: []string{"http://crl1.example.com/ca1.crl", "http://crl2.example.com/ca1.crl"},

		ExtraExtensions: []pkix.Extension{
			{
				Id:    []int{1, 2, 3, 4},
				Value: extraExtensionData,
			},
			// This extension should override the SubjectKeyId, above.
			{
				Id:       oidExtensionSubjectKeyId,
				Critical: false,
				Value:    []byte{0x04, 0x04, 4, 3, 2, 1},
			},
		},
	}
	pubKey, _ = priv.Public().(*PublicKey)
	ok, _ = CreateCertificateToPem("cert.pem", &template, &template, pubKey, privKey)
	if ok != true {
		fmt.Printf("failed to create cert file\n")
	}
	cert, err := ReadCertificateFromPem("cert.pem")
	if err != nil {
		fmt.Printf("failed to read cert file")
	}
	err = cert.CheckSignature(cert.SignatureAlgorithm, cert.RawTBSCertificate, cert.Signature)
	if err != nil {
		log.Fatal(err)
	} else {
		fmt.Printf("CheckSignature ok\n")
	}
}

func BenchmarkSM2(t *testing.B) {
	t.ReportAllocs()
	for i := 0; i < t.N; i++ {
		_, err := GenerateKey() // 生成密钥对
		if err != nil {
			log.Fatal(err)
		}
	}
}
<commit_msg>sm2性能测试代码完善<commit_after>/*
Copyright Suzhou Tongji Fintech Research Institute 2017 All Rights Reserved.
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 sm2

import (
	"crypto/rand"
	"crypto/x509/pkix"
	"encoding/asn1"
	"fmt"
	"io/ioutil"
	"log"
	"math/big"
	"net"
	"os"
	"testing"
	"time"
)

func TestSm2(t *testing.T) {
	priv, err := GenerateKey() // 生成密钥对
	if err != nil {
		log.Fatal(err)
	}
	ok, err := WritePrivateKeytoPem("priv.pem", priv, nil) // 生成密钥文件
	if ok != true {
		log.Fatal(err)
	}
	pubKey, _ := priv.Public().(*PublicKey)
	ok, err = WritePublicKeytoPem("pub.pem", pubKey, nil) // 生成公钥文件
	if ok != true {
		log.Fatal(err)
	}
	msg := []byte("test")
	err = ioutil.WriteFile("ifile", msg, os.FileMode(0644)) // 生成测试文件
	if err != nil {
		log.Fatal(err)
	}
	privKey, err := ReadPrivateKeyFromPem("priv.pem", nil) // 读取密钥
	if err != nil {
		log.Fatal(err)
	}
	pubKey, err = ReadPublicKeyFromPem("pub.pem", nil) // 读取公钥
	if err != nil {
		log.Fatal(err)
	}
	msg, _ = ioutil.ReadFile("ifile")                // 从文件读取数据
	sign, err := privKey.Sign(rand.Reader, msg, nil) // 签名
	if err != nil {
		log.Fatal(err)
	}
	err = ioutil.WriteFile("ofile", sign, os.FileMode(0644))
	if err != nil {
		log.Fatal(err)
	}
	signdata, _ := ioutil.ReadFile("ofile")
	ok = privKey.Verify(msg, signdata) // 密钥验证
	if ok != true {
		fmt.Printf("Verify error\n")
	} else {
		fmt.Printf("Verify ok\n")
	}
	ok = pubKey.Verify(msg, signdata) // 公钥验证
	if ok != true {
		fmt.Printf("Verify error\n")
	} else {
		fmt.Printf("Verify ok\n")
	}
	templateReq := CertificateRequest{
		Subject: pkix.Name{
			CommonName:   "test.example.com",
			Organization: []string{"Test"},
		},
		//		SignatureAlgorithm: ECDSAWithSHA256,
		SignatureAlgorithm: SM2WithSM3,
	}
	_, err = CreateCertificateRequestToPem("req.pem", &templateReq, privKey)
	if err != nil {
		log.Fatal(err)
	}
	req, err := ReadCertificateRequestFromPem("req.pem")
	if err != nil {
		log.Fatal(err)
	}
	err = req.CheckSignature()
	if err != nil {
		log.Fatal(err)
	} else {
		fmt.Printf("CheckSignature ok\n")
	}
	testExtKeyUsage := []ExtKeyUsage{ExtKeyUsageClientAuth, ExtKeyUsageServerAuth}
	testUnknownExtKeyUsage := []asn1.ObjectIdentifier{[]int{1, 2, 3}, []int{2, 59, 1}}
	extraExtensionData := []byte("extra extension")
	commonName := "test.example.com"
	template := Certificate{
		// SerialNumber is negative to ensure that negative
		// values are parsed. This is due to the prevalence of
		// buggy code that produces certificates with negative
		// serial numbers.
		SerialNumber: big.NewInt(-1),
		Subject: pkix.Name{
			CommonName:   commonName,
			Organization: []string{"TEST"},
			Country:      []string{"China"},
			ExtraNames: []pkix.AttributeTypeAndValue{
				{
					Type:  []int{2, 5, 4, 42},
					Value: "Gopher",
				},
				// This should override the Country, above.
				{
					Type:  []int{2, 5, 4, 6},
					Value: "NL",
				},
			},
		},
		NotBefore: time.Unix(1000, 0),
		NotAfter:  time.Unix(100000, 0),

		//		SignatureAlgorithm: ECDSAWithSHA256,
		SignatureAlgorithm: SM2WithSM3,

		SubjectKeyId: []byte{1, 2, 3, 4},
		KeyUsage:     KeyUsageCertSign,

		ExtKeyUsage:        testExtKeyUsage,
		UnknownExtKeyUsage: testUnknownExtKeyUsage,

		BasicConstraintsValid: true,
		IsCA: true,

		OCSPServer:            []string{"http://ocsp.example.com"},
		IssuingCertificateURL: []string{"http://crt.example.com/ca1.crt"},

		DNSNames:       []string{"test.example.com"},
		EmailAddresses: []string{"gopher@golang.org"},
		IPAddresses:    []net.IP{net.IPv4(127, 0, 0, 1).To4(), net.ParseIP("2001:4860:0:2001::68")},

		PolicyIdentifiers:   []asn1.ObjectIdentifier{[]int{1, 2, 3}},
		PermittedDNSDomains: []string{".example.com", "example.com"},

		CRLDistributionPoints: []string{"http://crl1.example.com/ca1.crl", "http://crl2.example.com/ca1.crl"},

		ExtraExtensions: []pkix.Extension{
			{
				Id:    []int{1, 2, 3, 4},
				Value: extraExtensionData,
			},
			// This extension should override the SubjectKeyId, above.
			{
				Id:       oidExtensionSubjectKeyId,
				Critical: false,
				Value:    []byte{0x04, 0x04, 4, 3, 2, 1},
			},
		},
	}
	pubKey, _ = priv.Public().(*PublicKey)
	ok, _ = CreateCertificateToPem("cert.pem", &template, &template, pubKey, privKey)
	if ok != true {
		fmt.Printf("failed to create cert file\n")
	}
	cert, err := ReadCertificateFromPem("cert.pem")
	if err != nil {
		fmt.Printf("failed to read cert file")
	}
	err = cert.CheckSignature(cert.SignatureAlgorithm, cert.RawTBSCertificate, cert.Signature)
	if err != nil {
		log.Fatal(err)
	} else {
		fmt.Printf("CheckSignature ok\n")
	}
}

func BenchmarkSM2(t *testing.B) {
	t.ReportAllocs()
	for i := 0; i < t.N; i++ {
		priv, err := GenerateKey() // 生成密钥对
		if err != nil {
			log.Fatal(err)
		}
		msg := []byte("test")
		sign, err := priv.Sign(rand.Reader, msg, nil) // 签名
		if err != nil {
			log.Fatal(err)
		}
		ok := priv.Verify(msg, sign) // 密钥验证
		if ok != true {
			fmt.Printf("Verify error\n")
		} else {
			fmt.Printf("Verify ok\n")
		}
	}
}
<|endoftext|>


<commit_before>package manage

import (
	"fmt"
	"os"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"text/template"
	"time"

	"github.com/antlinker/alog/buffer"
	"github.com/antlinker/alog/log"
	"github.com/antlinker/alog/store"
	"github.com/antlinker/alog/utils"
)

// NewLogManage 创建新的LogManage实例
func NewLogManage(config *log.LogConfig) log.LogManage {
	manage := &_LogManage{
		Config: config,
	}
	manage.Template = map[log.TmplKey]*template.Template{
		log.TmplConsole:     template.Must(template.New("").Parse(config.Console.Item.Tmpl)),
		log.TmplConsoleTime: template.Must(template.New("").Parse(config.Console.Item.TimeTmpl)),
	}
	switch config.Global.Buffer.Engine {
	case log.REDIS_BUFFER:
		var redisConfig log.RedisConfig
		redisStore := config.Store.Redis
		if redisStore != nil {
			if v, ok := redisStore[config.Global.Buffer.TargetStore]; ok {
				redisConfig = v
			}
		}
		manage.Buffer = buffer.NewRedisBuffer(redisConfig)
	default:
		manage.Buffer = buffer.NewMemoryBuffer()
	}
	manageStore := make(map[string]log.LogStore)
	if fileStore := config.Store.File; fileStore != nil {
		for k, v := range fileStore {
			manageStore[k] = store.NewFileStore(v)
		}
	}
	if elasticStore := config.Store.Elastic; elasticStore != nil {
		for k, v := range elasticStore {
			manageStore[k] = store.NewElasticStore(v)
		}
	}
	if mongoStore := config.Store.Mongo; mongoStore != nil {
		for k, v := range mongoStore {
			manageStore[k] = store.NewMongoStore(v)
		}
	}
	manage.Store = manageStore

	if config.Global.IsEnabled == 1 {
		go manage.execStore()
	}

	return manage
}

// _LogManage
type _LogManage struct {
	sync.RWMutex
	gID        uint64
	storeTotal int64
	Config     *log.LogConfig
	Template   map[log.TmplKey]*template.Template
	Buffer     log.LogBuffer
	Store      map[string]log.LogStore
}

func (lm *_LogManage) Write(level log.LogLevel, tag log.LogTag, v ...interface{}) {
	msg := fmt.Sprint(v...)
	lm.writeMsg(level, tag, msg)
}

func (lm *_LogManage) Writef(level log.LogLevel, tag log.LogTag, format string, v ...interface{}) {
	msg := fmt.Sprintf(format, v...)
	lm.writeMsg(level, tag, msg)
}

func (lm *_LogManage) Console(level log.LogLevel, tag log.LogTag, v ...interface{}) {
	msg := fmt.Sprint(v...)
	lm.writeConsole(level, tag, msg)
}

func (lm *_LogManage) Consolef(level log.LogLevel, tag log.LogTag, format string, v ...interface{}) {
	msg := fmt.Sprintf(format, v...)
	lm.writeConsole(level, tag, msg)
}

func (lm *_LogManage) writeConsole(level log.LogLevel, tag log.LogTag, msg string) {
	item := lm.logItem(level, tag, msg)
	lm.console(&item)
}

func (lm *_LogManage) TotalNum() int64 {
	return lm.storeTotal
}

func (lm *_LogManage) writeMsg(level log.LogLevel, tag log.LogTag, msg string) {
	item := lm.logItem(level, tag, msg)
	if lm.isPrint(&item) {
		lm.console(&item)
	}
	lm.Buffer.Push(item)
}

func (lm *_LogManage) logItem(level log.LogLevel, tag log.LogTag, msg string) log.LogItem {
	item := log.LogItem{
		ID:      atomic.AddUint64(&lm.gID, 1),
		Level:   level,
		Time:    time.Now(),
		Tag:     tag,
		Message: msg,
	}
	if lm.Config.Global.ShowFile == 1 {
		item.File = lm.file()
	}
	return item
}

func (lm *_LogManage) file() log.LogFile {
	var logFile log.LogFile
	pc, file, line, ok := runtime.Caller(lm.Config.Global.FileCaller)
	if !ok {
		logFile.Name = "???"
		logFile.FuncName = "???"
		return logFile
	}
	logFile.Name = file
	logFile.ShortName = utils.SubstrByStartAfter(file, "/")
	logFile.Line = line
	logFile.FuncName = utils.SubstrByStartAfter(runtime.FuncForPC(pc).Name(), "/")
	return logFile
}

func (lm *_LogManage) isPrint(item *log.LogItem) bool {
	var isPrint bool
	switch lm.Config.Global.Rule {
	case log.AlwaysRule:
		isPrint = lm.isEitherTrue(item, func(lm *_LogManage, item *log.LogItem) bool {
			return lm.Config.Global.IsPrint == 1
		}, func(lm *_LogManage, item *log.LogItem) bool {
			var v bool
			lm.tags(item, func(t log.TagConfig) bool {
				if t.Config.IsPrint == 1 {
					v = true
					return true
				}
				return false
			})
			return v
		}, func(lm *_LogManage, item *log.LogItem) bool {
			var v bool
			lm.levels(item, func(l log.LevelConfig) bool {
				if l.Config.IsPrint == 1 {
					v = true
					return true
				}
				return false
			})
			return v
		})
	case log.GlobalRule:
		isPrint = lm.Config.Global.IsPrint == 1
	case log.TagRule:
		lm.tags(item, func(t log.TagConfig) bool {
			if t.Config.IsPrint == 1 {
				isPrint = true
				return true
			}
			return false
		})
	case log.LevelRule:
		lm.levels(item, func(l log.LevelConfig) bool {
			if l.Config.IsPrint == 1 {
				isPrint = true
				return true
			}
			return false
		})
	}
	return isPrint
}

func (lm *_LogManage) isEitherTrue(item *log.LogItem, fns ...func(*_LogManage, *log.LogItem) bool) bool {
	for _, fn := range fns {
		if fn(lm, item) {
			return true
		}
	}
	return false
}

func (lm *_LogManage) tags(item *log.LogItem, fn func(log.TagConfig) bool) {
	for _, tag := range lm.Config.Tags {
		for _, tagName := range tag.Names {
			if tagName == (*item).Tag {
				if fn(tag) {
					return
				}
				break
			}
		}
	}
}

func (lm *_LogManage) levels(item *log.LogItem, fn func(log.LevelConfig) bool) {
	for _, lev := range lm.Config.Levels {
		for _, levVal := range lev.Values {
			if (*item).Level == levVal {
				if fn(lev) {
					return
				}
				break
			}
		}
	}
}

func (lm *_LogManage) console(item *log.LogItem) {
	if lm.Config.Console.Level <= (*item).Level {
		lm.stdout(lm.Template[log.TmplConsole], lm.Template[log.TmplConsoleTime], item)
	}
}

func (lm *_LogManage) systemLog(tag log.LogLevel, msg string) {
	item := lm.logItem(tag, "SYSTEM", msg)
	lm.stdout(log.DefaultSystemTmpl, log.DefaultConsoleTimeTmpl, &item)
}

func (lm *_LogManage) stdout(tmpl, timetmpl interface{}, item *log.LogItem) {
	info := log.ParseLogItem(tmpl, timetmpl, item)
	os.Stdout.WriteString(info)
}

func (lm *_LogManage) execStore() {
	interval := time.Duration(lm.Config.Global.Interval) * time.Second
	time.AfterFunc(interval, func() {
		lm.store()
		lm.execStore()
	})
}

func (lm *_LogManage) store() {
	for {
		item, err := lm.Buffer.Pop()
		if err != nil {
			panic(err)
		}
		if item == nil {
			break
		}
		targets := lm.storeTargets(item)
		l := len(targets)
		if l == 0 {
			continue
		}
		wg := new(sync.WaitGroup)
		wg.Add(l)
		for i, l := 0, len(targets); i < l; i++ {
			go lm.writeStore(wg, targets[i], item)
		}
		wg.Wait()
		atomic.AddInt64(&lm.storeTotal, 1)
	}
}

func (lm *_LogManage) target(target map[string]string, ts string) {
	tsa := strings.Split(ts, ",")
	for i := 0; i < len(tsa); i++ {
		t := tsa[i]
		if t == "" {
			continue
		}
		if _, ok := target[t]; !ok {
			target[t] = t
		}
	}
}

func (lm *_LogManage) storeTargets(item *log.LogItem) (targets []string) {
	target := make(map[string]string)
	rule := lm.Config.Global.Rule

	if rule == log.AlwaysRule || rule == log.GlobalRule {
		if lm.Config.Global.Level <= item.Level {
			lm.target(target, lm.Config.Global.TargetStore)
		}
	}
	if rule == log.AlwaysRule || rule == log.TagRule {
		lm.tags(item, func(tag log.TagConfig) bool {
			if tag.Level <= item.Level {
				lm.target(target, tag.Config.TargetStore)
			}
			return false
		})
	}
	if rule == log.AlwaysRule || rule == log.LevelRule {
		lm.levels(item, func(lev log.LevelConfig) bool {
			lm.target(target, lev.Config.TargetStore)
			return false
		})
	}

	for k := range target {
		targets = append(targets, k)
	}

	return
}

func (lm *_LogManage) writeStore(wg *sync.WaitGroup, target string, item *log.LogItem) {
	defer wg.Done()
	store, ok := lm.Store[target]
	if !ok {
		return
	}
	err := store.Store(item)
	if err != nil {
		lm.systemLog(log.FATAL, fmt.Sprintf("Write store error:%s", err.Error()))
	}
}
<commit_msg>Fixed manage<commit_after>package manage

import (
	"fmt"
	"os"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"text/template"
	"time"

	"github.com/antlinker/alog/buffer"
	"github.com/antlinker/alog/log"
	"github.com/antlinker/alog/store"
	"github.com/antlinker/alog/utils"
)

// NewLogManage 创建新的LogManage实例
func NewLogManage(config *log.LogConfig) log.LogManage {
	manage := &_LogManage{
		Config: config,
	}
	manage.Template = map[log.TmplKey]*template.Template{
		log.TmplConsole:     template.Must(template.New("").Parse(config.Console.Item.Tmpl)),
		log.TmplConsoleTime: template.Must(template.New("").Parse(config.Console.Item.TimeTmpl)),
	}
	switch config.Global.Buffer.Engine {
	case log.REDIS_BUFFER:
		var redisConfig log.RedisConfig
		redisStore := config.Store.Redis
		if redisStore != nil {
			if v, ok := redisStore[config.Global.Buffer.TargetStore]; ok {
				redisConfig = v
			}
		}
		manage.Buffer = buffer.NewRedisBuffer(redisConfig)
	default:
		manage.Buffer = buffer.NewMemoryBuffer()
	}
	manageStore := make(map[string]log.LogStore)
	if fileStore := config.Store.File; fileStore != nil {
		for k, v := range fileStore {
			manageStore[k] = store.NewFileStore(v)
		}
	}
	if elasticStore := config.Store.Elastic; elasticStore != nil {
		for k, v := range elasticStore {
			manageStore[k] = store.NewElasticStore(v)
		}
	}
	if mongoStore := config.Store.Mongo; mongoStore != nil {
		for k, v := range mongoStore {
			manageStore[k] = store.NewMongoStore(v)
		}
	}
	manage.Store = manageStore

	if config.Global.IsEnabled == 1 {
		go manage.execStore()
	}

	return manage
}

// _LogManage
type _LogManage struct {
	sync.RWMutex
	gID        uint64
	storeTotal int64
	Config     *log.LogConfig
	Template   map[log.TmplKey]*template.Template
	Buffer     log.LogBuffer
	Store      map[string]log.LogStore
}

func (lm *_LogManage) Write(level log.LogLevel, tag log.LogTag, v ...interface{}) {
	if lm.Config.Global.Level >= level {
		return
	}
	msg := fmt.Sprint(v...)
	lm.writeMsg(level, tag, msg)
}

func (lm *_LogManage) Writef(level log.LogLevel, tag log.LogTag, format string, v ...interface{}) {
	if lm.Config.Global.Level >= level {
		return
	}
	msg := fmt.Sprintf(format, v...)
	lm.writeMsg(level, tag, msg)
}

func (lm *_LogManage) Console(level log.LogLevel, tag log.LogTag, v ...interface{}) {
	msg := fmt.Sprint(v...)
	lm.writeConsole(level, tag, msg)
}

func (lm *_LogManage) Consolef(level log.LogLevel, tag log.LogTag, format string, v ...interface{}) {
	msg := fmt.Sprintf(format, v...)
	lm.writeConsole(level, tag, msg)
}

func (lm *_LogManage) writeConsole(level log.LogLevel, tag log.LogTag, msg string) {
	item := lm.logItem(level, tag, msg)
	lm.console(&item)
}

func (lm *_LogManage) TotalNum() int64 {
	return lm.storeTotal
}

func (lm *_LogManage) writeMsg(level log.LogLevel, tag log.LogTag, msg string) {
	item := lm.logItem(level, tag, msg)
	if lm.isPrint(&item) {
		lm.console(&item)
	}
	lm.Buffer.Push(item)
}

func (lm *_LogManage) logItem(level log.LogLevel, tag log.LogTag, msg string) log.LogItem {
	item := log.LogItem{
		ID:      atomic.AddUint64(&lm.gID, 1),
		Level:   level,
		Time:    time.Now(),
		Tag:     tag,
		Message: msg,
	}
	if lm.Config.Global.ShowFile == 1 {
		item.File = lm.file()
	}
	return item
}

func (lm *_LogManage) file() log.LogFile {
	var logFile log.LogFile
	pc, file, line, ok := runtime.Caller(lm.Config.Global.FileCaller)
	if !ok {
		logFile.Name = "???"
		logFile.FuncName = "???"
		return logFile
	}
	logFile.Name = file
	logFile.ShortName = utils.SubstrByStartAfter(file, "/")
	logFile.Line = line
	logFile.FuncName = utils.SubstrByStartAfter(runtime.FuncForPC(pc).Name(), "/")
	return logFile
}

func (lm *_LogManage) isPrint(item *log.LogItem) bool {
	var isPrint bool
	switch lm.Config.Global.Rule {
	case log.AlwaysRule:
		isPrint = lm.isEitherTrue(item, func(lm *_LogManage, item *log.LogItem) bool {
			return lm.Config.Global.IsPrint == 1
		}, func(lm *_LogManage, item *log.LogItem) bool {
			var v bool
			lm.tags(item, func(t log.TagConfig) bool {
				if t.Config.IsPrint == 1 {
					v = true
					return true
				}
				return false
			})
			return v
		}, func(lm *_LogManage, item *log.LogItem) bool {
			var v bool
			lm.levels(item, func(l log.LevelConfig) bool {
				if l.Config.IsPrint == 1 {
					v = true
					return true
				}
				return false
			})
			return v
		})
	case log.GlobalRule:
		isPrint = lm.Config.Global.IsPrint == 1
	case log.TagRule:
		lm.tags(item, func(t log.TagConfig) bool {
			if t.Config.IsPrint == 1 {
				isPrint = true
				return true
			}
			return false
		})
	case log.LevelRule:
		lm.levels(item, func(l log.LevelConfig) bool {
			if l.Config.IsPrint == 1 {
				isPrint = true
				return true
			}
			return false
		})
	}
	return isPrint
}

func (lm *_LogManage) isEitherTrue(item *log.LogItem, fns ...func(*_LogManage, *log.LogItem) bool) bool {
	for _, fn := range fns {
		if fn(lm, item) {
			return true
		}
	}
	return false
}

func (lm *_LogManage) tags(item *log.LogItem, fn func(log.TagConfig) bool) {
	for _, tag := range lm.Config.Tags {
		for _, tagName := range tag.Names {
			if tagName == (*item).Tag {
				if fn(tag) {
					return
				}
				break
			}
		}
	}
}

func (lm *_LogManage) levels(item *log.LogItem, fn func(log.LevelConfig) bool) {
	for _, lev := range lm.Config.Levels {
		for _, levVal := range lev.Values {
			if (*item).Level == levVal {
				if fn(lev) {
					return
				}
				break
			}
		}
	}
}

func (lm *_LogManage) console(item *log.LogItem) {
	if lm.Config.Console.Level <= (*item).Level {
		lm.stdout(lm.Template[log.TmplConsole], lm.Template[log.TmplConsoleTime], item)
	}
}

func (lm *_LogManage) systemLog(tag log.LogLevel, msg string) {
	item := lm.logItem(tag, "SYSTEM", msg)
	lm.stdout(log.DefaultSystemTmpl, log.DefaultConsoleTimeTmpl, &item)
}

func (lm *_LogManage) stdout(tmpl, timetmpl interface{}, item *log.LogItem) {
	info := log.ParseLogItem(tmpl, timetmpl, item)
	os.Stdout.WriteString(info)
}

func (lm *_LogManage) execStore() {
	interval := time.Duration(lm.Config.Global.Interval) * time.Second
	time.AfterFunc(interval, func() {
		lm.store()
		lm.execStore()
	})
}

func (lm *_LogManage) store() {
	for {
		item, err := lm.Buffer.Pop()
		if err != nil {
			panic(err)
		}
		if item == nil {
			break
		}
		targets := lm.storeTargets(item)
		l := len(targets)
		if l == 0 {
			continue
		}
		for i, l := 0, len(targets); i < l; i++ {
			lm.writeStore(targets[i], item)
		}
		atomic.AddInt64(&lm.storeTotal, 1)
		if item.Level == log.FATAL {
			os.Exit(1)
		}
	}
}

func (lm *_LogManage) target(target map[string]string, ts string) {
	tsa := strings.Split(ts, ",")
	for i := 0; i < len(tsa); i++ {
		t := tsa[i]
		if t == "" {
			continue
		}
		if _, ok := target[t]; !ok {
			target[t] = t
		}
	}
}

func (lm *_LogManage) storeTargets(item *log.LogItem) (targets []string) {
	target := make(map[string]string)
	rule := lm.Config.Global.Rule

	if rule == log.AlwaysRule || rule == log.GlobalRule {
		if lm.Config.Global.Level <= item.Level {
			lm.target(target, lm.Config.Global.TargetStore)
		}
	}
	if rule == log.AlwaysRule || rule == log.TagRule {
		lm.tags(item, func(tag log.TagConfig) bool {
			if tag.Level <= item.Level {
				lm.target(target, tag.Config.TargetStore)
			}
			return false
		})
	}
	if rule == log.AlwaysRule || rule == log.LevelRule {
		lm.levels(item, func(lev log.LevelConfig) bool {
			lm.target(target, lev.Config.TargetStore)
			return false
		})
	}

	for k := range target {
		targets = append(targets, k)
	}

	return
}

func (lm *_LogManage) writeStore(target string, item *log.LogItem) {
	store, ok := lm.Store[target]
	if !ok {
		return
	}
	err := store.Store(item)
	if err != nil {
		lm.systemLog(log.FATAL, fmt.Sprintf("Write store error:%s", err.Error()))
	}
}
<|endoftext|>


<commit_before>package dnews

import (
	"database/sql"
	"fmt"

	// postgresql
	"github.com/lib/pq"
	"github.com/qbit/pgenv"
)

// DBConnect returns a connection to the database
func DBConnect() (*sql.DB, error) {
	var cstr = pgenv.ConnStr{}
	cstr.SetDefaults()

	return sql.Open("postgres", cstr.ToString())
}

// Auth checks a user's username / password for login
func Auth(u string, p string) (*User, error) {
	var user = &User{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	err = db.QueryRow(`select id, created, fname, lname, email, username, (hash = crypt($1, hash)) as authed, admin from users where username = $2`, p, u).Scan(&user.ID, &user.Created, &user.FName, &user.LName, &user.Email, &user.User, &user.Authed, &user.Admin)
	if err != nil {
		return nil, err
	}

	defer db.Close()

	return user, nil
}

// GetRawArticle returns the raw markdown for a given article
func GetRawArticle(slug string) (*Article, error) {

	var a = Article{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	err = db.QueryRow(`
SELECT
 body
from articles
where
  slug = $1
`, slug).Scan(&a.Body)
	if err != nil {
		return nil, err
	}

	defer db.Close()

	return &a, nil
}

// GetBugs grabs all the bugs in the db
func GetBugs() (*Bugs, error) {
	var bs = Bugs{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	rows, err := db.Query(`select * from bugs`)
	if err != nil {
		return nil, err
	}

	for rows.Next() {
		var b = Bug{}
		rows.Scan(&b.ID, &b.Created, &b.Name, &b.Descr, &b.URL)

		bs = append(bs, &b)
	}

	return &bs, nil
}

// GetArticle returns the raw markdown for a given article
func GetArticle(slug string) (*Article, error) {
	var a = Article{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	err = db.QueryRow(`
SELECT
 articles.id,
 slug,
 published,
 title,
 body,
 key,
 email,
 fname,
 lname,
 sig
from articles
join users on
  (articles.authorid = users.id)
join pubkeys on
  (pubkeys.userid = users.id)
where
  articles.slug = $1
`, slug).Scan(&a.ID, &a.Slug, &a.Date, &a.Title, &a.Body, &a.Author.Pubkey, &a.Author.Email, &a.Author.FName, &a.Author.LName, &a.Signature)
	if err != nil {
		return nil, err
	}

	t, err := GetTags(a.ID, db)
	if err != nil {
		return nil, err
	}

	a.Tags = t
	a.HTML()

	defer db.Close()

	return &a, nil
}

// GetTagIDS takes a list of tag names and returns a set of tag ids
func GetTagIDS(s []string, db *sql.DB) (tagIDS []int, err error) {
	sql := `
select
  id
from tags
where
  name = ANY($1)
`
	rows, err := db.Query(sql, pq.Array(s))
	if err != nil {
		return nil, err
	}

	defer rows.Close()

	for rows.Next() {
		var i int
		rows.Scan(&i)
		tagIDS = append(tagIDS, i)
	}

	return tagIDS, nil
}

// GetAllTags returns all the tags in the DB
func GetAllTags() (Tags, error) {
	var ts = Tags{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	rows, err := db.Query(`select id, created, name from tags`)

	defer rows.Close()
	defer db.Close()

	for rows.Next() {
		var t = Tag{}
		err := rows.Scan(&t.ID, &t.Created, &t.Name)
		if err != nil {
			return nil, err
		}
		ts = append(ts, &t)
	}

	return ts, nil
}

// GetAllUsers gets all the users in the DB
func GetAllUsers() (Users, error) {
	var us = Users{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	rows, err := db.Query(`select id, created, fname, lname, email, username, admin from users`)

	defer rows.Close()
	defer db.Close()

	for rows.Next() {
		var u = User{}
		err := rows.Scan(&u.ID, &u.Created, &u.FName, &u.LName, &u.Email, &u.User, &u.Admin)
		if err != nil {
			return nil, err
		}
		us = append(us, &u)
	}

	return us, nil
}

// GetTags returns tags for a given article
func GetTags(id int, db *sql.DB) (Tags, error) {
	var ts = Tags{}
	rows, err := db.Query(`
		select
		tags.id,
		tags.name
		from article_tags
		join tags on
		(article_tags.tagid = tags.id)
		where
		articleid = $1
		`, id)
	if err != nil {
		return nil, err
	}

	defer rows.Close()

	for rows.Next() {
		var t = Tag{}
		err := rows.Scan(&t.ID, &t.Name)
		if err != nil {
			return nil, err
		}
		ts = append(ts, &t)
	}
	return ts, nil
}

// GetUserIDByEmail takes a users email and returns a User that is associated with it
func GetUserIDByEmail(e string, db *sql.DB) (id *int, err error) {
	err = db.QueryRow(`
select id from users where email = $1
`, e).Scan(&id)
	if err != nil {
		return nil, err
	}

	return id, nil
}

// GetArticlesByTag tags a tag and returns all the matching articles
func GetArticlesByTag(t string) (Articles, error) {
	var as = Articles{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	rows, err := db.Query(`
		SELECT
		articles.id,
		slug,
		published,
		title,
		body,
		key,
		email,
		fname,
		lname,
		sig
		from articles
		join users on
		(articles.authorid = users.id)
		join pubkeys on
		(pubkeys.userid = users.id)
		join article_tags on
		(article_tags.articleid = articles.id)
		join tags on
		(article_tags.tagid = tags.id)
		where
		live = true and
		tags.name = $1
		order by published desc
		`, t)
	if err != nil {
		return nil, err
	}

	defer rows.Close()
	defer db.Close()

	for rows.Next() {
		var a = Article{}
		err := rows.Scan(&a.ID, &a.Slug, &a.Date, &a.Title, &a.Body, &a.Author.Pubkey, &a.Author.Email, &a.Author.FName, &a.Author.LName, &a.Signature)
		if err != nil {
			return nil, err
		}
		t, err := GetTags(a.ID, db)
		if err != nil {
			return nil, err
		}
		a.Tags = t
		a.Verify(a.Author.Pubkey)
		a.HTML()
		as = append(as, &a)
	}

	return as, nil
}

// GetNArticles returns N most recent articles from the DB
func GetNArticles(n int) (Articles, error) {
	var as = Articles{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	rows, err := db.Query(`
		SELECT
		articles.id,
		slug,
		published,
		title,
		body,
		key,
		email,
		fname,
		lname,
		sig
		from articles
		join users on
		(articles.authorid = users.id)
		join pubkeys on
		(pubkeys.userid = users.id)
		where
		live = true
		order by published desc
		limit $1
		`, n)
	if err != nil {
		return nil, err
	}

	defer rows.Close()
	defer db.Close()

	for rows.Next() {
		var a = Article{}
		err := rows.Scan(&a.ID, &a.Slug, &a.Date, &a.Title, &a.Body, &a.Author.Pubkey, &a.Author.Email, &a.Author.FName, &a.Author.LName, &a.Signature)
		if err != nil {
			return nil, err
		}
		t, err := GetTags(a.ID, db)
		if err != nil {
			return nil, err
		}
		a.Tags = t
		a.Verify(a.Author.Pubkey)
		a.HTML()
		as = append(as, &a)
	}

	return as, nil
}

// SearchArticles uses pg's TS stuff to query all the articles for passed in values
func SearchArticles(query string, limit int) (Articles, error) {
	var as = Articles{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	rows, err := db.Query(`
		SELECT
		aid as id,
		slug,
		published,
		title,
		body,
		key,
		email,
		fname,
		lname,
		sig,
		ts_headline(body, q) as headline,
		rank
		FROM (
			SELECT
			*,
			articles.id as aid,
			ts_rank_cd(tsv, q) as rank
			FROM articles
			join users on
			(articles.authorid = users.id)
			join pubkeys on
			(pubkeys.userid = users.id), plainto_tsquery($1) q
			WHERE tsv @@ q
			ORDER BY rank DESC
			LIMIT $2) AS foo;
		`, query, limit)
	if err != nil {
		return nil, err
	}

	defer rows.Close()
	defer db.Close()

	for rows.Next() {
		var a = Article{}
		err := rows.Scan(&a.ID, &a.Slug, &a.Date, &a.Title, &a.Body, &a.Author.Pubkey, &a.Author.Email, &a.Author.FName, &a.Author.LName, &a.Signature, &a.Headline, &a.Rank)
		if err != nil {
			return nil, err
		}

		a.Verify(a.Author.Pubkey)
		a.HTML()

		as = append(as, &a)
	}

	return as, nil
}

// InsertUser takes a User and inserts them into the database
func InsertUser(u User) (*int, error) {
	var id int
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	err = db.QueryRow(`INSERT INTO users (fname, lname, email, username, hash) values ($1, $2, $3, (select hash($4)))`, u.FName, u.LName, u.Email, u.User, u.Pass).Scan(&id)
	if err != nil {
		return nil, err
	}
	return &id, nil
}

// InsertArticle takes an Article and inserts it into the db, it will verify the Author exists
// prior to inserting
func InsertArticle(a Article) (*int, error) {
	var id int
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}

	uid, err := AssignUser(a.Author.Email, db)
	if err != nil {
		return nil, err
	}

	a.AuthorID = *uid

	fmt.Printf("AuthorID: %d\n", a.AuthorID)

	err = db.QueryRow(`INSERT INTO articles (title, body, created, live, sig, authorid) values ($1, $2, $3, $4, $5, $6) returning id`, a.Title, a.Body, a.Date, a.Live, a.Signature, a.AuthorID).Scan(&id)
	if err != nil {
		return nil, err
	}

	a.ID = id

	tags, err := GetTagIDS(a.Tags.Join(), db)
	if err != nil {
		return nil, err
	}

	err = AssignTags(tags, a.ID, db)
	if err != nil {
		return nil, err
	}

	return &id, nil
}

// AssignUser takes a article (with user already assigned), gets the ID of said user from the db, and creates
// the association assuming the user exists in the db.
func AssignUser(e string, db *sql.DB) (*int, error) {
	uid, err := GetUserIDByEmail(e, db)
	if err != nil {
		return nil, err
	}

	return uid, nil
}

// AssignTags takes a set of tag ids and an article id and creates the association in the article_tags table
func AssignTags(ts []int, id int, db *sql.DB) error {
	txn, err := db.Begin()
	if err != nil {
		return err
	}

	stmt, err := txn.Prepare(pq.CopyIn("article_tags", "articleid", "tagid"))
	if err != nil {
		return err
	}

	for _, tid := range ts {
		_, err = stmt.Exec(id, tid)
		if err != nil {
			return err
		}
	}
	_, err = stmt.Exec()
	if err != nil {
		return err
	}

	err = stmt.Close()
	if err != nil {
		return err
	}

	err = txn.Commit()
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>close connections<commit_after>package dnews

import (
	"database/sql"
	"fmt"

	// postgresql
	"github.com/lib/pq"
	"github.com/qbit/pgenv"
)

// DBConnect returns a connection to the database
func DBConnect() (*sql.DB, error) {
	var cstr = pgenv.ConnStr{}
	cstr.SetDefaults()

	return sql.Open("postgres", cstr.ToString())
}

// Auth checks a user's username / password for login
func Auth(u string, p string) (*User, error) {
	var user = &User{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	err = db.QueryRow(`select id, created, fname, lname, email, username, (hash = crypt($1, hash)) as authed, admin from users where username = $2`, p, u).Scan(&user.ID, &user.Created, &user.FName, &user.LName, &user.Email, &user.User, &user.Authed, &user.Admin)
	if err != nil {
		return nil, err
	}

	defer db.Close()

	return user, nil
}

// GetRawArticle returns the raw markdown for a given article
func GetRawArticle(slug string) (*Article, error) {

	var a = Article{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	err = db.QueryRow(`
SELECT
 body
from articles
where
  slug = $1
`, slug).Scan(&a.Body)
	if err != nil {
		return nil, err
	}

	defer db.Close()

	return &a, nil
}

// GetBugs grabs all the bugs in the db
func GetBugs() (*Bugs, error) {
	var bs = Bugs{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	rows, err := db.Query(`select * from bugs`)

	defer rows.Close()
	defer db.Close()

	if err != nil {
		return nil, err
	}

	for rows.Next() {
		var b = Bug{}
		rows.Scan(&b.ID, &b.Created, &b.Name, &b.Descr, &b.URL)

		bs = append(bs, &b)
	}

	return &bs, nil
}

// GetArticle returns the raw markdown for a given article
func GetArticle(slug string) (*Article, error) {
	var a = Article{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	err = db.QueryRow(`
SELECT
 articles.id,
 slug,
 published,
 title,
 body,
 key,
 email,
 fname,
 lname,
 sig
from articles
join users on
  (articles.authorid = users.id)
join pubkeys on
  (pubkeys.userid = users.id)
where
  articles.slug = $1
`, slug).Scan(&a.ID, &a.Slug, &a.Date, &a.Title, &a.Body, &a.Author.Pubkey, &a.Author.Email, &a.Author.FName, &a.Author.LName, &a.Signature)
	if err != nil {
		return nil, err
	}

	t, err := GetTags(a.ID, db)
	if err != nil {
		return nil, err
	}

	a.Tags = t
	a.HTML()

	defer db.Close()

	return &a, nil
}

// GetTagIDS takes a list of tag names and returns a set of tag ids
func GetTagIDS(s []string, db *sql.DB) (tagIDS []int, err error) {
	sql := `
select
  id
from tags
where
  name = ANY($1)
`
	rows, err := db.Query(sql, pq.Array(s))
	if err != nil {
		return nil, err
	}

	defer rows.Close()

	for rows.Next() {
		var i int
		rows.Scan(&i)
		tagIDS = append(tagIDS, i)
	}

	return tagIDS, nil
}

// GetAllTags returns all the tags in the DB
func GetAllTags() (Tags, error) {
	var ts = Tags{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	rows, err := db.Query(`select id, created, name from tags`)

	defer rows.Close()
	defer db.Close()

	for rows.Next() {
		var t = Tag{}
		err := rows.Scan(&t.ID, &t.Created, &t.Name)
		if err != nil {
			return nil, err
		}
		ts = append(ts, &t)
	}

	return ts, nil
}

// GetAllUsers gets all the users in the DB
func GetAllUsers() (Users, error) {
	var us = Users{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	rows, err := db.Query(`select id, created, fname, lname, email, username, admin from users`)

	defer rows.Close()
	defer db.Close()

	for rows.Next() {
		var u = User{}
		err := rows.Scan(&u.ID, &u.Created, &u.FName, &u.LName, &u.Email, &u.User, &u.Admin)
		if err != nil {
			return nil, err
		}
		us = append(us, &u)
	}

	return us, nil
}

// GetTags returns tags for a given article
func GetTags(id int, db *sql.DB) (Tags, error) {
	var ts = Tags{}
	rows, err := db.Query(`
		select
		tags.id,
		tags.name
		from article_tags
		join tags on
		(article_tags.tagid = tags.id)
		where
		articleid = $1
		`, id)
	if err != nil {
		return nil, err
	}

	defer rows.Close()

	for rows.Next() {
		var t = Tag{}
		err := rows.Scan(&t.ID, &t.Name)
		if err != nil {
			return nil, err
		}
		ts = append(ts, &t)
	}
	return ts, nil
}

// GetUserIDByEmail takes a users email and returns a User that is associated with it
func GetUserIDByEmail(e string, db *sql.DB) (id *int, err error) {
	err = db.QueryRow(`
select id from users where email = $1
`, e).Scan(&id)
	if err != nil {
		return nil, err
	}

	return id, nil
}

// GetArticlesByTag tags a tag and returns all the matching articles
func GetArticlesByTag(t string) (Articles, error) {
	var as = Articles{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	rows, err := db.Query(`
		SELECT
		articles.id,
		slug,
		published,
		title,
		body,
		key,
		email,
		fname,
		lname,
		sig
		from articles
		join users on
		(articles.authorid = users.id)
		join pubkeys on
		(pubkeys.userid = users.id)
		join article_tags on
		(article_tags.articleid = articles.id)
		join tags on
		(article_tags.tagid = tags.id)
		where
		live = true and
		tags.name = $1
		order by published desc
		`, t)
	if err != nil {
		return nil, err
	}

	defer rows.Close()
	defer db.Close()

	for rows.Next() {
		var a = Article{}
		err := rows.Scan(&a.ID, &a.Slug, &a.Date, &a.Title, &a.Body, &a.Author.Pubkey, &a.Author.Email, &a.Author.FName, &a.Author.LName, &a.Signature)
		if err != nil {
			return nil, err
		}
		t, err := GetTags(a.ID, db)
		if err != nil {
			return nil, err
		}
		a.Tags = t
		a.Verify(a.Author.Pubkey)
		a.HTML()
		as = append(as, &a)
	}

	return as, nil
}

// GetNArticles returns N most recent articles from the DB
func GetNArticles(n int) (Articles, error) {
	var as = Articles{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	rows, err := db.Query(`
		SELECT
		articles.id,
		slug,
		published,
		title,
		body,
		key,
		email,
		fname,
		lname,
		sig
		from articles
		join users on
		(articles.authorid = users.id)
		join pubkeys on
		(pubkeys.userid = users.id)
		where
		live = true
		order by published desc
		limit $1
		`, n)
	if err != nil {
		return nil, err
	}

	defer rows.Close()
	defer db.Close()

	for rows.Next() {
		var a = Article{}
		err := rows.Scan(&a.ID, &a.Slug, &a.Date, &a.Title, &a.Body, &a.Author.Pubkey, &a.Author.Email, &a.Author.FName, &a.Author.LName, &a.Signature)
		if err != nil {
			return nil, err
		}
		t, err := GetTags(a.ID, db)
		if err != nil {
			return nil, err
		}
		a.Tags = t
		a.Verify(a.Author.Pubkey)
		a.HTML()
		as = append(as, &a)
	}

	return as, nil
}

// SearchArticles uses pg's TS stuff to query all the articles for passed in values
func SearchArticles(query string, limit int) (Articles, error) {
	var as = Articles{}
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	rows, err := db.Query(`
		SELECT
		aid as id,
		slug,
		published,
		title,
		body,
		key,
		email,
		fname,
		lname,
		sig,
		ts_headline(body, q) as headline,
		rank
		FROM (
			SELECT
			*,
			articles.id as aid,
			ts_rank_cd(tsv, q) as rank
			FROM articles
			join users on
			(articles.authorid = users.id)
			join pubkeys on
			(pubkeys.userid = users.id), plainto_tsquery($1) q
			WHERE tsv @@ q
			ORDER BY rank DESC
			LIMIT $2) AS foo;
		`, query, limit)
	if err != nil {
		return nil, err
	}

	defer rows.Close()
	defer db.Close()

	for rows.Next() {
		var a = Article{}
		err := rows.Scan(&a.ID, &a.Slug, &a.Date, &a.Title, &a.Body, &a.Author.Pubkey, &a.Author.Email, &a.Author.FName, &a.Author.LName, &a.Signature, &a.Headline, &a.Rank)
		if err != nil {
			return nil, err
		}

		a.Verify(a.Author.Pubkey)
		a.HTML()

		as = append(as, &a)
	}

	return as, nil
}

// InsertUser takes a User and inserts them into the database
func InsertUser(u User) (*int, error) {
	var id int
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}
	defer db.Close()
	err = db.QueryRow(`INSERT INTO users (fname, lname, email, username, hash) values ($1, $2, $3, (select hash($4)))`, u.FName, u.LName, u.Email, u.User, u.Pass).Scan(&id)
	if err != nil {
		return nil, err
	}
	return &id, nil
}

// InsertArticle takes an Article and inserts it into the db, it will verify the Author exists
// prior to inserting
func InsertArticle(a Article) (*int, error) {
	var id int
	db, err := DBConnect()
	if err != nil {
		return nil, err
	}

	defer db.Close()

	uid, err := AssignUser(a.Author.Email, db)
	if err != nil {
		return nil, err
	}

	a.AuthorID = *uid

	fmt.Printf("AuthorID: %d\n", a.AuthorID)

	err = db.QueryRow(`INSERT INTO articles (title, body, created, live, sig, authorid) values ($1, $2, $3, $4, $5, $6) returning id`, a.Title, a.Body, a.Date, a.Live, a.Signature, a.AuthorID).Scan(&id)
	if err != nil {
		return nil, err
	}

	a.ID = id

	tags, err := GetTagIDS(a.Tags.Join(), db)
	if err != nil {
		return nil, err
	}

	err = AssignTags(tags, a.ID, db)
	if err != nil {
		return nil, err
	}

	return &id, nil
}

// AssignUser takes a article (with user already assigned), gets the ID of said user from the db, and creates
// the association assuming the user exists in the db.
func AssignUser(e string, db *sql.DB) (*int, error) {
	uid, err := GetUserIDByEmail(e, db)
	if err != nil {
		return nil, err
	}

	return uid, nil
}

// AssignTags takes a set of tag ids and an article id and creates the association in the article_tags table
func AssignTags(ts []int, id int, db *sql.DB) error {
	txn, err := db.Begin()
	if err != nil {
		return err
	}

	stmt, err := txn.Prepare(pq.CopyIn("article_tags", "articleid", "tagid"))
	if err != nil {
		return err
	}

	for _, tid := range ts {
		_, err = stmt.Exec(id, tid)
		if err != nil {
			return err
		}
	}
	_, err = stmt.Exec()
	if err != nil {
		return err
	}

	err = stmt.Close()
	if err != nil {
		return err
	}

	err = txn.Commit()
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package gogitlab

import (
	"encoding/json"
)

const (
	projects_url         = "/projects"                         // Get a list of projects owned by the authenticated user
	projects_search_url  = "/projects/search/:query"           // Search for projects by name
	project_url          = "/projects/:id"                     // Get a specific project, identified by project ID or NAME
	project_url_events   = "/projects/:id/events"              // Get project events
	project_url_branches = "/projects/:id/repository/branches" // Lists all branches of a project
	project_url_members  = "/projects/:id/members"             // List project team members
	project_url_member   = "/projects/:id/members/:user_id"    // Get project team member
)

type Member struct {
	Id        int
	Username  string
	Email     string
	Name      string
	State     string
	CreatedAt string `json:"created_at,omitempty"`
	// AccessLevel int
}

type Namespace struct {
	Id          int
	Name        string
	Path        string
	Description string
	Owner_Id    int
	Created_At  string
	Updated_At  string
}

// A gitlab project
type Project struct {
	Id                   int        `json:"id,omitempty"`
	Name                 string     `json:"name,omitempty"`
	Description          string     `json:"description,omitempty"`
	DefaultBranch        string     `json:"default_branch,omitempty"`
	Owner                *Member    `json:"owner,omitempty"`
	Public               bool       `json:"public,omitempty"`
	Path                 string     `json:"path,omitempty"`
	PathWithNamespace    string     `json:"path_with_namespace,omitempty"`
	IssuesEnabled        bool       `json:"issues_enabled,omitempty"`
	MergeRequestsEnabled bool       `json:"merge_requests_enabled,omitempty"`
	WallEnabled          bool       `json:"wall_enabled,omitempty"`
	WikiEnabled          bool       `json:"wiki_enabled,omitempty"`
	CreatedAtRaw         string     `json:"created_at,omitempty"`
	Namespace            *Namespace `json:"namespace,omitempty"`
	SshRepoUrl           string     `json:"ssh_url_to_repo"`
	HttpRepoUrl          string     `json:"http_url_to_repo"`
}

/*
Get a list of projects owned by the authenticated user.
*/
func (g *Gitlab) Projects() ([]*Project, error) {

	url := g.ResourceUrl(projects_url, nil)

	var projects []*Project

	contents, err := g.buildAndExecRequest("GET", url, nil)
	if err == nil {
		err = json.Unmarshal(contents, &projects)
	}

	return projects, err
}

/*
Get a specific project, identified by project ID or NAME,
which is owned by the authentication user.
Namespaced project may be retrieved by specifying the namespace
and its project name like this:

	`namespace%2Fproject-name`

*/
func (g *Gitlab) Project(id string) (*Project, error) {

	url, opaque := g.ResourceUrlRaw(project_url, map[string]string{":id": id})

	var project *Project

	contents, err := g.buildAndExecRequestRaw("GET", url, opaque, nil)
	if err == nil {
		err = json.Unmarshal(contents, &project)
	}

	return project, err
}

/*
Lists all branches of a project.
*/
func (g *Gitlab) ProjectBranches(id string) ([]*Branch, error) {

	url, opaque := g.ResourceUrlRaw(project_url_branches, map[string]string{":id": id})

	var branches []*Branch

	contents, err := g.buildAndExecRequestRaw("GET", url, opaque, nil)
	if err == nil {
		err = json.Unmarshal(contents, &branches)
	}

	return branches, err
}

func (g *Gitlab) ProjectMembers(id string) ([]*Member, error) {
	url, opaque := g.ResourceUrlRaw(project_url_members, map[string]string{":id": id})

	var members []*Member

	contents, err := g.buildAndExecRequestRaw("GET", url, opaque, nil)
	if err == nil {
		err = json.Unmarshal(contents, &members)
	}

	return members, err
}
<commit_msg>Added minimal support for the project API<commit_after>package gogitlab

import (
	"encoding/json"
)

const (
	projects_url         = "/projects"                         // Get a list of projects owned by the authenticated user
	projects_search_url  = "/projects/search/:query"           // Search for projects by name
	project_url          = "/projects/:id"                     // Get a specific project, identified by project ID or NAME
	project_url_events   = "/projects/:id/events"              // Get project events
	project_url_branches = "/projects/:id/repository/branches" // Lists all branches of a project
	project_url_members  = "/projects/:id/members"             // List project team members
	project_url_member   = "/projects/:id/members/:user_id"    // Get project team member
)

type Member struct {
	Id        int
	Username  string
	Email     string
	Name      string
	State     string
	CreatedAt string `json:"created_at,omitempty"`
	// AccessLevel int
}

type Namespace struct {
	Id          int
	Name        string
	Path        string
	Description string
	Owner_Id    int
	Created_At  string
	Updated_At  string
}

// A gitlab project
type Project struct {
	Id                   int        `json:"id,omitempty"`
	Name                 string     `json:"name,omitempty"`
	Description          string     `json:"description,omitempty"`
	DefaultBranch        string     `json:"default_branch,omitempty"`
	Owner                *Member    `json:"owner,omitempty"`
	Public               bool       `json:"public,omitempty"`
	Path                 string     `json:"path,omitempty"`
	PathWithNamespace    string     `json:"path_with_namespace,omitempty"`
	IssuesEnabled        bool       `json:"issues_enabled,omitempty"`
	MergeRequestsEnabled bool       `json:"merge_requests_enabled,omitempty"`
	WallEnabled          bool       `json:"wall_enabled,omitempty"`
	WikiEnabled          bool       `json:"wiki_enabled,omitempty"`
	CreatedAtRaw         string     `json:"created_at,omitempty"`
	Namespace            *Namespace `json:"namespace,omitempty"`
	SshRepoUrl           string     `json:"ssh_url_to_repo"`
	HttpRepoUrl          string     `json:"http_url_to_repo"`
}

// A gitlab project
type ProjectRequest struct {
	Name      string `json:"name"`
	Namespace int    `json:"namespace_id,omitempty"`
}

/*
Get a list of projects owned by the authenticated user.
*/
func (g *Gitlab) Projects() ([]*Project, error) {

	url := g.ResourceUrl(projects_url, nil)

	var projects []*Project

	contents, err := g.buildAndExecRequest("GET", url, nil)
	if err == nil {
		err = json.Unmarshal(contents, &projects)
	}

	return projects, err
}

/*
Get a specific project, identified by project ID or NAME,
which is owned by the authentication user.
Namespaced project may be retrieved by specifying the namespace
and its project name like this:

	`namespace%2Fproject-name`

*/
func (g *Gitlab) Project(id string) (*Project, error) {

	url, opaque := g.ResourceUrlRaw(project_url, map[string]string{":id": id})

	var project *Project

	contents, err := g.buildAndExecRequestRaw("GET", url, opaque, nil)
	if err == nil {
		err = json.Unmarshal(contents, &project)
	}

	return project, err
}

/*
Lists all branches of a project.
*/
func (g *Gitlab) ProjectBranches(id string) ([]*Branch, error) {

	url, opaque := g.ResourceUrlRaw(project_url_branches, map[string]string{":id": id})

	var branches []*Branch

	contents, err := g.buildAndExecRequestRaw("GET", url, opaque, nil)
	if err == nil {
		err = json.Unmarshal(contents, &branches)
	}

	return branches, err
}

func (g *Gitlab) ProjectMembers(id string) ([]*Member, error) {
	url, opaque := g.ResourceUrlRaw(project_url_members, map[string]string{":id": id})

	var members []*Member

	contents, err := g.buildAndExecRequestRaw("GET", url, opaque, nil)
	if err == nil {
		err = json.Unmarshal(contents, &members)
	}

	return members, err
}

/*
Add new project .

    POST /projects

Parameters:

    CreateProject	Object containing the ID and NAMESPACE/PROJECT_NAME of a project

*/
func (g *Gitlab) AddProject(req *ProjectRequest) (project *Project, err error) {

	params := map[string]string{
		":id": "1234",
	}
	u := g.ResourceUrl(projects_url, params)

	encodedRequest, err := json.Marshal(req)
	if err != nil {
		return
	}

	data, err := g.buildAndExecRequest("POST", u, encodedRequest)
	if err != nil {
		return
	}

	project = new(Project)
	err = json.Unmarshal(data, project)
	if err != nil {
		panic(err)
	}
	return

}
<|endoftext|>


<commit_before>package cfschema

import (
	"encoding/json"
)

const (
	PropertyFormatDate                = "date"
	PropertyFormatDateTime            = "date-time"
	PropertyFormatEmail               = "email"
	PropertyFormatHostname            = "hostname"
	PropertyFormatIdnEmail            = "idn-email"
	PropertyFormatIdnHostname         = "idn-hostname"
	PropertyFormatIpv4                = "ipv4"
	PropertyFormatIpv6                = "ipv6"
	PropertyFormatIri                 = "iri"
	PropertyFormatIriReference        = "iri-reference"
	PropertyFormatJsonPointer         = "json-pointer"
	PropertyFormatRegex               = "regex"
	PropertyFormatRelativeJsonPointer = "relative-json-pointer"
	PropertyFormatTime                = "time"
	PropertyFormatUri                 = "uri"
	PropertyFormatUriReference        = "uri-reference"
	PropertyFormatUriTemplate         = "uri-template"
)

const (
	PropertyTypeArray   = "array"
	PropertyTypeBoolean = "boolean"
	PropertyTypeInteger = "integer"
	PropertyTypeNull    = "null"
	PropertyTypeNumber  = "number"
	PropertyTypeObject  = "object"
	PropertyTypeString  = "string"
)

// Property represents the CloudFormation Resource Schema customization for Definitions and Properties.
type Property struct {
	AdditionalProperties *bool                `json:"additionalProperties,omitempty"`
	Comment              *string              `json:"$comment,omitempty"`
	Description          *string              `json:"description,omitempty"`
	Enum                 []interface{}        `json:"enum,omitempty"`
	Examples             []interface{}        `json:"examples,omitempty"`
	Format               *string              `json:"format,omitempty"`
	InsertionOrder       *bool                `json:"insertionOrder,omitempty"`
	Items                *Property            `json:"items,omitempty"`
	MaxItems             *int                 `json:"maxItems,omitempty"`
	MaxLength            *int                 `json:"maxLength,omitempty"`
	MinItems             *int                 `json:"minItems,omitempty"`
	MinLength            *int                 `json:"minLength,omitempty"`
	Pattern              *string              `json:"pattern,omitempty"`
	PatternProperties    map[string]*Property `json:"patternProperties,omitempty"`
	Properties           map[string]*Property `json:"properties,omitempty"`
	Ref                  *Reference           `json:"$ref,omitempty"`
	Required             []string             `json:"required,omitempty"`
	Type                 *Type                `json:"type,omitempty"`
	UniqueItems          *bool                `json:"uniqueItems,omitempty"`
}

// String returns a string representation of Property.
func (p *Property) String() string {
	if p == nil {
		return ""
	}

	b, _ := json.MarshalIndent(p, "", "  ")

	return string(b)
}

func (p *Property) IsRequired(name string) bool {
	if p == nil {
		return false
	}

	for _, req := range p.Required {
		if req == name {
			return true
		}
	}

	return false
}
<commit_msg>Add 'maximum' and 'minimum' to Property.<commit_after>package cfschema

import (
	"encoding/json"
)

const (
	PropertyFormatDate                = "date"
	PropertyFormatDateTime            = "date-time"
	PropertyFormatEmail               = "email"
	PropertyFormatHostname            = "hostname"
	PropertyFormatIdnEmail            = "idn-email"
	PropertyFormatIdnHostname         = "idn-hostname"
	PropertyFormatIpv4                = "ipv4"
	PropertyFormatIpv6                = "ipv6"
	PropertyFormatIri                 = "iri"
	PropertyFormatIriReference        = "iri-reference"
	PropertyFormatJsonPointer         = "json-pointer"
	PropertyFormatRegex               = "regex"
	PropertyFormatRelativeJsonPointer = "relative-json-pointer"
	PropertyFormatTime                = "time"
	PropertyFormatUri                 = "uri"
	PropertyFormatUriReference        = "uri-reference"
	PropertyFormatUriTemplate         = "uri-template"
)

const (
	PropertyTypeArray   = "array"
	PropertyTypeBoolean = "boolean"
	PropertyTypeInteger = "integer"
	PropertyTypeNull    = "null"
	PropertyTypeNumber  = "number"
	PropertyTypeObject  = "object"
	PropertyTypeString  = "string"
)

// Property represents the CloudFormation Resource Schema customization for Definitions and Properties.
type Property struct {
	AdditionalProperties *bool                `json:"additionalProperties,omitempty"`
	Comment              *string              `json:"$comment,omitempty"`
	Description          *string              `json:"description,omitempty"`
	Enum                 []interface{}        `json:"enum,omitempty"`
	Examples             []interface{}        `json:"examples,omitempty"`
	Format               *string              `json:"format,omitempty"`
	InsertionOrder       *bool                `json:"insertionOrder,omitempty"`
	Items                *Property            `json:"items,omitempty"`
	Maximum              *int                 `json:"maximum,omitempty"`
	MaxItems             *int                 `json:"maxItems,omitempty"`
	MaxLength            *int                 `json:"maxLength,omitempty"`
	Minimum              *int                 `json:"minimum,omitempty"`
	MinItems             *int                 `json:"minItems,omitempty"`
	MinLength            *int                 `json:"minLength,omitempty"`
	Pattern              *string              `json:"pattern,omitempty"`
	PatternProperties    map[string]*Property `json:"patternProperties,omitempty"`
	Properties           map[string]*Property `json:"properties,omitempty"`
	Ref                  *Reference           `json:"$ref,omitempty"`
	Required             []string             `json:"required,omitempty"`
	Type                 *Type                `json:"type,omitempty"`
	UniqueItems          *bool                `json:"uniqueItems,omitempty"`
}

// String returns a string representation of Property.
func (p *Property) String() string {
	if p == nil {
		return ""
	}

	b, _ := json.MarshalIndent(p, "", "  ")

	return string(b)
}

func (p *Property) IsRequired(name string) bool {
	if p == nil {
		return false
	}

	for _, req := range p.Required {
		if req == name {
			return true
		}
	}

	return false
}
<|endoftext|>


<commit_before>// Copyright 2015 Michal Witkowski. All Rights Reserved.
// See LICENSE for licensing terms.

// Package watcher provides an etcd-backed Watcher for syncing FlagSet state with etcd.

package watcher

import (
	"fmt"
	"math/rand"
	"strings"
	"time"

	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
	etcd "github.com/coreos/etcd/client"
	"github.com/spf13/pflag"
	"github.com/mwitkow/go-flagz"
)

var(
	errNoValue = fmt.Errorf("no value in Node")
	errFlagNotDynamic = fmt.Errorf("flag is not dynamic")
)

// Watcher syncs updates from etcd into a given FlagSet.
type Watcher struct {
	client    etcd.Client
	etcdKeys  etcd.KeysAPI
	flagSet   *pflag.FlagSet
	logger    loggerCompatible
	etcdPath  string
	lastIndex uint64
	watching  bool
	context   context.Context
	cancel    context.CancelFunc
}

// Minimum logger interface needed.
// Default "log" and "logrus" should support these.
type loggerCompatible interface {
	Printf(format string, v ...interface{})
}

// New constructs a new Watcher
func New(set *pflag.FlagSet, keysApi etcd.KeysAPI, etcdPath string, logger loggerCompatible) (*Watcher, error) {
	set.
	if !strings.HasSuffix(etcdPath, "/") {
		etcdPath = etcdPath + "/"
	}
	u := &Watcher{
		flagSet:   set,
		etcdKeys:  keysApi,
		etcdPath:  etcdPath,
		logger:    logger,
		lastIndex: 0,
		watching:  false,
	}
	u.context, u.cancel = context.WithCancel(context.Background())
	return u, nil
}

// Initialize performs the initial read of etcd and sets all flags (dynamic and static) into FlagSet.
func (u *Watcher) Initialize() error {
	if u.lastIndex != 0 {
		return fmt.Errorf("flagz: already initialized.")
	}
	return u.readAllFlags(/* onlyDynamic */ false)
}

// Start kicks off the go routine that syncs dynamic flags from etcd to FlagSet.
func (u *Watcher) Start() error {
	if u.lastIndex == 0 {
		return fmt.Errorf("flagz: not initialized")
	}
	if u.watching {
		return fmt.Errorf("flagz: already watching")
	}
	u.watching = true
	go u.watchForUpdates()
	return nil
}

// Stops the auto-updating go-routine.
func (u *Watcher) Stop() error {
	if !u.watching {
		return fmt.Errorf("flagz: not watching")
	}
	u.logger.Printf("flagz: stopping")
	u.cancel()
	return nil
}

func (u *Watcher) readAllFlags(onlyDynamic bool) error {
	resp, err := u.etcdKeys.Get(u.context, u.etcdPath, &etcd.GetOptions{Recursive: true, Sort: true})
	if err != nil {
		return err
	}
	u.lastIndex = resp.Index
	errorStrings := []string{}
	for _, node := range resp.Node.Nodes {
		flagName, err := u.nodeToFlagName(node)
		if err != nil {
			u.logger.Printf("flagz: ignoring: %v", err)
			continue
		}
		if err := u.setFlag(flagName, node.Value, onlyDynamic); err != nil && err != errNoValue {
			errorStrings = append(errorStrings, err.Error())
		}
	}
	if len(errorStrings) > 0 {
		return fmt.Errorf("flagz: encountered %d errors while parsing flags from etcd: \n  %v",
			len(errorStrings), strings.Join(errorStrings, "\n"))
	}
	return nil
}

func (u *Watcher) setFlag(flagName string, value string, onlyDynamic bool) error {
	if value == "" {
		return errNoValue
	}
	flag := u.flagSet.Lookup(flagName)
	if flag == nil {
		return fmt.Errorf("flag=%v was not found", flagName)
	}
	if onlyDynamic && !flagz.IsFlagDynamic(flag) {
		return errFlagNotDynamic
	}
	// do not call flag.Value.Set, instead go through flagSet.Set to change "changed" state.
	return u.flagSet.Set(flagName, value)
}

func (u *Watcher) watchForUpdates() error {
	// We need to implement our own watcher because the one in go-etcd doesn't handle errorcode 400 and 401.
	// See https://github.com/coreos/etcd/blob/master/Documentation/errorcode.md
	// And https://coreos.com/etcd/docs/2.0.8/api.html#waiting-for-a-change
	watcher := u.etcdKeys.Watcher(u.etcdPath, &etcd.WatcherOptions{AfterIndex: u.lastIndex, Recursive: true})
	u.logger.Printf("flagz: watcher started")
	for u.watching {
		resp, err := watcher.Next(u.context)
		if etcdErr, ok := err.(etcd.Error); ok && etcdErr.Code == etcd.ErrorCodeEventIndexCleared {
			// Our index is out of the Etcd Log. Reread everything and reset index.
			u.logger.Printf("flagz: handling Etcd Index error by re-reading everything: %v", err)
			time.Sleep(200 * time.Millisecond)
			u.readAllFlags(/* onlyDynamic */ true)
			watcher = u.etcdKeys.Watcher(u.etcdPath, &etcd.WatcherOptions{AfterIndex: u.lastIndex, Recursive: true})
			continue
		} else if clusterErr, ok := err.(*etcd.ClusterError); ok {
			u.logger.Printf("flagz: etcd ClusterError. Will retry. %v", clusterErr.Detail())
			time.Sleep(100 * time.Millisecond)
			continue
		} else if err == context.DeadlineExceeded {
			u.logger.Printf("flagz: deadline exceeded which watching for changes, continuing watching")
			continue
		} else if err == context.Canceled {
			break
		} else if err != nil {
			u.logger.Printf("flagz: wicked etcd error. Restarting watching after some time. %v", err)
			// Etcd started dropping watchers, or is re-electing. Give it some time.
			randOffsetMs := int(500 * rand.Float32())
			time.Sleep(1*time.Second + time.Duration(randOffsetMs)*time.Millisecond)
			continue
		}
		u.lastIndex = resp.Node.ModifiedIndex
		flagName, err := u.nodeToFlagName(resp.Node)
		if err != nil {
			u.logger.Printf("flagz: ignoring %v at etcdindex=%v", err, u.lastIndex)
			continue
		}
		err = u.setFlag(flagName, resp.Node.Value, /*onlyDynamic*/ true)
		if err == errNoValue {
			u.logger.Printf("flagz: ignoring action=%v on flag=%v at etcdindex=%v", resp.Action, flagName, u.lastIndex)
			continue
		} else if err == errFlagNotDynamic {
			u.logger.Printf("flagz: ignoring updating flag=%v at etcdindex=%v, because of: %v", flagName, u.lastIndex, err)
		} else if err != nil {
			u.logger.Printf("flagz: failed updating flag=%v at etcdindex=%v, because of: %v", flagName, u.lastIndex, err)
			u.rollbackEtcdValue(flagName, resp)
		} else {
			u.logger.Printf("flagz: updated flag=%v to value=%v at etcdindex=%v", flagName, resp.Node.Value, u.lastIndex)
		}
	}
	u.logger.Printf("flagz: watcher exited")
	return nil
}

func (u *Watcher) rollbackEtcdValue(flagName string, resp *etcd.Response) {
	var err error
	if resp.PrevNode != nil {
		// It's just a new value that's wrong, roll back to prevNode value atomically.
		_, err = u.etcdKeys.Set(u.context, resp.Node.Key, resp.PrevNode.Value, &etcd.SetOptions{PrevIndex: u.lastIndex})
	} else {
		_, err = u.etcdKeys.Delete(u.context, resp.Node.Key, &etcd.DeleteOptions{PrevIndex: u.lastIndex})
	}
	if etcdErr, ok := err.(etcd.Error); ok && etcdErr.Code == etcd.ErrorCodeTestFailed {
		// Someone probably rolled it back in the mean time.
		u.logger.Printf("flagz: rolled back flag=%v was changed by someone else. All good.", flagName)
	} else if err != nil {
		u.logger.Printf("flagz: rolling back flagz=%v failed: %v", flagName, err)
	} else {
		u.logger.Printf("flagz: rolled back flagz=%v to correct state. All good.", flagName)
	}
}

func (u *Watcher) nodeToFlagName(node *etcd.Node) (string, error) {
	if node.Dir {
		return "", fmt.Errorf("key '%v' is a directory entry", node.Key)
	}
	if !strings.HasPrefix(node.Key, u.etcdPath) {
		return "", fmt.Errorf("key '%v' doesn't start with etcd path '%v'", node.Key, u.etcdPath)
	}
	truncated := strings.TrimPrefix(node.Key, u.etcdPath)
	if strings.Count(truncated, "/") > 0 {
		return "", fmt.Errorf("key '%v' isn't a direct leaf of etcd path '%v'", node.Key, u.etcdPath)
	}
	return truncated, nil
}

<commit_msg>fix typo in watcher<commit_after>// Copyright 2015 Michal Witkowski. All Rights Reserved.
// See LICENSE for licensing terms.

// Package watcher provides an etcd-backed Watcher for syncing FlagSet state with etcd.

package watcher

import (
	"fmt"
	"math/rand"
	"strings"
	"time"

	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
	etcd "github.com/coreos/etcd/client"
	"github.com/spf13/pflag"
	"github.com/mwitkow/go-flagz"
)

var(
	errNoValue = fmt.Errorf("no value in Node")
	errFlagNotDynamic = fmt.Errorf("flag is not dynamic")
)

// Watcher syncs updates from etcd into a given FlagSet.
type Watcher struct {
	client    etcd.Client
	etcdKeys  etcd.KeysAPI
	flagSet   *pflag.FlagSet
	logger    loggerCompatible
	etcdPath  string
	lastIndex uint64
	watching  bool
	context   context.Context
	cancel    context.CancelFunc
}

// Minimum logger interface needed.
// Default "log" and "logrus" should support these.
type loggerCompatible interface {
	Printf(format string, v ...interface{})
}

// New constructs a new Watcher
func New(set *pflag.FlagSet, keysApi etcd.KeysAPI, etcdPath string, logger loggerCompatible) (*Watcher, error) {
	if !strings.HasSuffix(etcdPath, "/") {
		etcdPath = etcdPath + "/"
	}
	u := &Watcher{
		flagSet:   set,
		etcdKeys:  keysApi,
		etcdPath:  etcdPath,
		logger:    logger,
		lastIndex: 0,
		watching:  false,
	}
	u.context, u.cancel = context.WithCancel(context.Background())
	return u, nil
}

// Initialize performs the initial read of etcd and sets all flags (dynamic and static) into FlagSet.
func (u *Watcher) Initialize() error {
	if u.lastIndex != 0 {
		return fmt.Errorf("flagz: already initialized.")
	}
	return u.readAllFlags(/* onlyDynamic */ false)
}

// Start kicks off the go routine that syncs dynamic flags from etcd to FlagSet.
func (u *Watcher) Start() error {
	if u.lastIndex == 0 {
		return fmt.Errorf("flagz: not initialized")
	}
	if u.watching {
		return fmt.Errorf("flagz: already watching")
	}
	u.watching = true
	go u.watchForUpdates()
	return nil
}

// Stops the auto-updating go-routine.
func (u *Watcher) Stop() error {
	if !u.watching {
		return fmt.Errorf("flagz: not watching")
	}
	u.logger.Printf("flagz: stopping")
	u.cancel()
	return nil
}

func (u *Watcher) readAllFlags(onlyDynamic bool) error {
	resp, err := u.etcdKeys.Get(u.context, u.etcdPath, &etcd.GetOptions{Recursive: true, Sort: true})
	if err != nil {
		return err
	}
	u.lastIndex = resp.Index
	errorStrings := []string{}
	for _, node := range resp.Node.Nodes {
		flagName, err := u.nodeToFlagName(node)
		if err != nil {
			u.logger.Printf("flagz: ignoring: %v", err)
			continue
		}
		if err := u.setFlag(flagName, node.Value, onlyDynamic); err != nil && err != errNoValue {
			errorStrings = append(errorStrings, err.Error())
		}
	}
	if len(errorStrings) > 0 {
		return fmt.Errorf("flagz: encountered %d errors while parsing flags from etcd: \n  %v",
			len(errorStrings), strings.Join(errorStrings, "\n"))
	}
	return nil
}

func (u *Watcher) setFlag(flagName string, value string, onlyDynamic bool) error {
	if value == "" {
		return errNoValue
	}
	flag := u.flagSet.Lookup(flagName)
	if flag == nil {
		return fmt.Errorf("flag=%v was not found", flagName)
	}
	if onlyDynamic && !flagz.IsFlagDynamic(flag) {
		return errFlagNotDynamic
	}
	// do not call flag.Value.Set, instead go through flagSet.Set to change "changed" state.
	return u.flagSet.Set(flagName, value)
}

func (u *Watcher) watchForUpdates() error {
	// We need to implement our own watcher because the one in go-etcd doesn't handle errorcode 400 and 401.
	// See https://github.com/coreos/etcd/blob/master/Documentation/errorcode.md
	// And https://coreos.com/etcd/docs/2.0.8/api.html#waiting-for-a-change
	watcher := u.etcdKeys.Watcher(u.etcdPath, &etcd.WatcherOptions{AfterIndex: u.lastIndex, Recursive: true})
	u.logger.Printf("flagz: watcher started")
	for u.watching {
		resp, err := watcher.Next(u.context)
		if etcdErr, ok := err.(etcd.Error); ok && etcdErr.Code == etcd.ErrorCodeEventIndexCleared {
			// Our index is out of the Etcd Log. Reread everything and reset index.
			u.logger.Printf("flagz: handling Etcd Index error by re-reading everything: %v", err)
			time.Sleep(200 * time.Millisecond)
			u.readAllFlags(/* onlyDynamic */ true)
			watcher = u.etcdKeys.Watcher(u.etcdPath, &etcd.WatcherOptions{AfterIndex: u.lastIndex, Recursive: true})
			continue
		} else if clusterErr, ok := err.(*etcd.ClusterError); ok {
			u.logger.Printf("flagz: etcd ClusterError. Will retry. %v", clusterErr.Detail())
			time.Sleep(100 * time.Millisecond)
			continue
		} else if err == context.DeadlineExceeded {
			u.logger.Printf("flagz: deadline exceeded which watching for changes, continuing watching")
			continue
		} else if err == context.Canceled {
			break
		} else if err != nil {
			u.logger.Printf("flagz: wicked etcd error. Restarting watching after some time. %v", err)
			// Etcd started dropping watchers, or is re-electing. Give it some time.
			randOffsetMs := int(500 * rand.Float32())
			time.Sleep(1*time.Second + time.Duration(randOffsetMs)*time.Millisecond)
			continue
		}
		u.lastIndex = resp.Node.ModifiedIndex
		flagName, err := u.nodeToFlagName(resp.Node)
		if err != nil {
			u.logger.Printf("flagz: ignoring %v at etcdindex=%v", err, u.lastIndex)
			continue
		}
		err = u.setFlag(flagName, resp.Node.Value, /*onlyDynamic*/ true)
		if err == errNoValue {
			u.logger.Printf("flagz: ignoring action=%v on flag=%v at etcdindex=%v", resp.Action, flagName, u.lastIndex)
			continue
		} else if err == errFlagNotDynamic {
			u.logger.Printf("flagz: ignoring updating flag=%v at etcdindex=%v, because of: %v", flagName, u.lastIndex, err)
		} else if err != nil {
			u.logger.Printf("flagz: failed updating flag=%v at etcdindex=%v, because of: %v", flagName, u.lastIndex, err)
			u.rollbackEtcdValue(flagName, resp)
		} else {
			u.logger.Printf("flagz: updated flag=%v to value=%v at etcdindex=%v", flagName, resp.Node.Value, u.lastIndex)
		}
	}
	u.logger.Printf("flagz: watcher exited")
	return nil
}

func (u *Watcher) rollbackEtcdValue(flagName string, resp *etcd.Response) {
	var err error
	if resp.PrevNode != nil {
		// It's just a new value that's wrong, roll back to prevNode value atomically.
		_, err = u.etcdKeys.Set(u.context, resp.Node.Key, resp.PrevNode.Value, &etcd.SetOptions{PrevIndex: u.lastIndex})
	} else {
		_, err = u.etcdKeys.Delete(u.context, resp.Node.Key, &etcd.DeleteOptions{PrevIndex: u.lastIndex})
	}
	if etcdErr, ok := err.(etcd.Error); ok && etcdErr.Code == etcd.ErrorCodeTestFailed {
		// Someone probably rolled it back in the mean time.
		u.logger.Printf("flagz: rolled back flag=%v was changed by someone else. All good.", flagName)
	} else if err != nil {
		u.logger.Printf("flagz: rolling back flagz=%v failed: %v", flagName, err)
	} else {
		u.logger.Printf("flagz: rolled back flagz=%v to correct state. All good.", flagName)
	}
}

func (u *Watcher) nodeToFlagName(node *etcd.Node) (string, error) {
	if node.Dir {
		return "", fmt.Errorf("key '%v' is a directory entry", node.Key)
	}
	if !strings.HasPrefix(node.Key, u.etcdPath) {
		return "", fmt.Errorf("key '%v' doesn't start with etcd path '%v'", node.Key, u.etcdPath)
	}
	truncated := strings.TrimPrefix(node.Key, u.etcdPath)
	if strings.Count(truncated, "/") > 0 {
		return "", fmt.Errorf("key '%v' isn't a direct leaf of etcd path '%v'", node.Key, u.etcdPath)
	}
	return truncated, nil
}

<|endoftext|>