<commit_before>// Copyright 2011 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"errors"
	"os"
)

func query(filename, query string, bufSize int) (res []string, err error) {
	file, err := os.OpenFile(filename, os.O_RDWR, 0)
	if err != nil {
		return
	}
	defer file.Close()

	_, err = file.WriteString(query)
	if err != nil {
		return
	}
	_, err = file.Seek(0, 0)
	if err != nil {
		return
	}
	buf := make([]byte, bufSize)
	for {
		n, _ := file.Read(buf)
		if n <= 0 {
			break
		}
		res = append(res, string(buf[:n]))
	}
	return
}

func queryCS(net, host, service string) (res []string, err error) {
	switch net {
	case "tcp4", "tcp6":
		net = "tcp"
	case "udp4", "udp6":
		net = "udp"
	}
	if host == "" {
		host = "*"
	}
	return query("/net/cs", net+"!"+host+"!"+service, 128)
}

func queryCS1(net string, ip IP, port int) (clone, dest string, err error) {
	ips := "*"
	if len(ip) != 0 && !ip.IsUnspecified() {
		ips = ip.String()
	}
	lines, err := queryCS(net, ips, itoa(port))
	if err != nil {
		return
	}
	f := getFields(lines[0])
	if len(f) < 2 {
		return "", "", errors.New("net: bad response from ndb/cs")
	}
	clone, dest = f[0], f[1]
	return
}

func queryDNS(addr string, typ string) (res []string, err error) {
	return query("/net/dns", addr+" "+typ, 1024)
}

// toLower returns a lower-case version of in. Restricting us to
// ASCII is sufficient to handle the IP protocol names and allow
// us to not depend on the strings and unicode packages.
func toLower(in string) string {
	for _, c := range in {
		if 'A' <= c && c <= 'Z' {
			// Has upper case; need to fix.
			out := []byte(in)
			for i := 0; i < len(in); i++ {
				c := in[i]
				if 'A' <= c && c <= 'Z' {
					c += 'a' - 'A'
				}
				out[i] = c
			}
			return string(out)
		}
	}
	return in
}

// lookupProtocol looks up IP protocol name and returns
// the corresponding protocol number.
func lookupProtocol(name string) (proto int, err error) {
	lines, err := query("/net/cs", "!protocol="+toLower(name), 128)
	if err != nil {
		return 0, err
	}
	unknownProtoError := errors.New("unknown IP protocol specified: " + name)
	if len(lines) == 0 {
		return 0, unknownProtoError
	}
	f := getFields(lines[0])
	if len(f) < 2 {
		return 0, unknownProtoError
	}
	s := f[1]
	if n, _, ok := dtoi(s, byteIndex(s, '=')+1); ok {
		return n, nil
	}
	return 0, unknownProtoError
}

func lookupHost(host string) (addrs []string, err error) {
	// Use /net/cs instead of /net/dns because cs knows about
	// host names in local network (e.g. from /lib/ndb/local)
	lines, err := queryCS("net", host, "1")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 2 {
			continue
		}
		addr := f[1]
		if i := byteIndex(addr, '!'); i >= 0 {
			addr = addr[:i] // remove port
		}
		if ParseIP(addr) == nil {
			continue
		}
		addrs = append(addrs, addr)
	}
	return
}

func lookupIP(host string) (ips []IP, err error) {
	addrs, err := LookupHost(host)
	if err != nil {
		return
	}
	for _, addr := range addrs {
		if ip := ParseIP(addr); ip != nil {
			ips = append(ips, ip)
		}
	}
	return
}

func lookupPort(network, service string) (port int, err error) {
	switch network {
	case "tcp4", "tcp6":
		network = "tcp"
	case "udp4", "udp6":
		network = "udp"
	}
	lines, err := queryCS(network, "127.0.0.1", service)
	if err != nil {
		return
	}
	unknownPortError := &AddrError{"unknown port", network + "/" + service}
	if len(lines) == 0 {
		return 0, unknownPortError
	}
	f := getFields(lines[0])
	if len(f) < 2 {
		return 0, unknownPortError
	}
	s := f[1]
	if i := byteIndex(s, '!'); i >= 0 {
		s = s[i+1:] // remove address
	}
	if n, _, ok := dtoi(s, 0); ok {
		return n, nil
	}
	return 0, unknownPortError
}

func lookupCNAME(name string) (cname string, err error) {
	lines, err := queryDNS(name, "cname")
	if err != nil {
		return
	}
	if len(lines) > 0 {
		if f := getFields(lines[0]); len(f) >= 3 {
			return f[2] + ".", nil
		}
	}
	return "", errors.New("net: bad response from ndb/dns")
}

func lookupSRV(service, proto, name string) (cname string, addrs []*SRV, err error) {
	var target string
	if service == "" && proto == "" {
		target = name
	} else {
		target = "_" + service + "._" + proto + "." + name
	}
	lines, err := queryDNS(target, "srv")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 6 {
			continue
		}
		port, _, portOk := dtoi(f[4], 0)
		priority, _, priorityOk := dtoi(f[3], 0)
		weight, _, weightOk := dtoi(f[2], 0)
		if !(portOk && priorityOk && weightOk) {
			continue
		}
		addrs = append(addrs, &SRV{f[5], uint16(port), uint16(priority), uint16(weight)})
		cname = f[0]
	}
	byPriorityWeight(addrs).sort()
	return
}

func lookupMX(name string) (mx []*MX, err error) {
	lines, err := queryDNS(name, "mx")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 4 {
			continue
		}
		if pref, _, ok := dtoi(f[2], 0); ok {
			mx = append(mx, &MX{f[3], uint16(pref)})
		}
	}
	byPref(mx).sort()
	return
}

func lookupNS(name string) (ns []*NS, err error) {
	lines, err := queryDNS(name, "ns")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 3 {
			continue
		}
		ns = append(ns, &NS{f[2]})
	}
	return
}

func lookupTXT(name string) (txt []string, err error) {
	lines, err := queryDNS(name, "txt")
	if err != nil {
		return
	}
	for _, line := range lines {
		if i := byteIndex(line, '\t'); i >= 0 {
			txt = append(txt, line[i+1:])
		}
	}
	return
}

func lookupAddr(addr string) (name []string, err error) {
	arpa, err := reverseaddr(addr)
	if err != nil {
		return
	}
	lines, err := queryDNS(arpa, "ptr")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 3 {
			continue
		}
		name = append(name, f[2])
	}
	return
}
<commit_msg>plan9: lookup query must seek to offset 0 before reading or        writing /net/dns or /net/cs (see nbd(8)).<commit_after>// Copyright 2011 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"errors"
	"os"
)

func query(filename, query string, bufSize int) (res []string, err error) {
	file, err := os.OpenFile(filename, os.O_RDWR, 0)
	if err != nil {
		return
	}
	defer file.Close()

	_, err = file.Seek(0, 0)
	if err != nil {
		return
	}
	_, err = file.WriteString(query)
	if err != nil {
		return
	}
	_, err = file.Seek(0, 0)
	if err != nil {
		return
	}
	buf := make([]byte, bufSize)
	for {
		n, _ := file.Read(buf)
		if n <= 0 {
			break
		}
		res = append(res, string(buf[:n]))
	}
	return
}

func queryCS(net, host, service string) (res []string, err error) {
	switch net {
	case "tcp4", "tcp6":
		net = "tcp"
	case "udp4", "udp6":
		net = "udp"
	}
	if host == "" {
		host = "*"
	}
	return query("/net/cs", net+"!"+host+"!"+service, 128)
}

func queryCS1(net string, ip IP, port int) (clone, dest string, err error) {
	ips := "*"
	if len(ip) != 0 && !ip.IsUnspecified() {
		ips = ip.String()
	}
	lines, err := queryCS(net, ips, itoa(port))
	if err != nil {
		return
	}
	f := getFields(lines[0])
	if len(f) < 2 {
		return "", "", errors.New("net: bad response from ndb/cs")
	}
	clone, dest = f[0], f[1]
	return
}

func queryDNS(addr string, typ string) (res []string, err error) {
	return query("/net/dns", addr+" "+typ, 1024)
}

// toLower returns a lower-case version of in. Restricting us to
// ASCII is sufficient to handle the IP protocol names and allow
// us to not depend on the strings and unicode packages.
func toLower(in string) string {
	for _, c := range in {
		if 'A' <= c && c <= 'Z' {
			// Has upper case; need to fix.
			out := []byte(in)
			for i := 0; i < len(in); i++ {
				c := in[i]
				if 'A' <= c && c <= 'Z' {
					c += 'a' - 'A'
				}
				out[i] = c
			}
			return string(out)
		}
	}
	return in
}

// lookupProtocol looks up IP protocol name and returns
// the corresponding protocol number.
func lookupProtocol(name string) (proto int, err error) {
	lines, err := query("/net/cs", "!protocol="+toLower(name), 128)
	if err != nil {
		return 0, err
	}
	unknownProtoError := errors.New("unknown IP protocol specified: " + name)
	if len(lines) == 0 {
		return 0, unknownProtoError
	}
	f := getFields(lines[0])
	if len(f) < 2 {
		return 0, unknownProtoError
	}
	s := f[1]
	if n, _, ok := dtoi(s, byteIndex(s, '=')+1); ok {
		return n, nil
	}
	return 0, unknownProtoError
}

func lookupHost(host string) (addrs []string, err error) {
	// Use /net/cs instead of /net/dns because cs knows about
	// host names in local network (e.g. from /lib/ndb/local)
	lines, err := queryCS("net", host, "1")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 2 {
			continue
		}
		addr := f[1]
		if i := byteIndex(addr, '!'); i >= 0 {
			addr = addr[:i] // remove port
		}
		if ParseIP(addr) == nil {
			continue
		}
		addrs = append(addrs, addr)
	}
	return
}

func lookupIP(host string) (ips []IP, err error) {
	addrs, err := LookupHost(host)
	if err != nil {
		return
	}
	for _, addr := range addrs {
		if ip := ParseIP(addr); ip != nil {
			ips = append(ips, ip)
		}
	}
	return
}

func lookupPort(network, service string) (port int, err error) {
	switch network {
	case "tcp4", "tcp6":
		network = "tcp"
	case "udp4", "udp6":
		network = "udp"
	}
	lines, err := queryCS(network, "127.0.0.1", service)
	if err != nil {
		return
	}
	unknownPortError := &AddrError{"unknown port", network + "/" + service}
	if len(lines) == 0 {
		return 0, unknownPortError
	}
	f := getFields(lines[0])
	if len(f) < 2 {
		return 0, unknownPortError
	}
	s := f[1]
	if i := byteIndex(s, '!'); i >= 0 {
		s = s[i+1:] // remove address
	}
	if n, _, ok := dtoi(s, 0); ok {
		return n, nil
	}
	return 0, unknownPortError
}

func lookupCNAME(name string) (cname string, err error) {
	lines, err := queryDNS(name, "cname")
	if err != nil {
		return
	}
	if len(lines) > 0 {
		if f := getFields(lines[0]); len(f) >= 3 {
			return f[2] + ".", nil
		}
	}
	return "", errors.New("net: bad response from ndb/dns")
}

func lookupSRV(service, proto, name string) (cname string, addrs []*SRV, err error) {
	var target string
	if service == "" && proto == "" {
		target = name
	} else {
		target = "_" + service + "._" + proto + "." + name
	}
	lines, err := queryDNS(target, "srv")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 6 {
			continue
		}
		port, _, portOk := dtoi(f[4], 0)
		priority, _, priorityOk := dtoi(f[3], 0)
		weight, _, weightOk := dtoi(f[2], 0)
		if !(portOk && priorityOk && weightOk) {
			continue
		}
		addrs = append(addrs, &SRV{f[5], uint16(port), uint16(priority), uint16(weight)})
		cname = f[0]
	}
	byPriorityWeight(addrs).sort()
	return
}

func lookupMX(name string) (mx []*MX, err error) {
	lines, err := queryDNS(name, "mx")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 4 {
			continue
		}
		if pref, _, ok := dtoi(f[2], 0); ok {
			mx = append(mx, &MX{f[3], uint16(pref)})
		}
	}
	byPref(mx).sort()
	return
}

func lookupNS(name string) (ns []*NS, err error) {
	lines, err := queryDNS(name, "ns")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 3 {
			continue
		}
		ns = append(ns, &NS{f[2]})
	}
	return
}

func lookupTXT(name string) (txt []string, err error) {
	lines, err := queryDNS(name, "txt")
	if err != nil {
		return
	}
	for _, line := range lines {
		if i := byteIndex(line, '\t'); i >= 0 {
			txt = append(txt, line[i+1:])
		}
	}
	return
}

func lookupAddr(addr string) (name []string, err error) {
	arpa, err := reverseaddr(addr)
	if err != nil {
		return
	}
	lines, err := queryDNS(arpa, "ptr")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 3 {
			continue
		}
		name = append(name, f[2])
	}
	return
}
<|endoftext|>


<commit_before>package remind

import (
	"github.com/nmeum/marvin/irc"
	"github.com/nmeum/marvin/modules"
	"strings"
	"time"
)

type Module struct {
	TimeLimit int `json:"limit"`
}

func Init(moduleSet *modules.ModuleSet) {
	moduleSet.Register(new(Module))
}

func (m *Module) Name() string {
	return "remind"
}

func (m *Module) Help() string {
	return "!remind DURATION MSG"
}

func (m *Module) Defaults() {
	m.TimeLimit = 10
}

func (m *Module) Load(client *irc.Client) error {
	client.CmdHook("privmsg", func(c *irc.Client, msg irc.Message) error {
		splited := strings.Split(msg.Data, " ")
		if len(splited) < 3 || splited[0] != "!remind" {
			return nil
		}

		duration, err := time.ParseDuration(splited[1])
		if err != nil {
			return err
		}

		limit := (time.Duration)(m.TimeLimit) * time.Hour
		if duration > limit {
			return c.Write("NOTICE %s :%d exceeds the limit of %d",
				duration.Hours(), limit.Hours())
		}

		time.AfterFunc(duration, func() {
			c.Write("PRIVMSG %s :Reminder: %s",
				msg.Sender.Name, splited[2])
		})

		return c.Write("NOTICE %s :Reminder setup for %s",
			msg.Receiver, splited[1])
	})

	return nil
}
<commit_msg>Improve output of remind module<commit_after>package remind

import (
	"github.com/nmeum/marvin/irc"
	"github.com/nmeum/marvin/modules"
	"strings"
	"time"
)

type Module struct {
	TimeLimit int `json:"limit"`
}

func Init(moduleSet *modules.ModuleSet) {
	moduleSet.Register(new(Module))
}

func (m *Module) Name() string {
	return "remind"
}

func (m *Module) Help() string {
	return "!remind DURATION MSG"
}

func (m *Module) Defaults() {
	m.TimeLimit = 10
}

func (m *Module) Load(client *irc.Client) error {
	client.CmdHook("privmsg", func(c *irc.Client, msg irc.Message) error {
		splited := strings.Split(msg.Data, " ")
		if len(splited) < 3 || splited[0] != "!remind" {
			return nil
		}

		rawDuration := splited[1]
		duration, err := time.ParseDuration(rawDuration)
		if err != nil {
			return err
		}

		limit := (time.Duration)(m.TimeLimit) * time.Hour
		if duration > limit {
			return c.Write("NOTICE %s :%d exceeds the limit of %d",
				duration.Hours(), limit.Hours())
		}

		reminder := strings.Join(splited[2:], " ")
		time.AfterFunc(duration, func() {
			c.Write("PRIVMSG %s :Reminder: %s",
				msg.Sender.Name, reminder)
		})

		return c.Write("NOTICE %s :Reminder setup for %s",
			msg.Receiver, rawDuration)
	})

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"errors"
	"os"
)

func query(filename, query string, bufSize int) (res []string, err error) {
	file, err := os.OpenFile(filename, os.O_RDWR, 0)
	if err != nil {
		return
	}
	defer file.Close()

	_, err = file.WriteString(query)
	if err != nil {
		return
	}
	_, err = file.Seek(0, 0)
	if err != nil {
		return
	}
	buf := make([]byte, bufSize)
	for {
		n, _ := file.Read(buf)
		if n <= 0 {
			break
		}
		res = append(res, string(buf[:n]))
	}
	return
}

func queryCS(net, host, service string) (res []string, err error) {
	switch net {
	case "tcp4", "tcp6":
		net = "tcp"
	case "udp4", "udp6":
		net = "udp"
	}
	if host == "" {
		host = "*"
	}
	return query("/net/cs", net+"!"+host+"!"+service, 128)
}

func queryCS1(net string, ip IP, port int) (clone, dest string, err error) {
	ips := "*"
	if len(ip) != 0 && !ip.IsUnspecified() {
		ips = ip.String()
	}
	lines, err := queryCS(net, ips, itoa(port))
	if err != nil {
		return
	}
	f := getFields(lines[0])
	if len(f) < 2 {
		return "", "", errors.New("net: bad response from ndb/cs")
	}
	clone, dest = f[0], f[1]
	return
}

func queryDNS(addr string, typ string) (res []string, err error) {
	return query("/net/dns", addr+" "+typ, 1024)
}

// lookupProtocol looks up IP protocol name and returns
// the corresponding protocol number.
func lookupProtocol(name string) (proto int, err error) {
	lines, err := query("/net/cs", "!protocol="+name, 128)
	if err != nil {
		return 0, err
	}
	unknownProtoError := errors.New("unknown IP protocol specified: " + name)
	if len(lines) == 0 {
		return 0, unknownProtoError
	}
	f := getFields(lines[0])
	if len(f) < 2 {
		return 0, unknownProtoError
	}
	s := f[1]
	if n, _, ok := dtoi(s, byteIndex(s, '=')+1); ok {
		return n, nil
	}
	return 0, unknownProtoError
}

func lookupHost(host string) (addrs []string, err error) {
	// Use /net/cs instead of /net/dns because cs knows about
	// host names in local network (e.g. from /lib/ndb/local)
	lines, err := queryCS("tcp", host, "1")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 2 {
			continue
		}
		addr := f[1]
		if i := byteIndex(addr, '!'); i >= 0 {
			addr = addr[:i] // remove port
		}
		if ParseIP(addr) == nil {
			continue
		}
		addrs = append(addrs, addr)
	}
	return
}

func lookupIP(host string) (ips []IP, err error) {
	addrs, err := LookupHost(host)
	if err != nil {
		return
	}
	for _, addr := range addrs {
		if ip := ParseIP(addr); ip != nil {
			ips = append(ips, ip)
		}
	}
	return
}

func lookupPort(network, service string) (port int, err error) {
	switch network {
	case "tcp4", "tcp6":
		network = "tcp"
	case "udp4", "udp6":
		network = "udp"
	}
	lines, err := queryCS(network, "127.0.0.1", service)
	if err != nil {
		return
	}
	unknownPortError := &AddrError{"unknown port", network + "/" + service}
	if len(lines) == 0 {
		return 0, unknownPortError
	}
	f := getFields(lines[0])
	if len(f) < 2 {
		return 0, unknownPortError
	}
	s := f[1]
	if i := byteIndex(s, '!'); i >= 0 {
		s = s[i+1:] // remove address
	}
	if n, _, ok := dtoi(s, 0); ok {
		return n, nil
	}
	return 0, unknownPortError
}

func lookupCNAME(name string) (cname string, err error) {
	lines, err := queryDNS(name, "cname")
	if err != nil {
		return
	}
	if len(lines) > 0 {
		if f := getFields(lines[0]); len(f) >= 3 {
			return f[2] + ".", nil
		}
	}
	return "", errors.New("net: bad response from ndb/dns")
}

func lookupSRV(service, proto, name string) (cname string, addrs []*SRV, err error) {
	var target string
	if service == "" && proto == "" {
		target = name
	} else {
		target = "_" + service + "._" + proto + "." + name
	}
	lines, err := queryDNS(target, "srv")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 6 {
			continue
		}
		port, _, portOk := dtoi(f[2], 0)
		priority, _, priorityOk := dtoi(f[3], 0)
		weight, _, weightOk := dtoi(f[4], 0)
		if !(portOk && priorityOk && weightOk) {
			continue
		}
		addrs = append(addrs, &SRV{f[5], uint16(port), uint16(priority), uint16(weight)})
		cname = f[0]
	}
	byPriorityWeight(addrs).sort()
	return
}

func lookupMX(name string) (mx []*MX, err error) {
	lines, err := queryDNS(name, "mx")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 4 {
			continue
		}
		if pref, _, ok := dtoi(f[2], 0); ok {
			mx = append(mx, &MX{f[3], uint16(pref)})
		}
	}
	byPref(mx).sort()
	return
}

func lookupNS(name string) (ns []*NS, err error) {
	lines, err := queryDNS(name, "ns")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 3 {
			continue
		}
		ns = append(ns, &NS{f[2]})
	}
	return
}

func lookupTXT(name string) (txt []string, err error) {
	lines, err := queryDNS(name, "txt")
	if err != nil {
		return
	}
	for _, line := range lines {
		if i := byteIndex(line, '\t'); i >= 0 {
			txt = append(txt, line[i+1:])
		}
	}
	return
}

func lookupAddr(addr string) (name []string, err error) {
	arpa, err := reverseaddr(addr)
	if err != nil {
		return
	}
	lines, err := queryDNS(arpa, "ptr")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 3 {
			continue
		}
		name = append(name, f[2])
	}
	return
}
<commit_msg>net: fix LookupSRV ordering on plan 9<commit_after>// Copyright 2011 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"errors"
	"os"
)

func query(filename, query string, bufSize int) (res []string, err error) {
	file, err := os.OpenFile(filename, os.O_RDWR, 0)
	if err != nil {
		return
	}
	defer file.Close()

	_, err = file.WriteString(query)
	if err != nil {
		return
	}
	_, err = file.Seek(0, 0)
	if err != nil {
		return
	}
	buf := make([]byte, bufSize)
	for {
		n, _ := file.Read(buf)
		if n <= 0 {
			break
		}
		res = append(res, string(buf[:n]))
	}
	return
}

func queryCS(net, host, service string) (res []string, err error) {
	switch net {
	case "tcp4", "tcp6":
		net = "tcp"
	case "udp4", "udp6":
		net = "udp"
	}
	if host == "" {
		host = "*"
	}
	return query("/net/cs", net+"!"+host+"!"+service, 128)
}

func queryCS1(net string, ip IP, port int) (clone, dest string, err error) {
	ips := "*"
	if len(ip) != 0 && !ip.IsUnspecified() {
		ips = ip.String()
	}
	lines, err := queryCS(net, ips, itoa(port))
	if err != nil {
		return
	}
	f := getFields(lines[0])
	if len(f) < 2 {
		return "", "", errors.New("net: bad response from ndb/cs")
	}
	clone, dest = f[0], f[1]
	return
}

func queryDNS(addr string, typ string) (res []string, err error) {
	return query("/net/dns", addr+" "+typ, 1024)
}

// lookupProtocol looks up IP protocol name and returns
// the corresponding protocol number.
func lookupProtocol(name string) (proto int, err error) {
	lines, err := query("/net/cs", "!protocol="+name, 128)
	if err != nil {
		return 0, err
	}
	unknownProtoError := errors.New("unknown IP protocol specified: " + name)
	if len(lines) == 0 {
		return 0, unknownProtoError
	}
	f := getFields(lines[0])
	if len(f) < 2 {
		return 0, unknownProtoError
	}
	s := f[1]
	if n, _, ok := dtoi(s, byteIndex(s, '=')+1); ok {
		return n, nil
	}
	return 0, unknownProtoError
}

func lookupHost(host string) (addrs []string, err error) {
	// Use /net/cs instead of /net/dns because cs knows about
	// host names in local network (e.g. from /lib/ndb/local)
	lines, err := queryCS("tcp", host, "1")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 2 {
			continue
		}
		addr := f[1]
		if i := byteIndex(addr, '!'); i >= 0 {
			addr = addr[:i] // remove port
		}
		if ParseIP(addr) == nil {
			continue
		}
		addrs = append(addrs, addr)
	}
	return
}

func lookupIP(host string) (ips []IP, err error) {
	addrs, err := LookupHost(host)
	if err != nil {
		return
	}
	for _, addr := range addrs {
		if ip := ParseIP(addr); ip != nil {
			ips = append(ips, ip)
		}
	}
	return
}

func lookupPort(network, service string) (port int, err error) {
	switch network {
	case "tcp4", "tcp6":
		network = "tcp"
	case "udp4", "udp6":
		network = "udp"
	}
	lines, err := queryCS(network, "127.0.0.1", service)
	if err != nil {
		return
	}
	unknownPortError := &AddrError{"unknown port", network + "/" + service}
	if len(lines) == 0 {
		return 0, unknownPortError
	}
	f := getFields(lines[0])
	if len(f) < 2 {
		return 0, unknownPortError
	}
	s := f[1]
	if i := byteIndex(s, '!'); i >= 0 {
		s = s[i+1:] // remove address
	}
	if n, _, ok := dtoi(s, 0); ok {
		return n, nil
	}
	return 0, unknownPortError
}

func lookupCNAME(name string) (cname string, err error) {
	lines, err := queryDNS(name, "cname")
	if err != nil {
		return
	}
	if len(lines) > 0 {
		if f := getFields(lines[0]); len(f) >= 3 {
			return f[2] + ".", nil
		}
	}
	return "", errors.New("net: bad response from ndb/dns")
}

func lookupSRV(service, proto, name string) (cname string, addrs []*SRV, err error) {
	var target string
	if service == "" && proto == "" {
		target = name
	} else {
		target = "_" + service + "._" + proto + "." + name
	}
	lines, err := queryDNS(target, "srv")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 6 {
			continue
		}
		port, _, portOk := dtoi(f[4], 0)
		priority, _, priorityOk := dtoi(f[3], 0)
		weight, _, weightOk := dtoi(f[2], 0)
		if !(portOk && priorityOk && weightOk) {
			continue
		}
		addrs = append(addrs, &SRV{f[5], uint16(port), uint16(priority), uint16(weight)})
		cname = f[0]
	}
	byPriorityWeight(addrs).sort()
	return
}

func lookupMX(name string) (mx []*MX, err error) {
	lines, err := queryDNS(name, "mx")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 4 {
			continue
		}
		if pref, _, ok := dtoi(f[2], 0); ok {
			mx = append(mx, &MX{f[3], uint16(pref)})
		}
	}
	byPref(mx).sort()
	return
}

func lookupNS(name string) (ns []*NS, err error) {
	lines, err := queryDNS(name, "ns")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 3 {
			continue
		}
		ns = append(ns, &NS{f[2]})
	}
	return
}

func lookupTXT(name string) (txt []string, err error) {
	lines, err := queryDNS(name, "txt")
	if err != nil {
		return
	}
	for _, line := range lines {
		if i := byteIndex(line, '\t'); i >= 0 {
			txt = append(txt, line[i+1:])
		}
	}
	return
}

func lookupAddr(addr string) (name []string, err error) {
	arpa, err := reverseaddr(addr)
	if err != nil {
		return
	}
	lines, err := queryDNS(arpa, "ptr")
	if err != nil {
		return
	}
	for _, line := range lines {
		f := getFields(line)
		if len(f) < 3 {
			continue
		}
		name = append(name, f[2])
	}
	return
}
<|endoftext|>


<commit_before>package spec

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"time"

	"github.com/Songmu/retry"
	"github.com/mackerelio/golib/logging"
)

// This Generator collects metadata about cloud instances.
// Currently EC2 and GCE are supported.
// EC2: http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/AESDG-chapter-instancedata.html
// GCE: https://developers.google.com/compute/docs/metadata
// AzureVM: https://docs.microsoft.com/azure/virtual-machines/virtual-machines-instancemetadataservice-overview

// CloudGenerator definition
type CloudGenerator struct {
	CloudMetaGenerator
}

// CloudMetaGenerator interface of metadata generator for each cloud platform
type CloudMetaGenerator interface {
	Generate() (interface{}, error)
	SuggestCustomIdentifier() (string, error)
}

// Key is a root key for the generator.
func (g *CloudGenerator) Key() string {
	return "cloud"
}

var cloudLogger = logging.GetLogger("spec.cloud")

var ec2BaseURL, gceMetaURL, azureVMBaseURL *url.URL

func init() {
	ec2BaseURL, _ = url.Parse("http://169.254.169.254/latest/meta-data")
	gceMetaURL, _ = url.Parse("http://metadata.google.internal/computeMetadata/v1/?recursive=true")
	azureVMBaseURL, _ = url.Parse("http://169.254.169.254/metadata/instance")
}

var timeout = 100 * time.Millisecond

// SuggestCloudGenerator returns suitable CloudGenerator
func SuggestCloudGenerator() *CloudGenerator {
	if isEC2() {
		return &CloudGenerator{&EC2Generator{ec2BaseURL}}
	}
	if isGCE() {
		return &CloudGenerator{&GCEGenerator{gceMetaURL}}
	}

	return nil
}

func httpCli() *http.Client {
	return &http.Client{
		Timeout: timeout,
		Transport: &http.Transport{
			// don't use HTTP_PROXY when requesting cloud instance metadata APIs
			Proxy: nil,
		},
	}
}

func isEC2() bool {
	// If the OS is Linux, check /sys/hypervisor/uuid file first.
	// ref. http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/identify_ec2_instances.html
	if runtime.GOOS == "linux" {
		data, err := ioutil.ReadFile("/sys/hypervisor/uuid")
		if err != nil {
			// Probably not EC2.
			return false
		}
		// Probably not EC2.
		if !strings.HasPrefix(string(data), "ec2") {
			return false
		}
	}

	res := false
	cl := httpCli()
	err := retry.Retry(3, 2*timeout, func() error {
		// '/ami-id` is probably an AWS specific URL
		resp, err := cl.Get(ec2BaseURL.String() + "/ami-id")
		if err != nil {
			return err
		}
		defer resp.Body.Close()

		res = resp.StatusCode == 200
		return nil
	})

	if err == nil {
		return res
	}
	return false
}

func isGCE() bool {
	_, err := requestGCEMeta()
	return err == nil
}

// Note: May want to check without using the API.
func isAzure() bool {
	cl := httpCli()
	// '/vmId` is probably Azure VM specific URL
	req, err := http.NewRequest("GET", azureVMBaseURL.String()+"/compute/vmId?api-version=2017-04-02&format=text", nil)
	if err != nil {
		return false
	}
	req.Header.Set("Metadata", "true")

	resp, err := cl.Do(req)
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	return resp.StatusCode == 200
}

func requestGCEMeta() ([]byte, error) {
	cl := httpCli()
	req, err := http.NewRequest("GET", gceMetaURL.String(), nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Metadata-Flavor", "Google")

	resp, err := cl.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("failed to request gce meta. response code: %d", resp.StatusCode)
	}
	return ioutil.ReadAll(resp.Body)
}

// EC2Generator meta generator for EC2
type EC2Generator struct {
	baseURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *EC2Generator) Generate() (interface{}, error) {
	cl := httpCli()

	metadataKeys := []string{
		"instance-id",
		"instance-type",
		"placement/availability-zone",
		"security-groups",
		"ami-id",
		"hostname",
		"local-hostname",
		"public-hostname",
		"local-ipv4",
		"public-ipv4",
		"reservation-id",
	}

	metadata := make(map[string]string)

	for _, key := range metadataKeys {
		resp, err := cl.Get(g.baseURL.String() + "/" + key)
		if err != nil {
			cloudLogger.Debugf("This host may not be running on EC2. Error while reading '%s'", key)
			return nil, nil
		}
		defer resp.Body.Close()
		if resp.StatusCode == 200 {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				cloudLogger.Errorf("Results of requesting metadata cannot be read: '%s'", err)
				break
			}
			metadata[key] = string(body)
			cloudLogger.Debugf("results %s:%s", key, string(body))
		} else {
			cloudLogger.Debugf("Status code of the result of requesting metadata '%s' is '%d'", key, resp.StatusCode)
		}
	}

	results := make(map[string]interface{})
	results["provider"] = "ec2"
	results["metadata"] = metadata

	return results, nil
}

// SuggestCustomIdentifier suggests the identifier of the EC2 instance
func (g *EC2Generator) SuggestCustomIdentifier() (string, error) {
	cl := httpCli()
	key := "instance-id"
	resp, err := cl.Get(g.baseURL.String() + "/" + key)
	if err != nil {
		return "", fmt.Errorf("error while retrieving instance-id")
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("failed to request instance-id. response code: %d", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("results of requesting instance-id cannot be read: '%s'", err)
	}
	instanceID := string(body)
	if instanceID == "" {
		return "", fmt.Errorf("invalid instance id")
	}
	return instanceID + ".ec2.amazonaws.com", nil
}

// GCEGenerator generate for GCE
type GCEGenerator struct {
	metaURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *GCEGenerator) Generate() (interface{}, error) {
	bytes, err := requestGCEMeta()
	if err != nil {
		return nil, err
	}
	var data gceMeta
	json.Unmarshal(bytes, &data)
	return data.toGeneratorResults(), nil
}

type gceInstance struct {
	Zone         string
	InstanceType string `json:"machineType"`
	Hostname     string
	InstanceID   uint64 `json:"id"`
}

type gceProject struct {
	ProjectID        string
	NumericProjectID uint64
}

type gceMeta struct {
	Instance *gceInstance
	Project  *gceProject
}

func (g gceMeta) toGeneratorMeta() map[string]string {
	meta := make(map[string]string)

	lastS := func(s string) string {
		ss := strings.Split(s, "/")
		return ss[len(ss)-1]
	}

	if ins := g.Instance; ins != nil {
		meta["hostname"] = ins.Hostname
		meta["instance-id"] = fmt.Sprint(ins.InstanceID)
		meta["instance-type"] = lastS(ins.InstanceType)
		meta["zone"] = lastS(ins.Zone)
	}

	if proj := g.Project; proj != nil {
		meta["projectId"] = proj.ProjectID
	}

	return meta
}

func (g gceMeta) toGeneratorResults() interface{} {
	results := make(map[string]interface{})
	results["provider"] = "gce"
	results["metadata"] = g.toGeneratorMeta()

	return results
}

// SuggestCustomIdentifier for GCE is not implemented yet
func (g *GCEGenerator) SuggestCustomIdentifier() (string, error) {
	return "", nil
}

// AzureVMGenerator meta generator for Azure VM
type AzureVMGenerator struct {
	baseURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *AzureVMGenerator) Generate() (interface{}, error) {
	metadataComputeKeys := []string{
		"location",
		"name",
		"offer",
		"osType",
		"platformFaultDomain",
		"platformUpdateDomain",
		"publisher",
		"sku",
		"version",
		"vmId",
		"vmSize",
	}

	ipAddressKeys := []string{
		"privateIpAddress",
		"publicIpAddress",
	}

	metadata := make(map[string]string)
	metadata = retrieveAzureVMMetadata(metadata, g.baseURL.String(), "/compute/", metadataComputeKeys)
	metadata = retrieveAzureVMMetadata(metadata, g.baseURL.String(), "/network/interface/0/ipv4/ipAddress/0/", ipAddressKeys)

	results := make(map[string]interface{})
	results["provider"] = "AzureVM"
	results["metadata"] = metadata

	return results, nil
}

func retrieveAzureVMMetadata(metadataMap map[string]string, baseURL string, urlSuffix string, keys []string) map[string]string {
	cl := httpCli()

	for _, key := range keys {
		req, err := http.NewRequest("GET", baseURL+urlSuffix+key+"?api-version=2017-04-02&format=text", nil)
		if err != nil {
			cloudLogger.Debugf("This host may not be running on Azure VM. Error while reading '%s'", key)
			return nil
		}

		req.Header.Set("Metadata", "true")

		resp, err := cl.Do(req)
		if err != nil {
			cloudLogger.Debugf("This host may not be running on Azure VM. Error while reading '%s'", key)
			return nil
		}
		defer resp.Body.Close()

		if resp.StatusCode == 200 {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				cloudLogger.Errorf("Results of requesting metadata cannot be read: '%s'", err)
				break
			}
			metadataMap[key] = string(body)
			cloudLogger.Debugf("results %s:%s", key, string(body))
		} else {
			cloudLogger.Debugf("Status code of the result of requesting metadata '%s' is '%d'", key, resp.StatusCode)
		}
	}
	return metadataMap
}

// SuggestCustomIdentifier suggests the identifier of the Azure VM instance
func (g *AzureVMGenerator) SuggestCustomIdentifier() (string, error) {
	cl := httpCli()
	req, err := http.NewRequest("GET", azureVMBaseURL.String()+"/compute/vmId?api-version=2017-04-02&format=text", nil)
	if err != nil {
		return "", fmt.Errorf("error while retrieving vmId")
	}
	req.Header.Set("Metadata", "true")

	resp, err := cl.Do(req)
	if err != nil {
		return "", fmt.Errorf("error while retrieving vmId")
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return "", fmt.Errorf("failed to request vmId. response code: %d", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("results of requesting vmId cannot be read: '%s'", err)
	}
	instanceID := string(body)
	if instanceID == "" {
		return "", fmt.Errorf("invalid instance id")
	}
	return instanceID + ".virtual_machine.azure.microsoft.com", nil
}
<commit_msg>adjust retry interval for isEC2 check<commit_after>package spec

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"time"

	"github.com/Songmu/retry"
	"github.com/mackerelio/golib/logging"
)

// This Generator collects metadata about cloud instances.
// Currently EC2 and GCE are supported.
// EC2: http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/AESDG-chapter-instancedata.html
// GCE: https://developers.google.com/compute/docs/metadata
// AzureVM: https://docs.microsoft.com/azure/virtual-machines/virtual-machines-instancemetadataservice-overview

// CloudGenerator definition
type CloudGenerator struct {
	CloudMetaGenerator
}

// CloudMetaGenerator interface of metadata generator for each cloud platform
type CloudMetaGenerator interface {
	Generate() (interface{}, error)
	SuggestCustomIdentifier() (string, error)
}

// Key is a root key for the generator.
func (g *CloudGenerator) Key() string {
	return "cloud"
}

var cloudLogger = logging.GetLogger("spec.cloud")

var ec2BaseURL, gceMetaURL, azureVMBaseURL *url.URL

func init() {
	ec2BaseURL, _ = url.Parse("http://169.254.169.254/latest/meta-data")
	gceMetaURL, _ = url.Parse("http://metadata.google.internal/computeMetadata/v1/?recursive=true")
	azureVMBaseURL, _ = url.Parse("http://169.254.169.254/metadata/instance")
}

var timeout = 100 * time.Millisecond

// SuggestCloudGenerator returns suitable CloudGenerator
func SuggestCloudGenerator() *CloudGenerator {
	if isEC2() {
		return &CloudGenerator{&EC2Generator{ec2BaseURL}}
	}
	if isGCE() {
		return &CloudGenerator{&GCEGenerator{gceMetaURL}}
	}

	return nil
}

func httpCli() *http.Client {
	return &http.Client{
		Timeout: timeout,
		Transport: &http.Transport{
			// don't use HTTP_PROXY when requesting cloud instance metadata APIs
			Proxy: nil,
		},
	}
}

func isEC2() bool {
	// If the OS is Linux, check /sys/hypervisor/uuid file first.
	// ref. http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/identify_ec2_instances.html
	if runtime.GOOS == "linux" {
		data, err := ioutil.ReadFile("/sys/hypervisor/uuid")
		if err != nil {
			// Probably not EC2.
			return false
		}
		// Probably not EC2.
		if !strings.HasPrefix(string(data), "ec2") {
			return false
		}
	}

	res := false
	cl := httpCli()
	err := retry.Retry(3, 2*time.Second, func() error {
		// '/ami-id` is probably an AWS specific URL
		resp, err := cl.Get(ec2BaseURL.String() + "/ami-id")
		if err != nil {
			return err
		}
		defer resp.Body.Close()

		res = resp.StatusCode == 200
		return nil
	})

	if err == nil {
		return res
	}
	return false
}

func isGCE() bool {
	_, err := requestGCEMeta()
	return err == nil
}

// Note: May want to check without using the API.
func isAzure() bool {
	cl := httpCli()
	// '/vmId` is probably Azure VM specific URL
	req, err := http.NewRequest("GET", azureVMBaseURL.String()+"/compute/vmId?api-version=2017-04-02&format=text", nil)
	if err != nil {
		return false
	}
	req.Header.Set("Metadata", "true")

	resp, err := cl.Do(req)
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	return resp.StatusCode == 200
}

func requestGCEMeta() ([]byte, error) {
	cl := httpCli()
	req, err := http.NewRequest("GET", gceMetaURL.String(), nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Metadata-Flavor", "Google")

	resp, err := cl.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("failed to request gce meta. response code: %d", resp.StatusCode)
	}
	return ioutil.ReadAll(resp.Body)
}

// EC2Generator meta generator for EC2
type EC2Generator struct {
	baseURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *EC2Generator) Generate() (interface{}, error) {
	cl := httpCli()

	metadataKeys := []string{
		"instance-id",
		"instance-type",
		"placement/availability-zone",
		"security-groups",
		"ami-id",
		"hostname",
		"local-hostname",
		"public-hostname",
		"local-ipv4",
		"public-ipv4",
		"reservation-id",
	}

	metadata := make(map[string]string)

	for _, key := range metadataKeys {
		resp, err := cl.Get(g.baseURL.String() + "/" + key)
		if err != nil {
			cloudLogger.Debugf("This host may not be running on EC2. Error while reading '%s'", key)
			return nil, nil
		}
		defer resp.Body.Close()
		if resp.StatusCode == 200 {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				cloudLogger.Errorf("Results of requesting metadata cannot be read: '%s'", err)
				break
			}
			metadata[key] = string(body)
			cloudLogger.Debugf("results %s:%s", key, string(body))
		} else {
			cloudLogger.Debugf("Status code of the result of requesting metadata '%s' is '%d'", key, resp.StatusCode)
		}
	}

	results := make(map[string]interface{})
	results["provider"] = "ec2"
	results["metadata"] = metadata

	return results, nil
}

// SuggestCustomIdentifier suggests the identifier of the EC2 instance
func (g *EC2Generator) SuggestCustomIdentifier() (string, error) {
	cl := httpCli()
	key := "instance-id"
	resp, err := cl.Get(g.baseURL.String() + "/" + key)
	if err != nil {
		return "", fmt.Errorf("error while retrieving instance-id")
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("failed to request instance-id. response code: %d", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("results of requesting instance-id cannot be read: '%s'", err)
	}
	instanceID := string(body)
	if instanceID == "" {
		return "", fmt.Errorf("invalid instance id")
	}
	return instanceID + ".ec2.amazonaws.com", nil
}

// GCEGenerator generate for GCE
type GCEGenerator struct {
	metaURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *GCEGenerator) Generate() (interface{}, error) {
	bytes, err := requestGCEMeta()
	if err != nil {
		return nil, err
	}
	var data gceMeta
	json.Unmarshal(bytes, &data)
	return data.toGeneratorResults(), nil
}

type gceInstance struct {
	Zone         string
	InstanceType string `json:"machineType"`
	Hostname     string
	InstanceID   uint64 `json:"id"`
}

type gceProject struct {
	ProjectID        string
	NumericProjectID uint64
}

type gceMeta struct {
	Instance *gceInstance
	Project  *gceProject
}

func (g gceMeta) toGeneratorMeta() map[string]string {
	meta := make(map[string]string)

	lastS := func(s string) string {
		ss := strings.Split(s, "/")
		return ss[len(ss)-1]
	}

	if ins := g.Instance; ins != nil {
		meta["hostname"] = ins.Hostname
		meta["instance-id"] = fmt.Sprint(ins.InstanceID)
		meta["instance-type"] = lastS(ins.InstanceType)
		meta["zone"] = lastS(ins.Zone)
	}

	if proj := g.Project; proj != nil {
		meta["projectId"] = proj.ProjectID
	}

	return meta
}

func (g gceMeta) toGeneratorResults() interface{} {
	results := make(map[string]interface{})
	results["provider"] = "gce"
	results["metadata"] = g.toGeneratorMeta()

	return results
}

// SuggestCustomIdentifier for GCE is not implemented yet
func (g *GCEGenerator) SuggestCustomIdentifier() (string, error) {
	return "", nil
}

// AzureVMGenerator meta generator for Azure VM
type AzureVMGenerator struct {
	baseURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *AzureVMGenerator) Generate() (interface{}, error) {
	metadataComputeKeys := []string{
		"location",
		"name",
		"offer",
		"osType",
		"platformFaultDomain",
		"platformUpdateDomain",
		"publisher",
		"sku",
		"version",
		"vmId",
		"vmSize",
	}

	ipAddressKeys := []string{
		"privateIpAddress",
		"publicIpAddress",
	}

	metadata := make(map[string]string)
	metadata = retrieveAzureVMMetadata(metadata, g.baseURL.String(), "/compute/", metadataComputeKeys)
	metadata = retrieveAzureVMMetadata(metadata, g.baseURL.String(), "/network/interface/0/ipv4/ipAddress/0/", ipAddressKeys)

	results := make(map[string]interface{})
	results["provider"] = "AzureVM"
	results["metadata"] = metadata

	return results, nil
}

func retrieveAzureVMMetadata(metadataMap map[string]string, baseURL string, urlSuffix string, keys []string) map[string]string {
	cl := httpCli()

	for _, key := range keys {
		req, err := http.NewRequest("GET", baseURL+urlSuffix+key+"?api-version=2017-04-02&format=text", nil)
		if err != nil {
			cloudLogger.Debugf("This host may not be running on Azure VM. Error while reading '%s'", key)
			return nil
		}

		req.Header.Set("Metadata", "true")

		resp, err := cl.Do(req)
		if err != nil {
			cloudLogger.Debugf("This host may not be running on Azure VM. Error while reading '%s'", key)
			return nil
		}
		defer resp.Body.Close()

		if resp.StatusCode == 200 {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				cloudLogger.Errorf("Results of requesting metadata cannot be read: '%s'", err)
				break
			}
			metadataMap[key] = string(body)
			cloudLogger.Debugf("results %s:%s", key, string(body))
		} else {
			cloudLogger.Debugf("Status code of the result of requesting metadata '%s' is '%d'", key, resp.StatusCode)
		}
	}
	return metadataMap
}

// SuggestCustomIdentifier suggests the identifier of the Azure VM instance
func (g *AzureVMGenerator) SuggestCustomIdentifier() (string, error) {
	cl := httpCli()
	req, err := http.NewRequest("GET", azureVMBaseURL.String()+"/compute/vmId?api-version=2017-04-02&format=text", nil)
	if err != nil {
		return "", fmt.Errorf("error while retrieving vmId")
	}
	req.Header.Set("Metadata", "true")

	resp, err := cl.Do(req)
	if err != nil {
		return "", fmt.Errorf("error while retrieving vmId")
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return "", fmt.Errorf("failed to request vmId. response code: %d", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("results of requesting vmId cannot be read: '%s'", err)
	}
	instanceID := string(body)
	if instanceID == "" {
		return "", fmt.Errorf("invalid instance id")
	}
	return instanceID + ".virtual_machine.azure.microsoft.com", nil
}
<|endoftext|>


<commit_before>package libkb

//
// Code for encoding and decoding P3SKB-formatted keys. Also works for decoding
// general Keybase Packet types, but we only have P3SKB at present
//

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"github.com/keybase/go-triplesec"
	"golang.org/x/crypto/openpgp"
	"io"
	"os"
)

type P3SKB struct {
	Priv P3SKBPriv `codec:"priv"`
	Pub  []byte    `codec:"pub"`

	decodedPub      *PgpKeyBundle
	decryptedSecret *PgpKeyBundle
}

type P3SKBPriv struct {
	Data       []byte `codec:"data"`
	Encryption int    `codec:"encryption"`
}

func (key *PgpKeyBundle) ToP3SKB(tsec *triplesec.Cipher) (ret *P3SKB, err error) {

	ret = &P3SKB{}

	var pk, sk bytes.Buffer

	// Need to serialize Private first, because
	err = (*openpgp.Entity)(key).SerializePrivate(&sk, nil)
	if err != nil {
		return
	}
	if tsec != nil {
		ret.Priv.Data, err = tsec.Encrypt(sk.Bytes())
		ret.Priv.Encryption = int(triplesec.Version) // Version 3 is the current TripleSec version
		if err != nil {
			return
		}
	} else {
		ret.Priv.Data = sk.Bytes()
		ret.Priv.Encryption = 0
	}

	err = (*openpgp.Entity)(key).Serialize(&pk)
	if err != nil {
		return
	}
	ret.Pub = pk.Bytes()

	return
}

func (p *P3SKB) ToPacket() (ret *KeybasePacket, err error) {
	ret = &KeybasePacket{
		Version: KEYBASE_PACKET_V1,
		Tag:     TAG_P3SKB, // Keybase tags starts at 513 (OpenPGP are 0-30)
	}
	ret.Body = p
	err = ret.HashMe()
	return
}

func (p *P3SKB) GetPubKey() (key *PgpKeyBundle, err error) {
	if p.decodedPub == nil {
		key, err = ReadOneKeyFromBytes(p.Pub)
		p.decodedPub = key
	}
	return
}

func (p *P3SKB) VerboseDescription() (ret string, err error) {
	var key *PgpKeyBundle
	key, err = p.GetPubKey()
	if err == nil {
		ret = key.VerboseDescription()
	}
	return
}

func (p *P3SKB) UnlockSecretKey(tsec *triplesec.Cipher) (key *PgpKeyBundle, err error) {
	if key = p.decryptedSecret; key != nil {
		return
	}
	var unlocked []byte

	if p.Priv.Encryption == 0 {
		unlocked = p.Priv.Data
	} else if unlocked, err = tsec.Decrypt(p.Priv.Data); err != nil {
		if _, ok := err.(triplesec.BadPassphraseError); ok {
			err = PassphraseError{}
		}
		return
	}

	if key, err = ReadOneKeyFromBytes(unlocked); err != nil {
		return
	}

	if key.PrivateKey == nil {
		err = NoKeyError{"no private key found"}
	} else if !key.PrivateKey.Encrypted {
		err = BadKeyError{"PGP key material should be unencrypted"}
	} else {
		p.decryptedSecret = key
	}
	return
}

type P3SKBKeyringFile struct {
	filename string
	Blocks   []*P3SKB
	index    map[PgpFingerprint]*P3SKB
	dirty    bool
}

func NewP3SKBKeyringFile(n string) *P3SKBKeyringFile {
	return &P3SKBKeyringFile{
		filename: n,
		Blocks:   make([]*P3SKB, 0, 1),
		index:    make(map[PgpFingerprint]*P3SKB),
		dirty:    false,
	}
}

func (k *P3SKBKeyringFile) Load() (err error) {
	G.Log.Debug("+ Loading P3SKB keyring: %s", k.filename)
	var packets KeybasePackets
	var file *os.File
	if file, err = os.OpenFile(k.filename, os.O_RDONLY, 0); err == nil {
		stream := base64.NewDecoder(base64.StdEncoding, file)
		packets, err = DecodePackets(stream)
		tmp := file.Close()
		if err == nil && tmp != nil {
			err = tmp
		}
	}

	if err != nil {
		if os.IsNotExist(err) {
			G.Log.Debug("| Keybase secret keyring doesn't exist: %s", k.filename)
		} else {
			G.Log.Warning("Error opening %s: %s", k.filename, err.Error())
		}

	} else if err == nil {
		k.Blocks, err = packets.ToListOfP3SKBs()
	}

	G.Log.Debug("- Loaded P3SKB keyring: %s -> %s", k.filename, ErrToOk(err))
	return
}

func (k *P3SKBKeyringFile) Index() (err error) {
	for _, b := range k.Blocks {
		var key *PgpKeyBundle
		key, err = b.GetPubKey()
		if err != nil {
			return
		}
		fp := key.GetFingerprint()
		k.index[fp] = b
	}
	G.Log.Debug("| Indexed %d secret keys", len(k.Blocks))
	return
}

func (k *P3SKBKeyringFile) LoadAndIndex() error {
	err := k.Load()
	if err == nil {
		err = k.Index()
	}
	return err
}

func (p KeybasePacket) ToP3SKB() (*P3SKB, error) {
	ret, ok := p.Body.(*P3SKB)
	if !ok {
		return nil, fmt.Errorf("Bad P3SKB packet")
	} else {
		return ret, nil
	}
}

func (s *P3SKB) ArmoredEncode() (ret string, err error) {
	var p *KeybasePacket
	if p, err = s.ToPacket(); err == nil {
		ret, err = p.ArmoredEncode()
	}
	return
}

func (f *P3SKBKeyringFile) Push(p3skb *P3SKB) error {
	k, err := p3skb.GetPubKey()
	if err != nil {
		return fmt.Errorf("Failed to get pubkey: %s", err.Error())
	}
	f.dirty = true
	f.Blocks = append(f.Blocks, p3skb)
	fp := PgpFingerprint(k.PrimaryKey.Fingerprint)
	f.index[fp] = p3skb
	return nil
}

func (f P3SKBKeyringFile) GetFilename() string { return f.filename }

func (f P3SKBKeyringFile) WriteTo(w io.Writer) error {
	G.Log.Debug("+ WriteTo")
	packets := make(KeybasePackets, len(f.Blocks))
	var err error
	for i, b := range f.Blocks {
		if packets[i], err = b.ToPacket(); err != nil {
			return err
		}
	}
	b64 := base64.NewEncoder(base64.StdEncoding, w)
	if err = packets.EncodeTo(b64); err != nil {
		G.Log.Warning("Encoding problem: %s", err.Error())
		return err
	}
	G.Log.Debug("- WriteTo")
	b64.Close()
	return nil
}

func (f *P3SKBKeyringFile) Save() error {
	if !f.dirty {
		return nil
	}
	err := SafeWriteToFile(*f)
	if err == nil {
		f.dirty = false
		G.Log.Info("Updated keyring %s", f.filename)
	}
	return err
}

func (p KeybasePackets) ToListOfP3SKBs() (ret []*P3SKB, err error) {
	ret = make([]*P3SKB, len(p))
	for i, e := range p {
		if k, ok := e.Body.(*P3SKB); ok {
			ret[i] = k
		} else {
			err = fmt.Errorf("Bad P3SKB sequence; got packet of wrong type")
			ret = nil
			break
		}
	}
	return
}

func (p *P3SKB) PromptAndUnlock(reason string) (ret *PgpKeyBundle, err error) {
	if ret = p.decryptedSecret; ret != nil {
		return
	}

	// First try the triplsec that we have loaded in (if at all)
	if tsec := G.LoginState.GetCachedTriplesec(); tsec != nil {
		ret, err = p.UnlockSecretKey(tsec)
		if err == nil {
		} else if _, ok := err.(PassphraseError); ok {
			err = nil
		} else {
			return
		}
	}

	var desc string
	if desc, err = p.VerboseDescription(); err != nil {
		return
	}

	unlocker := func(pw string) (ret *PgpKeyBundle, err error) {
		return
	}

	return KeyUnlocker{
		Tries:    5,
		Reason:   reason,
		KeyDesc:  desc,
		Unlocker: unlocker,
	}.Run()
}
<commit_msg>we now have an unlocker for P3SKB keys too<commit_after>package libkb

//
// Code for encoding and decoding P3SKB-formatted keys. Also works for decoding
// general Keybase Packet types, but we only have P3SKB at present
//

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"github.com/keybase/go-triplesec"
	"golang.org/x/crypto/openpgp"
	"io"
	"os"
)

type P3SKB struct {
	Priv P3SKBPriv `codec:"priv"`
	Pub  []byte    `codec:"pub"`

	decodedPub      *PgpKeyBundle
	decryptedSecret *PgpKeyBundle
}

type P3SKBPriv struct {
	Data       []byte `codec:"data"`
	Encryption int    `codec:"encryption"`
}

func (key *PgpKeyBundle) ToP3SKB(tsec *triplesec.Cipher) (ret *P3SKB, err error) {

	ret = &P3SKB{}

	var pk, sk bytes.Buffer

	// Need to serialize Private first, because
	err = (*openpgp.Entity)(key).SerializePrivate(&sk, nil)
	if err != nil {
		return
	}
	if tsec != nil {
		ret.Priv.Data, err = tsec.Encrypt(sk.Bytes())
		ret.Priv.Encryption = int(triplesec.Version) // Version 3 is the current TripleSec version
		if err != nil {
			return
		}
	} else {
		ret.Priv.Data = sk.Bytes()
		ret.Priv.Encryption = 0
	}

	err = (*openpgp.Entity)(key).Serialize(&pk)
	if err != nil {
		return
	}
	ret.Pub = pk.Bytes()

	return
}

func (p *P3SKB) ToPacket() (ret *KeybasePacket, err error) {
	ret = &KeybasePacket{
		Version: KEYBASE_PACKET_V1,
		Tag:     TAG_P3SKB, // Keybase tags starts at 513 (OpenPGP are 0-30)
	}
	ret.Body = p
	err = ret.HashMe()
	return
}

func (p *P3SKB) GetPubKey() (key *PgpKeyBundle, err error) {
	if p.decodedPub == nil {
		key, err = ReadOneKeyFromBytes(p.Pub)
		p.decodedPub = key
	}
	return
}

func (p *P3SKB) VerboseDescription() (ret string, err error) {
	var key *PgpKeyBundle
	key, err = p.GetPubKey()
	if err == nil {
		ret = key.VerboseDescription()
	}
	return
}

func (p *P3SKB) UnlockSecretKey(tsec *triplesec.Cipher) (key *PgpKeyBundle, err error) {
	if key = p.decryptedSecret; key != nil {
		return
	}
	var unlocked []byte

	if p.Priv.Encryption == 0 {
		unlocked = p.Priv.Data
	} else if unlocked, err = tsec.Decrypt(p.Priv.Data); err != nil {
		if _, ok := err.(triplesec.BadPassphraseError); ok {
			err = PassphraseError{}
		}
		return
	}

	if key, err = ReadOneKeyFromBytes(unlocked); err != nil {
		return
	}

	if key.PrivateKey == nil {
		err = NoKeyError{"no private key found"}
	} else if !key.PrivateKey.Encrypted {
		err = BadKeyError{"PGP key material should be unencrypted"}
	} else {
		p.decryptedSecret = key
	}
	return
}

type P3SKBKeyringFile struct {
	filename string
	Blocks   []*P3SKB
	index    map[PgpFingerprint]*P3SKB
	dirty    bool
}

func NewP3SKBKeyringFile(n string) *P3SKBKeyringFile {
	return &P3SKBKeyringFile{
		filename: n,
		Blocks:   make([]*P3SKB, 0, 1),
		index:    make(map[PgpFingerprint]*P3SKB),
		dirty:    false,
	}
}

func (k *P3SKBKeyringFile) Load() (err error) {
	G.Log.Debug("+ Loading P3SKB keyring: %s", k.filename)
	var packets KeybasePackets
	var file *os.File
	if file, err = os.OpenFile(k.filename, os.O_RDONLY, 0); err == nil {
		stream := base64.NewDecoder(base64.StdEncoding, file)
		packets, err = DecodePackets(stream)
		tmp := file.Close()
		if err == nil && tmp != nil {
			err = tmp
		}
	}

	if err != nil {
		if os.IsNotExist(err) {
			G.Log.Debug("| Keybase secret keyring doesn't exist: %s", k.filename)
		} else {
			G.Log.Warning("Error opening %s: %s", k.filename, err.Error())
		}

	} else if err == nil {
		k.Blocks, err = packets.ToListOfP3SKBs()
	}

	G.Log.Debug("- Loaded P3SKB keyring: %s -> %s", k.filename, ErrToOk(err))
	return
}

func (k *P3SKBKeyringFile) Index() (err error) {
	for _, b := range k.Blocks {
		var key *PgpKeyBundle
		key, err = b.GetPubKey()
		if err != nil {
			return
		}
		fp := key.GetFingerprint()
		k.index[fp] = b
	}
	G.Log.Debug("| Indexed %d secret keys", len(k.Blocks))
	return
}

func (k *P3SKBKeyringFile) LoadAndIndex() error {
	err := k.Load()
	if err == nil {
		err = k.Index()
	}
	return err
}

func (p KeybasePacket) ToP3SKB() (*P3SKB, error) {
	ret, ok := p.Body.(*P3SKB)
	if !ok {
		return nil, fmt.Errorf("Bad P3SKB packet")
	} else {
		return ret, nil
	}
}

func (s *P3SKB) ArmoredEncode() (ret string, err error) {
	var p *KeybasePacket
	if p, err = s.ToPacket(); err == nil {
		ret, err = p.ArmoredEncode()
	}
	return
}

func (f *P3SKBKeyringFile) Push(p3skb *P3SKB) error {
	k, err := p3skb.GetPubKey()
	if err != nil {
		return fmt.Errorf("Failed to get pubkey: %s", err.Error())
	}
	f.dirty = true
	f.Blocks = append(f.Blocks, p3skb)
	fp := PgpFingerprint(k.PrimaryKey.Fingerprint)
	f.index[fp] = p3skb
	return nil
}

func (f P3SKBKeyringFile) GetFilename() string { return f.filename }

func (f P3SKBKeyringFile) WriteTo(w io.Writer) error {
	G.Log.Debug("+ WriteTo")
	packets := make(KeybasePackets, len(f.Blocks))
	var err error
	for i, b := range f.Blocks {
		if packets[i], err = b.ToPacket(); err != nil {
			return err
		}
	}
	b64 := base64.NewEncoder(base64.StdEncoding, w)
	if err = packets.EncodeTo(b64); err != nil {
		G.Log.Warning("Encoding problem: %s", err.Error())
		return err
	}
	G.Log.Debug("- WriteTo")
	b64.Close()
	return nil
}

func (f *P3SKBKeyringFile) Save() error {
	if !f.dirty {
		return nil
	}
	err := SafeWriteToFile(*f)
	if err == nil {
		f.dirty = false
		G.Log.Info("Updated keyring %s", f.filename)
	}
	return err
}

func (p KeybasePackets) ToListOfP3SKBs() (ret []*P3SKB, err error) {
	ret = make([]*P3SKB, len(p))
	for i, e := range p {
		if k, ok := e.Body.(*P3SKB); ok {
			ret[i] = k
		} else {
			err = fmt.Errorf("Bad P3SKB sequence; got packet of wrong type")
			ret = nil
			break
		}
	}
	return
}

func (p *P3SKB) PromptAndUnlock(reason string) (ret *PgpKeyBundle, err error) {
	if ret = p.decryptedSecret; ret != nil {
		return
	}

	// First try the triplsec that we have loaded in (if at all)
	if tsec := G.LoginState.GetCachedTriplesec(); tsec != nil {
		ret, err = p.UnlockSecretKey(tsec)
		if err == nil {
		} else if _, ok := err.(PassphraseError); ok {
			err = nil
		} else {
			return
		}
	}

	var desc string
	if desc, err = p.VerboseDescription(); err != nil {
		return
	}

	unlocker := func(pw string) (ret *PgpKeyBundle, err error) {
		var tsec *triplesec.Cipher
		tsec, err = triplesec.NewCipher([]byte(pw), nil)
		if err == nil {
			ret, err = p.UnlockSecretKey(tsec)
		}
		return
	}

	return KeyUnlocker{
		Tries:    5,
		Reason:   reason,
		KeyDesc:  desc,
		Unlocker: unlocker,
	}.Run()
}
<|endoftext|>


<commit_before>package npdfpages

import (
	"bufio"
	"io"
	"os"
	"regexp"
)

var exp, _ = regexp.Compile(`/Page\b`)

func PagesInRuneReader(reader io.RuneReader) (pages int) {
	for exp.FindReaderIndex(reader) != nil {
		pages++
	}
	return
}

func PagesAtPath(path string) (pages int) {
	if reader, err := os.Open(path); err == nil {
		pages = PagesInRuneReader(bufio.NewReader(reader))
		reader.Close()
	}
	return
}
<commit_msg>10x speed improvement.<commit_after>package npdfpages

import (
	"bufio"
	"io"
	"os"
)

// guaranteed not to overlap or be adjacent
const match = "/Page\x00"

func Pages(reader io.ByteReader) (pages int) {
	i := 0
	for {
		b, err := reader.ReadByte()
		if err != nil {
			return
		}
	check:
		switch match[i] {
		case 0:
			if !(b >= 'A' && b <= 'Z' || b >= 'a' && b <= 'z') {
				pages++
			}
			i = 0
			goto check
		case b:
			i++
		default:
			i = 0
		}
	}
	// flow shouldn't reach this point
	return
}

func PagesAtPath(path string) (pages int) {
	if reader, err := os.Open(path); err == nil {
		reader.Chdir()
		pages = Pages(bufio.NewReader(reader))
		reader.Close()
	}
	return
}
<|endoftext|>


<commit_before>package sql

import (
	"database/sql"
	"errors"
	"fmt"
	"github.com/golang/glog"
	_ "github.com/lib/pq"
	"sync"
)

func NewPostgres() *Postgres {
	return &Postgres{
		User:            "ubuntu",
		Db:              "circle_test",
		Ssl:             false,
		DoCreateSchemas: true,
		DoUpdateSchemas: false,
	}
}

type Postgres struct {
	Host            string
	Port            int
	Db              string
	User            string
	Password        string
	Schemas         []*Schema
	DoCreateSchemas bool
	DoUpdateSchemas bool

	Ssl  bool
	conn *sql.DB
}

var (
	initialize_system sync.Once
)

func (this *Postgres) Conn() *sql.DB {
	return this.conn
}

func (this *Postgres) Open() error {
	glog.Infoln("Postgres connection string:", this.connection_string())

	db, err := sql.Open(string(POSTGRES), this.connection_string())
	if err != nil {
		return err
	}

	this.conn = db
	glog.Infoln("Connected:", this.conn)
	initialize_system.Do(func() {
		// bootstrap the system schema
		err1 := postgres_schema.Initialize(db)
		glog.Infoln("Initialized system schema:", err1)
		err2 := postgres_schema.PrepareStatements(db)
		glog.Infoln("Prepared statements:", err2)
	})
	err = sync_schemas(this.conn, this.Schemas, this.DoCreateSchemas, this.DoUpdateSchemas)
	if err != nil {
		return err
	}
	// prepare statements
	for _, s := range this.Schemas {
		err = s.PrepareStatements(this.conn)
		if err != nil {
			return err
		}
	}
	return nil
}

func (this *Postgres) Upsert(schema *Schema, update, insert StatementKey, args ...interface{}) error {
	return schema.Upsert(this.conn, update, insert, args...)
}

func (this *Postgres) Delete(schema *Schema, delete StatementKey, args ...interface{}) error {
	return schema.Delete(this.conn, delete, args...)
}

func (this *Postgres) GetOne(schema *Schema, get StatementKey, opt *Options, args ...interface{}) error {
	return schema.GetOne(this.conn, get, opt, args...)
}

func (this *Postgres) Close() error {
	if this.conn == nil {
		return ErrNotConnected
	}
	return this.conn.Close()
}

func (this *Postgres) DropAll() error {
	// This just drops everything... dangerous!
	for _, s := range this.Schemas {
		err := s.DropTables(this.conn)
		if err != nil {
			return err
		}
	}
	return nil
}

func (this *Postgres) TruncateAll() error {
	for _, s := range this.Schemas {
		for t, _ := range s.CreateTables {
			glog.Infoln("Truncating", t)
			_, err := this.conn.Exec("delete from " + t + " where true")
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func (this *Postgres) connection_string() string {
	cs := fmt.Sprintf("user=%s dbname=%s", this.User, this.Db)
	if this.Password != "" {
		cs += fmt.Sprintf(" password='%s'", this.Password)
	}
	if this.Host != "" {
		cs += fmt.Sprintf(" host=%s", this.Host)
	}
	if this.Port > 0 {
		cs += fmt.Sprintf(" port=%d", this.Port)
	}
	if this.Ssl {
		cs += " sslmode=verify-full"
	} else {
		cs += " sslmode=disable"
	}
	return cs
}

var postgres_schema = &Schema{
	Platform: POSTGRES,
	Name:     "system",
	Version:  1,
	CreateTables: map[string]string{
		"schema_versions": `
create table if not exists system_schema_versions (
    schema_name varchar,
    version     integer,
    commit_hash varchar null
)
		`,
	},
	PreparedStatements: map[StatementKey]Statement{
		kSelectVersionInfoBySchemaName: Statement{
			Query: `
select version, commit_hash from system_schema_versions
where schema_name=$1
`},
		kDeleteVersionInfo: Statement{
			Query: `
delete from system_schema_versions where schema_name=$1
`,
			Args: func(args ...interface{}) ([]interface{}, error) {
				if len(args) != 1 {
					return nil, errors.New("args-mismatch")
				}
				schema_name, ok := args[0].(string)
				if !ok {
					return nil, errors.New("bad-schema-name")
				}
				return []interface{}{
					schema_name,
				}, nil
			},
		},
		kInsertVersionInfo: Statement{
			Query: `
insert into system_schema_versions (schema_name, version, commit_hash)
values ($1, $2, $3)
`,
			Args: func(args ...interface{}) ([]interface{}, error) {
				if len(args) != 3 {
					return nil, errors.New("args-mismatch")
				}
				schema_name, ok := args[0].(string)
				if !ok {
					return nil, errors.New("bad-schema-name")
				}
				version, ok := args[1].(int)
				if !ok {
					return nil, errors.New("bad-version")
				}
				commit_hash, ok := args[2].(string)
				if !ok {
					return nil, errors.New("bad-commit-hash")
				}
				return []interface{}{
					schema_name,
					version,
					commit_hash,
				}, nil
			},
		},
		kUpdateVersionInfo: Statement{
			Query: `
update system_schema_versions
set version=$1, commit_hash=$2
where schema_name=$3
`,
			Args: func(args ...interface{}) ([]interface{}, error) {
				if len(args) != 3 {
					return nil, errors.New("args-mismatch")
				}
				schema_name, ok := args[0].(string)
				if !ok {
					return nil, errors.New("bad-schema-name")
				}
				version, ok := args[1].(int)
				if !ok {
					return nil, errors.New("bad-version")
				}
				commit_hash, ok := args[2].(string)
				if !ok {
					return nil, errors.New("bad-commit-hash")
				}
				return []interface{}{
					version,
					commit_hash,
					schema_name,
				}, nil
			},
		},
	},
}
<commit_msg>Panic when failed to initialize db<commit_after>package sql

import (
	"database/sql"
	"errors"
	"fmt"
	"github.com/golang/glog"
	_ "github.com/lib/pq"
	"sync"
)

func NewPostgres() *Postgres {
	return &Postgres{
		User:            "ubuntu",
		Db:              "circle_test",
		Ssl:             false,
		DoCreateSchemas: true,
		DoUpdateSchemas: false,
	}
}

type Postgres struct {
	Host            string
	Port            int
	Db              string
	User            string
	Password        string
	Schemas         []*Schema
	DoCreateSchemas bool
	DoUpdateSchemas bool

	Ssl  bool
	conn *sql.DB
}

var (
	initialize_system sync.Once
)

func (this *Postgres) Conn() *sql.DB {
	return this.conn
}

func (this *Postgres) Open() error {
	glog.Infoln("Postgres connection string:", this.connection_string())

	db, err := sql.Open(string(POSTGRES), this.connection_string())
	if err != nil {
		return err
	}

	this.conn = db
	glog.Infoln("Connected:", this.conn)
	initialize_system.Do(func() {
		// bootstrap the system schema
		err1 := postgres_schema.Initialize(db)
		glog.Infoln("Initialized system schema:", err1)
		if err1 != nil {
			panic(err1)
		}
		err2 := postgres_schema.PrepareStatements(db)
		glog.Infoln("Prepared statements:", err2)
		if err2 != nil {
			panic(err2)
		}
	})
	err = sync_schemas(this.conn, this.Schemas, this.DoCreateSchemas, this.DoUpdateSchemas)
	if err != nil {
		return err
	}
	// prepare statements
	for _, s := range this.Schemas {
		err = s.PrepareStatements(this.conn)
		if err != nil {
			return err
		}
	}
	return nil
}

func (this *Postgres) Upsert(schema *Schema, update, insert StatementKey, args ...interface{}) error {
	return schema.Upsert(this.conn, update, insert, args...)
}

func (this *Postgres) Delete(schema *Schema, delete StatementKey, args ...interface{}) error {
	return schema.Delete(this.conn, delete, args...)
}

func (this *Postgres) GetOne(schema *Schema, get StatementKey, opt *Options, args ...interface{}) error {
	return schema.GetOne(this.conn, get, opt, args...)
}

func (this *Postgres) Close() error {
	if this.conn == nil {
		return ErrNotConnected
	}
	return this.conn.Close()
}

func (this *Postgres) DropAll() error {
	// This just drops everything... dangerous!
	for _, s := range this.Schemas {
		err := s.DropTables(this.conn)
		if err != nil {
			return err
		}
	}
	return nil
}

func (this *Postgres) TruncateAll() error {
	for _, s := range this.Schemas {
		for t, _ := range s.CreateTables {
			glog.Infoln("Truncating", t)
			_, err := this.conn.Exec("delete from " + t + " where true")
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func (this *Postgres) connection_string() string {
	cs := fmt.Sprintf("user=%s dbname=%s", this.User, this.Db)
	if this.Password != "" {
		cs += fmt.Sprintf(" password='%s'", this.Password)
	}
	if this.Host != "" {
		cs += fmt.Sprintf(" host=%s", this.Host)
	}
	if this.Port > 0 {
		cs += fmt.Sprintf(" port=%d", this.Port)
	}
	if this.Ssl {
		cs += " sslmode=verify-full"
	} else {
		cs += " sslmode=disable"
	}
	return cs
}

var postgres_schema = &Schema{
	Platform: POSTGRES,
	Name:     "system",
	Version:  1,
	CreateTables: map[string]string{
		"schema_versions": `
create table if not exists system_schema_versions (
    schema_name varchar,
    version     integer,
    commit_hash varchar null
)
		`,
	},
	PreparedStatements: map[StatementKey]Statement{
		kSelectVersionInfoBySchemaName: Statement{
			Query: `
select version, commit_hash from system_schema_versions
where schema_name=$1
`},
		kDeleteVersionInfo: Statement{
			Query: `
delete from system_schema_versions where schema_name=$1
`,
			Args: func(args ...interface{}) ([]interface{}, error) {
				if len(args) != 1 {
					return nil, errors.New("args-mismatch")
				}
				schema_name, ok := args[0].(string)
				if !ok {
					return nil, errors.New("bad-schema-name")
				}
				return []interface{}{
					schema_name,
				}, nil
			},
		},
		kInsertVersionInfo: Statement{
			Query: `
insert into system_schema_versions (schema_name, version, commit_hash)
values ($1, $2, $3)
`,
			Args: func(args ...interface{}) ([]interface{}, error) {
				if len(args) != 3 {
					return nil, errors.New("args-mismatch")
				}
				schema_name, ok := args[0].(string)
				if !ok {
					return nil, errors.New("bad-schema-name")
				}
				version, ok := args[1].(int)
				if !ok {
					return nil, errors.New("bad-version")
				}
				commit_hash, ok := args[2].(string)
				if !ok {
					return nil, errors.New("bad-commit-hash")
				}
				return []interface{}{
					schema_name,
					version,
					commit_hash,
				}, nil
			},
		},
		kUpdateVersionInfo: Statement{
			Query: `
update system_schema_versions
set version=$1, commit_hash=$2
where schema_name=$3
`,
			Args: func(args ...interface{}) ([]interface{}, error) {
				if len(args) != 3 {
					return nil, errors.New("args-mismatch")
				}
				schema_name, ok := args[0].(string)
				if !ok {
					return nil, errors.New("bad-schema-name")
				}
				version, ok := args[1].(int)
				if !ok {
					return nil, errors.New("bad-version")
				}
				commit_hash, ok := args[2].(string)
				if !ok {
					return nil, errors.New("bad-commit-hash")
				}
				return []interface{}{
					version,
					commit_hash,
					schema_name,
				}, nil
			},
		},
	},
}
<|endoftext|>


<commit_before>package push

import (
	"strings"

	"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("push with only an app name", func() {
	var (
		appName string
	)

	BeforeEach(func() {
		appName = helpers.NewAppName()
	})

	Describe("app existence", func() {
		Context("when the app does not exist", func() {
			It("creates the app", func() {
				helpers.WithHelloWorldApp(func(dir string) {
					session := helpers.CustomCF(helpers.CFEnv{WorkingDirectory: dir}, PushCommandName, appName)
					Eventually(session).Should(Say("Getting app info\\.\\.\\."))
					Eventually(session).Should(Say("Creating app with these attributes\\.\\.\\."))
					Eventually(session).Should(Say("\\+\\s+name:\\s+%s", appName))
					Eventually(session).Should(Say("\\s+routes:"))
					Eventually(session).Should(Say("(?i)\\+\\s+%s.%s", appName, defaultSharedDomain()))
					Eventually(session).Should(Say("Mapping routes\\.\\.\\."))
					Eventually(session).Should(Say("Comparing local files to remote cache\\.\\.\\."))
					Eventually(session).Should(Say("Packaging files to upload\\.\\.\\."))
					Eventually(session).Should(Say("Uploading files\\.\\.\\."))
					Eventually(session).Should(Say("100.00%"))
					Eventually(session).Should(Say("Waiting for API to complete processing files\\.\\.\\."))
					Eventually(session).Should(Say("Staging complete"))
					Eventually(session).Should(Say("Waiting for app to start\\.\\.\\."))
					Eventually(session).Should(Say("requested state:\\s+started"))
					Eventually(session).Should(Exit(0))
				})

				session := helpers.CF("app", appName)
				Eventually(session).Should(Say("name:\\s+%s", appName))
				Eventually(session).Should(Exit(0))
			})
		})

		Context("when the app exists", func() {
			BeforeEach(func() {
				helpers.WithHelloWorldApp(func(dir string) {
					Eventually(helpers.CustomCF(helpers.CFEnv{WorkingDirectory: dir}, "push", appName)).Should(Exit(0))
				})
			})

			It("updates the app", func() {
				helpers.WithHelloWorldApp(func(dir string) {
					session := helpers.CustomCF(helpers.CFEnv{WorkingDirectory: dir}, PushCommandName, appName)
					Eventually(session).Should(Say("Getting app info\\.\\.\\."))
					Eventually(session).Should(Say("Updating app with these attributes\\.\\.\\."))
					Eventually(session).Should(Say("(?m)^\\s+name:\\s+%s$", appName))
					Eventually(session).Should(Say("\\s+routes:"))
					Eventually(session).Should(Say("(?mi)^\\s+%s.%s$", strings.ToLower(appName), defaultSharedDomain()))
					Eventually(session).Should(Say("Mapping routes\\.\\.\\."))
					Eventually(session).Should(Say("Uploading files\\.\\.\\."))
					Eventually(session).Should(Say("100.00%"))
					Eventually(session).Should(Say("Waiting for API to complete processing files\\.\\.\\."))
					Eventually(session).Should(Say("Stopping app\\.\\.\\."))
					Eventually(session).Should(Say("Staging complete"))
					Eventually(session).Should(Say("Waiting for app to start\\.\\.\\."))
					Eventually(session).Should(Say("requested state:\\s+started"))
					Eventually(session).Should(Exit(0))
				})

				session := helpers.CF("app", appName)
				Eventually(session).Should(Say("name:\\s+%s", appName))
				Eventually(session).Should(Exit(0))
			})
		})
	})
})
<commit_msg>use helper instead of looking for 'Staging complete'<commit_after>package push

import (
	"strings"

	"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("push with only an app name", func() {
	var (
		appName string
	)

	BeforeEach(func() {
		appName = helpers.NewAppName()
	})

	Describe("app existence", func() {
		Context("when the app does not exist", func() {
			It("creates the app", func() {
				helpers.WithHelloWorldApp(func(dir string) {
					session := helpers.CustomCF(helpers.CFEnv{WorkingDirectory: dir}, PushCommandName, appName)
					Eventually(session).Should(Say("Getting app info\\.\\.\\."))
					Eventually(session).Should(Say("Creating app with these attributes\\.\\.\\."))
					Eventually(session).Should(Say("\\+\\s+name:\\s+%s", appName))
					Eventually(session).Should(Say("\\s+routes:"))
					Eventually(session).Should(Say("(?i)\\+\\s+%s.%s", appName, defaultSharedDomain()))
					Eventually(session).Should(Say("Mapping routes\\.\\.\\."))
					Eventually(session).Should(Say("Comparing local files to remote cache\\.\\.\\."))
					Eventually(session).Should(Say("Packaging files to upload\\.\\.\\."))
					Eventually(session).Should(Say("Uploading files\\.\\.\\."))
					Eventually(session).Should(Say("100.00%"))
					Eventually(session).Should(Say("Waiting for API to complete processing files\\.\\.\\."))
					helpers.ConfirmStagingLogs(session)
					Eventually(session).Should(Say("Waiting for app to start\\.\\.\\."))
					Eventually(session).Should(Say("requested state:\\s+started"))
					Eventually(session).Should(Exit(0))
				})

				session := helpers.CF("app", appName)
				Eventually(session).Should(Say("name:\\s+%s", appName))
				Eventually(session).Should(Exit(0))
			})
		})

		Context("when the app exists", func() {
			BeforeEach(func() {
				helpers.WithHelloWorldApp(func(dir string) {
					Eventually(helpers.CustomCF(helpers.CFEnv{WorkingDirectory: dir}, "push", appName)).Should(Exit(0))
				})
			})

			It("updates the app", func() {
				helpers.WithHelloWorldApp(func(dir string) {
					session := helpers.CustomCF(helpers.CFEnv{WorkingDirectory: dir}, PushCommandName, appName)
					Eventually(session).Should(Say("Getting app info\\.\\.\\."))
					Eventually(session).Should(Say("Updating app with these attributes\\.\\.\\."))
					Eventually(session).Should(Say("(?m)^\\s+name:\\s+%s$", appName))
					Eventually(session).Should(Say("\\s+routes:"))
					Eventually(session).Should(Say("(?mi)^\\s+%s.%s$", strings.ToLower(appName), defaultSharedDomain()))
					Eventually(session).Should(Say("Mapping routes\\.\\.\\."))
					Eventually(session).Should(Say("Uploading files\\.\\.\\."))
					Eventually(session).Should(Say("100.00%"))
					Eventually(session).Should(Say("Waiting for API to complete processing files\\.\\.\\."))
					Eventually(session).Should(Say("Stopping app\\.\\.\\."))
					helpers.ConfirmStagingLogs(session)
					Eventually(session).Should(Say("Waiting for app to start\\.\\.\\."))
					Eventually(session).Should(Say("requested state:\\s+started"))
					Eventually(session).Should(Exit(0))
				})

				session := helpers.CF("app", appName)
				Eventually(session).Should(Say("name:\\s+%s", appName))
				Eventually(session).Should(Exit(0))
			})
		})
	})
})
<|endoftext|>


<commit_before>package main

import (
  "os"
  "fmt"
  "io"
  "bufio"
  "strings"
  "strconv"
  "path"
  "path/filepath"
  "errors"
  "bytes"
  "go/token"
  "go/parser"
  "go/printer"
  "golang.org/x/tools/go/ast/astutil"
)

var verbose bool = false
var log *os.File = os.Stderr

type Section struct {
  Kind uint
  Text string
}
const (
  StaticSection uint = iota
  CodeSection
)

var sections []*Section

//--- Linear pattern matcher

type Pattern struct {
  Text []rune
  Length, Pos int
}

func newPattern(s string) Pattern {
  runes := []rune(s)
  return Pattern{ Text: runes, Length: len(runes) }
}

// Next returns true if Pos advances past the last character of Text.
func (pattern *Pattern) Next(ch rune) bool {
  // If Pos is past the end of Text, reset it to the beginning.
  if pattern.Pos == pattern.Length {
    pattern.Pos = 0
  }
  // Try to match the current rune in Text.
  if ch == pattern.Text[pattern.Pos] {
    pattern.Pos++
  }
  // Check for a complete match.
  return pattern.Pos == pattern.Length
}

//--- Template parsing and output generation

type TemplateEntry struct {
  SitePath, HardPath string
  FileInfo os.FileInfo
  InsertionLine int
}

func (entry TemplateEntry) String() string {
  if entry.InsertionLine == 0 {
    return entry.SitePath
  }
  return fmt.Sprintf("-> line %d: %s", entry.InsertionLine, entry.SitePath)
}

func makeTemplateEntry(siteRoot, startDir, sitePath string,
    insertionLine int) (*TemplateEntry, error) {
  hardPath := makeHardPath(siteRoot, startDir, sitePath)
  fileInfo, error := os.Stat(hardPath)
  if error != nil {
    return nil, error
  }
  entry := TemplateEntry{
      SitePath: sitePath,
      HardPath: hardPath,
      FileInfo: fileInfo,
      InsertionLine: insertionLine,
    }
  return &entry, nil
}

func makeHardPath(siteRoot, startDir, sitePath string) string {
  // A hard path names a location in the physical file system rather than
  //  in the website's directory structure. It is either an absolute path
  //  or a relative path with respect to the directory containing the
  //  top-level template that is being parsed.
  var dir string
  if filepath.IsAbs(sitePath) {
    dir = siteRoot
  } else {
    dir = startDir
  }
  // Note that filepath.Join automatically performs filepath.Clean, thus
  //  returning a lexically unique form of the path. However, the path
  //  does not uniquely identify a file if it includes a symbolic link.
  //  Therefore, we cannot rely on string comparison to prevent cycles.
  hardPath := filepath.Join(dir, sitePath)
  return hardPath
}

func parse(templatePath string) error {

  // We resolve absolute paths with the website root.
  siteRoot := "/var/www/dd1"  // Stub. We'll get the real value from Apache.
  // We resolve relative paths using the starting directory.
  startDir := filepath.Dir(templatePath)

  // Make an insertion stack with a top-level entry.
  entry, error := makeTemplateEntry(siteRoot, startDir, templatePath, 0)
  if error != nil {
    return error
  }
  stack := []*TemplateEntry{ entry }
  return doParse(siteRoot, startDir, stack)
}

func doParse(siteRoot, startDir string, stack []*TemplateEntry) error {
  current := stack[len(stack)-1]
  if verbose {
    fmt.Fprintf(log, "// start \"%s\"\n", current.SitePath)
  }
  var topLevel bool
  if len(stack) == 1 {
    topLevel = true
  }

  // Check for an insertion cycle.
  for i := len(stack)-2; i >= 0; i-- {
    ancestor := stack[i]
    if os.SameFile(ancestor.FileInfo, current.FileInfo) {
      lines := []string{ "doParse: insertion cycle" }
      for j := i; j < len(stack); j++ {
        lines = append(lines, stack[j].String())
      }
      message := fmt.Sprintf(strings.Join(lines, "\n  "))
      return errors.New(message)
    }
  }

  var error error
  var file *os.File

  file, error = os.Open(current.HardPath)
  if error != nil {
    return error
  }

  reader := bufio.NewReader(file)
  writer := bufio.NewWriter(os.Stdout)
  defer writer.Flush()

  codePattern := newPattern("<?code")
  insertPattern := newPattern("<?insert")
  openPatterns := []*Pattern{ &codePattern, &insertPattern }
  var open *Pattern
  close := newPattern("?>")

  var buffer []rune
  var ch rune
  var size int
  countBytes, countRunes := 0, 0
  lineIndex := 1
  prefix := true

  for {
    ch, size, error = reader.ReadRune()
    if error == nil {
      buffer = append(buffer, ch)
      countBytes += size
      countRunes += 1
      if ch == '\n' {
        lineIndex += 1
      }
    } else {
      content := string(buffer)
      if topLevel {
        content = strings.TrimSpace(content)
      }
      emitStatic(content)
      break
    }

    if open == nil {
      for _, pattern := range openPatterns {
        if pattern.Next(ch) {
          open = pattern
          content := string(buffer[0:len(buffer)-open.Length])
          if prefix {
            if topLevel {
              content = strings.TrimSpace(content)
            }
            prefix = false
          }
          emitStatic(content)
          buffer = []rune{}
        }
      }
    } else {
      if close.Next(ch) {
        content := buffer[0:len(buffer)-close.Length]
        if open == &codePattern {
          emitCode(string(content))
        } else if open == &insertPattern {
          childPath := strings.TrimSpace(string(content))
          entry, error := makeTemplateEntry(siteRoot, startDir, childPath,
              lineIndex)
          if error != nil {
            return error
          }
          stack = append(stack, entry)
          error = doParse(siteRoot, startDir, stack)
          if error != nil {
            return error
          }
          stack = stack[0:len(stack)-1]
        }
        open = nil
        buffer = []rune{}
      }
    }
  }
  if verbose {
    fmt.Fprintf(log, "// finish \"%s\"\n", current.SitePath)
    fmt.Fprintf(log, "// read %d bytes, %d runes\n", countBytes, countRunes)
    fmt.Fprintf(log, "// finished on line %d\n", lineIndex)
  }
  if error == io.EOF {
    return nil
  }
  return error
}

func emitCode(content string) {
  sections = append(sections, &Section{ Kind: CodeSection, Text: content })
}

func emitStatic(content string) {
  if len(content) == 0 {
    return
  }
  from := 0
  for pos, ch := range content {
    if ch == '`' {
      if pos != from {
        raw := fmt.Sprintf("`%s`", content[from:pos])
        emitStaticChunk(raw)
      }
      emitStaticChunk("'`'")
      from = pos+1
    }
  }
  if from != len(content) {
    raw := fmt.Sprintf("`%s`", content[from:len(content)])
    emitStaticChunk(raw)
  }
}
func emitStaticChunk(chunk string) {
  sections = append(sections, &Section{ Kind: StaticSection, Text: chunk })
}

func main() {
  writer := bufio.NewWriter(os.Stdout)
  defer writer.Flush()

  numFiles := len(os.Args)-1
  if numFiles == 0 {
    writer.WriteString("No files specified.\n")
    return
  }
  for argIx := 1; argIx <= numFiles; argIx++ {
    // Parse the top-level template. The resulting source code goes into
    //  the global variable sections.
    templatePath := os.Args[argIx]
    error := parse(templatePath)
    if error == nil {
      // Concatenate the code sections and run them through the Go parser.
      output := bytes.Buffer{}
      for _, section := range sections {
        if section.Kind == CodeSection {
          fmt.Fprintf(&output, section.Text)
        } else {
          s := fmt.Sprintf(";fmt.Print(%s);", section.Text)
          fmt.Fprintf(&output, s)
        }
      }
      fileSet := token.NewFileSet()
      fileNode, error := parser.ParseFile(fileSet, "output", output.Bytes(),
          parser.ParseComments)
      if error == nil {
        // Before injecting an import statement, we have to check for
        //  various collisions. The astutil.AddNameImport function only
        //  checks for an import with the same path. There are more cases.
        //  0. no import has the name or path "fmt":
        //    import "fmt"
        //  1. an import has the path "fmt" with a different name:
        //    import "fmt"
        //  2. an import has the name "fmt" with a different path:
        //    import _fmt0 "fmt" (or _fmt1, ...) and write _fmt0.Printf
        //  3. an import has the name and path "fmt":
        //    do nothing
        maxCollisionCase := 0
        for _, importSpec := range fileNode.Imports {
          importPath, _ := strconv.Unquote(importSpec.Path.Value)
          var importName string
          if importSpec.Name == nil {
            importName = path.Base(importPath)
          } else {
            importName = importSpec.Name.Name
          }
          collisionCase := 0
          if importPath == "fmt" {
            if importName == "fmt" {
              collisionCase = 3
            } else {
              collisionCase = 1
            }
          } else if importName == "fmt" {
            collisionCase = 2
          }
          if collisionCase > maxCollisionCase {
            maxCollisionCase = collisionCase
          }
        }

        fmt.Fprintf(log, "maxCollisionCase: %d\n", maxCollisionCase)
        if maxCollisionCase <= 1 {
          astutil.AddImport(fileSet, fileNode, "fmt")
        } else if maxCollisionCase == 2 {
        }

        // Print with a custom configuration: soft tabs of two spaces each.
        config := printer.Config{ Mode: printer.UseSpaces, Tabwidth: 2 }
        (&config).Fprint(writer, fileSet, fileNode)
      } else {
        // If there was a Go parsing error, print our template-generated
        //  source code, followed by the Go parser's error message.
        writer.Write(output.Bytes())
        writer.WriteString("\n----------\n")
        writer.WriteString(fmt.Sprintf("Go parsing error: %s\n", error))
      }
    } else {
      writer.WriteString(fmt.Sprintf("template parsing error: %s\n", error))
    }
  }
}
<commit_msg>Generate code and insert print statements<commit_after>package main

import (
  "os"
  "fmt"
  "io"
  "bufio"
  "strings"
  "strconv"
  "path"
  "path/filepath"
  "errors"
  "bytes"
  "go/token"
  "go/parser"
  "go/printer"
  "golang.org/x/tools/go/ast/astutil"
)

var verbose bool = false
var log *os.File = os.Stderr

type Section struct {
  Kind uint
  Text string
}
const (
  StaticSection uint = iota
  CodeSection
)

var sections []*Section

//--- Linear pattern matcher

type Pattern struct {
  Text []rune
  Length, Pos int
}

func newPattern(s string) Pattern {
  runes := []rune(s)
  return Pattern{ Text: runes, Length: len(runes) }
}

// Next returns true if Pos advances past the last character of Text.
func (pattern *Pattern) Next(ch rune) bool {
  // If Pos is past the end of Text, reset it to the beginning.
  if pattern.Pos == pattern.Length {
    pattern.Pos = 0
  }
  // Try to match the current rune in Text.
  if ch == pattern.Text[pattern.Pos] {
    pattern.Pos++
  }
  // Check for a complete match.
  return pattern.Pos == pattern.Length
}

//--- Template parsing and output generation

type TemplateEntry struct {
  SitePath, HardPath string
  FileInfo os.FileInfo
  InsertionLine int
}

func (entry TemplateEntry) String() string {
  if entry.InsertionLine == 0 {
    return entry.SitePath
  }
  return fmt.Sprintf("-> line %d: %s", entry.InsertionLine, entry.SitePath)
}

func makeTemplateEntry(siteRoot, startDir, sitePath string,
    insertionLine int) (*TemplateEntry, error) {
  hardPath := makeHardPath(siteRoot, startDir, sitePath)
  fileInfo, error := os.Stat(hardPath)
  if error != nil {
    return nil, error
  }
  entry := TemplateEntry{
      SitePath: sitePath,
      HardPath: hardPath,
      FileInfo: fileInfo,
      InsertionLine: insertionLine,
    }
  return &entry, nil
}

func makeHardPath(siteRoot, startDir, sitePath string) string {
  // A hard path names a location in the physical file system rather than
  //  in the website's directory structure. It is either an absolute path
  //  or a relative path with respect to the directory containing the
  //  top-level template that is being parsed.
  var dir string
  if filepath.IsAbs(sitePath) {
    dir = siteRoot
  } else {
    dir = startDir
  }
  // Note that filepath.Join automatically performs filepath.Clean, thus
  //  returning a lexically unique form of the path. However, the path
  //  does not uniquely identify a file if it includes a symbolic link.
  //  Therefore, we cannot rely on string comparison to prevent cycles.
  hardPath := filepath.Join(dir, sitePath)
  return hardPath
}

func parse(templatePath string) error {

  // We resolve absolute paths with the website root.
  siteRoot := "/var/www/dd1"  // Stub. We'll get the real value from Apache.
  // We resolve relative paths using the starting directory.
  startDir := filepath.Dir(templatePath)

  // Make an insertion stack with a top-level entry.
  entry, error := makeTemplateEntry(siteRoot, startDir, templatePath, 0)
  if error != nil {
    return error
  }
  stack := []*TemplateEntry{ entry }
  return doParse(siteRoot, startDir, stack)
}

func doParse(siteRoot, startDir string, stack []*TemplateEntry) error {
  current := stack[len(stack)-1]
  if verbose {
    fmt.Fprintf(log, "// start \"%s\"\n", current.SitePath)
  }
  var topLevel bool
  if len(stack) == 1 {
    topLevel = true
  }

  // Check for an insertion cycle.
  for i := len(stack)-2; i >= 0; i-- {
    ancestor := stack[i]
    if os.SameFile(ancestor.FileInfo, current.FileInfo) {
      lines := []string{ "doParse: insertion cycle" }
      for j := i; j < len(stack); j++ {
        lines = append(lines, stack[j].String())
      }
      message := fmt.Sprintf(strings.Join(lines, "\n  "))
      return errors.New(message)
    }
  }

  var error error
  var file *os.File

  file, error = os.Open(current.HardPath)
  if error != nil {
    return error
  }

  reader := bufio.NewReader(file)
  writer := bufio.NewWriter(os.Stdout)
  defer writer.Flush()

  codePattern := newPattern("<?code")
  insertPattern := newPattern("<?insert")
  openPatterns := []*Pattern{ &codePattern, &insertPattern }
  var open *Pattern
  close := newPattern("?>")

  var buffer []rune
  var ch rune
  var size int
  countBytes, countRunes := 0, 0
  lineIndex := 1
  prefix := true

  for {
    ch, size, error = reader.ReadRune()
    if error == nil {
      buffer = append(buffer, ch)
      countBytes += size
      countRunes += 1
      if ch == '\n' {
        lineIndex += 1
      }
    } else {
      content := string(buffer)
      if topLevel {
        content = strings.TrimSpace(content)
      }
      emitStatic(content)
      break
    }

    if open == nil {
      for _, pattern := range openPatterns {
        if pattern.Next(ch) {
          open = pattern
          content := string(buffer[0:len(buffer)-open.Length])
          if prefix {
            if topLevel {
              content = strings.TrimSpace(content)
            }
            prefix = false
          }
          emitStatic(content)
          buffer = []rune{}
        }
      }
    } else {
      if close.Next(ch) {
        content := buffer[0:len(buffer)-close.Length]
        if open == &codePattern {
          emitCode(string(content))
        } else if open == &insertPattern {
          childPath := strings.TrimSpace(string(content))
          entry, error := makeTemplateEntry(siteRoot, startDir, childPath,
              lineIndex)
          if error != nil {
            return error
          }
          stack = append(stack, entry)
          error = doParse(siteRoot, startDir, stack)
          if error != nil {
            return error
          }
          stack = stack[0:len(stack)-1]
        }
        open = nil
        buffer = []rune{}
      }
    }
  }
  if verbose {
    fmt.Fprintf(log, "// finish \"%s\"\n", current.SitePath)
    fmt.Fprintf(log, "// read %d bytes, %d runes\n", countBytes, countRunes)
    fmt.Fprintf(log, "// finished on line %d\n", lineIndex)
  }
  if error == io.EOF {
    return nil
  }
  return error
}

func emitCode(content string) {
  sections = append(sections, &Section{ Kind: CodeSection, Text: content })
}

func emitStatic(content string) {
  if len(content) == 0 {
    return
  }
  from := 0
  for pos, ch := range content {
    if ch == '`' {
      if pos != from {
        raw := fmt.Sprintf("`%s`", content[from:pos])
        emitStaticChunk(raw)
      }
      emitStaticChunk("'`'")
      from = pos+1
    }
  }
  if from != len(content) {
    raw := fmt.Sprintf("`%s`", content[from:len(content)])
    emitStaticChunk(raw)
  }
}
func emitStaticChunk(chunk string) {
  sections = append(sections, &Section{ Kind: StaticSection, Text: chunk })
}

func main() {
  writer := bufio.NewWriter(os.Stdout)
  defer writer.Flush()

  numFiles := len(os.Args)-1
  if numFiles == 0 {
    writer.WriteString("No files specified.\n")
    return
  }
  for argIx := 1; argIx <= numFiles; argIx++ {
    // Parse the top-level template. The resulting source code goes into
    //  the global variable sections.
    templatePath := os.Args[argIx]

    // We parse the template to obtain code sections and static sections.
    error := parse(templatePath)
    if error != nil {
      writer.WriteString(fmt.Sprintf("Template parsing error: %s\n", error))
      continue
    }

    // Concatenate only the code sections.
    output := bytes.Buffer{}
    for _, section := range sections {
      if section.Kind == CodeSection {
        fmt.Fprintf(&output, section.Text)
      }
    }
    fileSet := token.NewFileSet()
    fileNode, error := parser.ParseFile(fileSet, "output", output.Bytes(),
        parser.ParseComments)
    if error != nil {
      writer.Write(output.Bytes())
      writer.WriteString(fmt.Sprintf(
          "\n---\nError parsing code sections: %s\n", error))
      continue
    }

    //  Has the package fmt been imported? Is the name "fmt" available?
    fmtIsImported := false
    var fmtImportName string       // use this if fmt has already been imported
    seenName := map[string]bool{}  // consult this if we have to import fmt

    for _, importSpec := range fileNode.Imports {
      importPath, _ := strconv.Unquote(importSpec.Path.Value)
      var importName string
      if importSpec.Name == nil {
        importName = path.Base(importPath)
      } else {
        importName = importSpec.Name.Name
      }
      seenName[importName] = true
      if !fmtIsImported && importPath == "fmt" && importName != "_" {
        fmtIsImported = true
        fmtImportName = importName
      }
    }

    var importAs, printPrefix string  // NB: these are "" by default
    if fmtIsImported {
      if fmtImportName != "." {  // no prefix is needed with a dot import
        printPrefix = fmtImportName+"."
      }
    } else {
      if !seenName["fmt"] {
        importAs = "fmt"
      } else {
        for i := 0; ; i++ {
          importAs = fmt.Sprintf("fmt_%d", i)
          _, found := seenName[importAs]
          if !found {
            break
          }
        }
      }
      printPrefix = importAs+"."
    }

    // Concatenate the code sections and static sections.
    output.Reset()
    for _, section := range sections {
      if section.Kind == CodeSection {
        fmt.Fprintf(&output, section.Text)
      } else {
        s := fmt.Sprintf(";%sPrint(%s);", printPrefix, section.Text)
        fmt.Fprintf(&output, s)
      }
    }
    // Have Go parse the entire template output.
    fileSet = token.NewFileSet()
    fileNode, error = parser.ParseFile(fileSet, "output", output.Bytes(),
        parser.ParseComments)
    if error != nil {
      writer.Write(output.Bytes())
      writer.WriteString(fmt.Sprintf(
          "\n---\nError parsing entire template output: %s\n", error))
      continue
    }
    // Finally, inject an import statement if necessary.
    if !fmtIsImported {
      if importAs == "fmt" {
        astutil.AddImport(fileSet, fileNode, "fmt")
      } else {
        astutil.AddNamedImport(fileSet, fileNode, importAs, "fmt")
      }
    }

    // Print with a custom configuration: soft tabs of two spaces each.
    config := printer.Config{ Mode: printer.UseSpaces, Tabwidth: 2 }
    (&config).Fprint(writer, fileSet, fileNode)
  }
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: MIT

package locale

import "golang.org/x/text/language"

var zhHans = map[string]string{
	// 与 flag 包相关的处理
	FlagUsage: `%s 是一个 RESTful API 文档生成工具

用法：
apidoc [options] [path]

参数：
%s

源代码采用 MIT 开源许可证，发布于 %s
详细信息可访问官网 %s
`,
	FlagHUsage:            "显示帮助信息",
	FlagVUsage:            "显示版本信息",
	FlagLUsage:            "显示所有支持的语言",
	FlagDUsage:            "根据目录下的内容生成配置文件",
	FlagVersionBuildWith:  "%s %s build with %s",
	FlagVersionCommitHash: "commit hash %s",

	VersionInCompatible: "当前程序与配置文件中指定的版本号不兼容",
	Complete:            "完成！文档保存在：%s，总用时：%v",
	ConfigWriteSuccess:  "配置内容成功写入 %s",

	// 错误信息，可能在地方用到
	ErrRequired:              "不能为空",
	ErrMustEmpty:             "只能为空",
	ErrInvalidFormat:         "格式不正确",
	ErrDirNotExists:          "目录不存在",
	ErrUnsupportedInputLang:  "不支持的输入语言：%s",
	ErrNotFoundEndFlag:       "找不到结束符号",
	ErrNotFoundSupportedLang: "该目录下没有支持的语言文件",
	ErrUnknownTag:            "不认识的标签",
	ErrDuplicateTag:          "重复的标签",
	ErrUnsupportedEncoding:   "不支持的编码方式",
	ErrDirIsEmpty:            "目录下没有需要解析的文件",
	ErrInvalidValue:          "无效的值",
	ErrInvalidOpenapi:        "openapi 内容错误：字段：%s；错误内容：%s",
	ErrPathNotMatchParams:    "地址参数不匹配",
	ErrDuplicateValue:        "重复的值",
	ErrMessage:               "%s 位于 %s",

	// logs
	InfoPrefix:    "[信息] ",
	WarnPrefix:    "[警告] ",
	ErrorPrefix:   "[错误] ",
	SuccessPrefix: "[成功] ",
}

func init() {
	addLocale(language.MustParse("zh-Hans"), zhHans)
}
<commit_msg>解决 zh-cn 无法匹配到  zh-hans 的错误<commit_after>// SPDX-License-Identifier: MIT

package locale

import "golang.org/x/text/language"

var zhHans = map[string]string{
	// 与 flag 包相关的处理
	FlagUsage: `%s 是一个 RESTful API 文档生成工具

用法：
apidoc [options] [path]

参数：
%s

源代码采用 MIT 开源许可证，发布于 %s
详细信息可访问官网 %s
`,
	FlagHUsage:            "显示帮助信息",
	FlagVUsage:            "显示版本信息",
	FlagLUsage:            "显示所有支持的语言",
	FlagDUsage:            "根据目录下的内容生成配置文件",
	FlagVersionBuildWith:  "%s %s build with %s",
	FlagVersionCommitHash: "commit hash %s",

	VersionInCompatible: "当前程序与配置文件中指定的版本号不兼容",
	Complete:            "完成！文档保存在：%s，总用时：%v",
	ConfigWriteSuccess:  "配置内容成功写入 %s",

	// 错误信息，可能在地方用到
	ErrRequired:              "不能为空",
	ErrMustEmpty:             "只能为空",
	ErrInvalidFormat:         "格式不正确",
	ErrDirNotExists:          "目录不存在",
	ErrUnsupportedInputLang:  "不支持的输入语言：%s",
	ErrNotFoundEndFlag:       "找不到结束符号",
	ErrNotFoundSupportedLang: "该目录下没有支持的语言文件",
	ErrUnknownTag:            "不认识的标签",
	ErrDuplicateTag:          "重复的标签",
	ErrUnsupportedEncoding:   "不支持的编码方式",
	ErrDirIsEmpty:            "目录下没有需要解析的文件",
	ErrInvalidValue:          "无效的值",
	ErrInvalidOpenapi:        "openapi 内容错误：字段：%s；错误内容：%s",
	ErrPathNotMatchParams:    "地址参数不匹配",
	ErrDuplicateValue:        "重复的值",
	ErrMessage:               "%s 位于 %s",

	// logs
	InfoPrefix:    "[信息] ",
	WarnPrefix:    "[警告] ",
	ErrorPrefix:   "[错误] ",
	SuccessPrefix: "[成功] ",
}

func init() {
	addLocale(language.MustParse("zh-Hans"), zhHans)

	// 大部分的系统都采用 zh-cn 作为语言标记，
	// 但是 golang.org/x/text 现在不能将 zh-cn 自动转换成 zh-hans
	addLocale(language.MustParse("zh-cn"), zhHans)
}
<|endoftext|>