<commit_before>// +build linux

package host

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/shirou/gopsutil/internal/common"
)

type LSB struct {
	ID          string
	Release     string
	Codename    string
	Description string
}

// from utmp.h
const USER_PROCESS = 7

func Info() (*InfoStat, error) {
	ret := &InfoStat{
		OS: runtime.GOOS,
	}

	hostname, err := os.Hostname()
	if err == nil {
		ret.Hostname = hostname
	}

	platform, family, version, err := PlatformInformation()
	if err == nil {
		ret.Platform = platform
		ret.PlatformFamily = family
		ret.PlatformVersion = version
	}
	system, role, err := Virtualization()
	if err == nil {
		ret.VirtualizationSystem = system
		ret.VirtualizationRole = role
	}
	boot, err := BootTime()
	if err == nil {
		ret.BootTime = boot
		ret.Uptime = uptime(boot)
	}

	return ret, nil
}

// BootTime returns the system boot time expressed in seconds since the epoch.
func BootTime() (uint64, error) {
	filename := common.HostProc("stat")
	lines, err := common.ReadLines(filename)
	if err != nil {
		return 0, err
	}
	for _, line := range lines {
		if strings.HasPrefix(line, "btime") {
			f := strings.Fields(line)
			if len(f) != 2 {
				return 0, fmt.Errorf("wrong btime format")
			}
			b, err := strconv.ParseInt(f[1], 10, 64)
			if err != nil {
				return 0, err
			}
			return uint64(b), nil
		}
	}

	return 0, fmt.Errorf("could not find btime")
}

func uptime(boot uint64) uint64 {
	return uint64(time.Now().Unix()) - boot
}

func Uptime() (uint64, error) {
	boot, err := BootTime()
	if err != nil {
		return 0, err
	}
	return uptime(boot), nil
}

func Users() ([]UserStat, error) {
	utmpfile := "/var/run/utmp"

	file, err := os.Open(utmpfile)
	if err != nil {
		return nil, err
	}

	buf, err := ioutil.ReadAll(file)
	if err != nil {
		return nil, err
	}

	count := len(buf) / sizeOfUtmp

	ret := make([]UserStat, 0, count)

	for i := 0; i < count; i++ {
		b := buf[i*sizeOfUtmp : (i+1)*sizeOfUtmp]

		var u utmp
		br := bytes.NewReader(b)
		err := binary.Read(br, binary.LittleEndian, &u)
		if err != nil {
			continue
		}
		if u.Type != USER_PROCESS {
			continue
		}
		user := UserStat{
			User:     common.IntToString(u.User[:]),
			Terminal: common.IntToString(u.Line[:]),
			Host:     common.IntToString(u.Host[:]),
			Started:  int(u.Tv.Sec),
		}
		ret = append(ret, user)
	}

	return ret, nil

}

func getLSB() (*LSB, error) {
	ret := &LSB{}
	if common.PathExists(common.HostEtc("lsb-release")) {
		contents, err := common.ReadLines(common.HostEtc("lsb-release"))
		if err != nil {
			return ret, err // return empty
		}
		for _, line := range contents {
			field := strings.Split(line, "=")
			if len(field) < 2 {
				continue
			}
			switch field[0] {
			case "DISTRIB_ID":
				ret.ID = field[1]
			case "DISTRIB_RELEASE":
				ret.Release = field[1]
			case "DISTRIB_CODENAME":
				ret.Codename = field[1]
			case "DISTRIB_DESCRIPTION":
				ret.Description = field[1]
			}
		}
	} else if common.PathExists("/usr/bin/lsb_release") {
		lsb_release, err := exec.LookPath("/usr/bin/lsb_release")
		if err != nil {
			return ret, err
		}
		out, err := invoke.Command(lsb_release)
		if err != nil {
			return ret, err
		}
		for _, line := range strings.Split(string(out), "\n") {
			field := strings.Split(line, ":")
			if len(field) < 2 {
				continue
			}
			switch field[0] {
			case "Distributor ID":
				ret.ID = field[1]
			case "Release":
				ret.Release = field[1]
			case "Codename":
				ret.Codename = field[1]
			case "Description":
				ret.Description = field[1]
			}
		}

	}

	return ret, nil
}

func PlatformInformation() (platform string, family string, version string, err error) {

	lsb, err := getLSB()
	if err != nil {
		lsb = &LSB{}
	}

	if common.PathExists(common.HostEtc("oracle-release")) {
		platform = "oracle"
		contents, err := common.ReadLines(common.HostEtc("oracle-release"))
		if err == nil {
			version = getRedhatishVersion(contents)
		}

	} else if common.PathExists(common.HostEtc("enterprise-release")) {
		platform = "oracle"
		contents, err := common.ReadLines(common.HostEtc("enterprise-release"))
		if err == nil {
			version = getRedhatishVersion(contents)
		}
	} else if common.PathExists(common.HostEtc("debian_version")) {
		if lsb.ID == "Ubuntu" {
			platform = "ubuntu"
			version = lsb.Release
		} else if lsb.ID == "LinuxMint" {
			platform = "linuxmint"
			version = lsb.Release
		} else {
			if common.PathExists("/usr/bin/raspi-config") {
				platform = "raspbian"
			} else {
				platform = "debian"
			}
			contents, err := common.ReadLines(common.HostEtc("debian_version"))
			if err == nil {
				version = contents[0]
			}
		}
	} else if common.PathExists(common.HostEtc("redhat-release")) {
		contents, err := common.ReadLines(common.HostEtc("redhat-release"))
		if err == nil {
			version = getRedhatishVersion(contents)
			platform = getRedhatishPlatform(contents)
		}
	} else if common.PathExists(common.HostEtc("system-release")) {
		contents, err := common.ReadLines(common.HostEtc("system-release"))
		if err == nil {
			version = getRedhatishVersion(contents)
			platform = getRedhatishPlatform(contents)
		}
	} else if common.PathExists(common.HostEtc("gentoo-release")) {
		platform = "gentoo"
		contents, err := common.ReadLines(common.HostEtc("gentoo-release"))
		if err == nil {
			version = getRedhatishVersion(contents)
		}
	} else if common.PathExists(common.HostEtc("SuSE-release")) {
		contents, err := common.ReadLines(common.HostEtc("SuSE-release"))
		if err == nil {
			version = getSuseVersion(contents)
			platform = getSusePlatform(contents)
		}
		// TODO: slackware detecion
	} else if common.PathExists(common.HostEtc("arch-release")) {
		platform = "arch"
		version = lsb.Release
	} else if lsb.ID == "RedHat" {
		platform = "redhat"
		version = lsb.Release
	} else if lsb.ID == "Amazon" {
		platform = "amazon"
		version = lsb.Release
	} else if lsb.ID == "ScientificSL" {
		platform = "scientific"
		version = lsb.Release
	} else if lsb.ID == "XenServer" {
		platform = "xenserver"
		version = lsb.Release
	} else if lsb.ID != "" {
		platform = strings.ToLower(lsb.ID)
		version = lsb.Release
	}

	switch platform {
	case "debian", "ubuntu", "linuxmint", "raspbian":
		family = "debian"
	case "fedora":
		family = "fedora"
	case "oracle", "centos", "redhat", "scientific", "enterpriseenterprise", "amazon", "xenserver", "cloudlinux", "ibm_powerkvm":
		family = "rhel"
	case "suse", "opensuse":
		family = "suse"
	case "gentoo":
		family = "gentoo"
	case "slackware":
		family = "slackware"
	case "arch":
		family = "arch"
	case "exherbo":
		family = "exherbo"
	}

	return platform, family, version, nil

}

func getRedhatishVersion(contents []string) string {
	c := strings.ToLower(strings.Join(contents, ""))

	if strings.Contains(c, "rawhide") {
		return "rawhide"
	}
	if matches := regexp.MustCompile(`release (\d[\d.]*)`).FindStringSubmatch(c); matches != nil {
		return matches[1]
	}
	return ""
}

func getRedhatishPlatform(contents []string) string {
	c := strings.ToLower(strings.Join(contents, ""))

	if strings.Contains(c, "red hat") {
		return "redhat"
	}
	f := strings.Split(c, " ")

	return f[0]
}

func getSuseVersion(contents []string) string {
	version := ""
	for _, line := range contents {
		if matches := regexp.MustCompile(`VERSION = ([\d.]+)`).FindStringSubmatch(line); matches != nil {
			version = matches[1]
		} else if matches := regexp.MustCompile(`PATCHLEVEL = ([\d]+)`).FindStringSubmatch(line); matches != nil {
			version = version + "." + matches[1]
		}
	}
	return version
}

func getSusePlatform(contents []string) string {
	c := strings.ToLower(strings.Join(contents, ""))
	if strings.Contains(c, "opensuse") {
		return "opensuse"
	}
	return "suse"
}

func Virtualization() (string, string, error) {
	var system string
	var role string

	filename := common.HostProc("xen")
	if common.PathExists(filename) {
		system = "xen"
		role = "guest" // assume guest

		if common.PathExists(filename + "/capabilities") {
			contents, err := common.ReadLines(filename + "/capabilities")
			if err == nil {
				if common.StringsHas(contents, "control_d") {
					role = "host"
				}
			}
		}
	}

	filename = common.HostProc("modules")
	if common.PathExists(filename) {
		contents, err := common.ReadLines(filename)
		if err == nil {
			if common.StringsContains(contents, "kvm") {
				system = "kvm"
				role = "host"
			} else if common.StringsContains(contents, "vboxdrv") {
				system = "vbox"
				role = "host"
			} else if common.StringsContains(contents, "vboxguest") {
				system = "vbox"
				role = "guest"
			}
		}
	}

	filename = common.HostProc("cpuinfo")
	if common.PathExists(filename) {
		contents, err := common.ReadLines(filename)
		if err == nil {
			if common.StringsHas(contents, "QEMU Virtual CPU") ||
				common.StringsHas(contents, "Common KVM processor") ||
				common.StringsHas(contents, "Common 32-bit KVM processor") {
				system = "kvm"
				role = "guest"
			}
		}
	}

	filename = common.HostProc()
	if common.PathExists(filename + "/bc/0") {
		system = "openvz"
		role = "host"
	} else if common.PathExists(filename + "/vz") {
		system = "openvz"
		role = "guest"
	}

	// not use dmidecode because it requires root
	if common.PathExists(filename + "/self/status") {
		contents, err := common.ReadLines(filename + "/self/status")
		if err == nil {

			if common.StringsHas(contents, "s_context:") ||
				common.StringsHas(contents, "VxID:") {
				system = "linux-vserver"
			}
			// TODO: guest or host
		}
	}

	if common.PathExists(filename + "/self/cgroup") {
		contents, err := common.ReadLines(filename + "/self/cgroup")
		if err == nil {
			if common.StringsHas(contents, "lxc") ||
				common.StringsHas(contents, "docker") {
				system = "lxc"
				role = "guest"
			} else if common.PathExists("/usr/bin/lxc-version") { // TODO: which
				system = "lxc"
				role = "host"
			}
		}
	}

	return system, role, nil
}
<commit_msg>[host]linux: add alpine linux to detect host.Info().<commit_after>// +build linux

package host

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"time"

	common "github.com/shirou/gopsutil/internal/common"
)

type LSB struct {
	ID          string
	Release     string
	Codename    string
	Description string
}

// from utmp.h
const USER_PROCESS = 7

func Info() (*InfoStat, error) {
	ret := &InfoStat{
		OS: runtime.GOOS,
	}

	hostname, err := os.Hostname()
	if err == nil {
		ret.Hostname = hostname
	}

	platform, family, version, err := PlatformInformation()
	if err == nil {
		ret.Platform = platform
		ret.PlatformFamily = family
		ret.PlatformVersion = version
	}
	system, role, err := Virtualization()
	if err == nil {
		ret.VirtualizationSystem = system
		ret.VirtualizationRole = role
	}
	boot, err := BootTime()
	if err == nil {
		ret.BootTime = boot
		ret.Uptime = uptime(boot)
	}

	return ret, nil
}

// BootTime returns the system boot time expressed in seconds since the epoch.
func BootTime() (uint64, error) {
	filename := common.HostProc("stat")
	lines, err := common.ReadLines(filename)
	if err != nil {
		return 0, err
	}
	for _, line := range lines {
		if strings.HasPrefix(line, "btime") {
			f := strings.Fields(line)
			if len(f) != 2 {
				return 0, fmt.Errorf("wrong btime format")
			}
			b, err := strconv.ParseInt(f[1], 10, 64)
			if err != nil {
				return 0, err
			}
			return uint64(b), nil
		}
	}

	return 0, fmt.Errorf("could not find btime")
}

func uptime(boot uint64) uint64 {
	return uint64(time.Now().Unix()) - boot
}

func Uptime() (uint64, error) {
	boot, err := BootTime()
	if err != nil {
		return 0, err
	}
	return uptime(boot), nil
}

func Users() ([]UserStat, error) {
	utmpfile := "/var/run/utmp"

	file, err := os.Open(utmpfile)
	if err != nil {
		return nil, err
	}

	buf, err := ioutil.ReadAll(file)
	if err != nil {
		return nil, err
	}

	count := len(buf) / sizeOfUtmp

	ret := make([]UserStat, 0, count)

	for i := 0; i < count; i++ {
		b := buf[i*sizeOfUtmp : (i+1)*sizeOfUtmp]

		var u utmp
		br := bytes.NewReader(b)
		err := binary.Read(br, binary.LittleEndian, &u)
		if err != nil {
			continue
		}
		if u.Type != USER_PROCESS {
			continue
		}
		user := UserStat{
			User:     common.IntToString(u.User[:]),
			Terminal: common.IntToString(u.Line[:]),
			Host:     common.IntToString(u.Host[:]),
			Started:  int(u.Tv.Sec),
		}
		ret = append(ret, user)
	}

	return ret, nil

}

func getLSB() (*LSB, error) {
	ret := &LSB{}
	if common.PathExists(common.HostEtc("lsb-release")) {
		contents, err := common.ReadLines(common.HostEtc("lsb-release"))
		if err != nil {
			return ret, err // return empty
		}
		for _, line := range contents {
			field := strings.Split(line, "=")
			if len(field) < 2 {
				continue
			}
			switch field[0] {
			case "DISTRIB_ID":
				ret.ID = field[1]
			case "DISTRIB_RELEASE":
				ret.Release = field[1]
			case "DISTRIB_CODENAME":
				ret.Codename = field[1]
			case "DISTRIB_DESCRIPTION":
				ret.Description = field[1]
			}
		}
	} else if common.PathExists("/usr/bin/lsb_release") {
		lsb_release, err := exec.LookPath("/usr/bin/lsb_release")
		if err != nil {
			return ret, err
		}
		out, err := invoke.Command(lsb_release)
		if err != nil {
			return ret, err
		}
		for _, line := range strings.Split(string(out), "\n") {
			field := strings.Split(line, ":")
			if len(field) < 2 {
				continue
			}
			switch field[0] {
			case "Distributor ID":
				ret.ID = field[1]
			case "Release":
				ret.Release = field[1]
			case "Codename":
				ret.Codename = field[1]
			case "Description":
				ret.Description = field[1]
			}
		}

	}

	return ret, nil
}

func PlatformInformation() (platform string, family string, version string, err error) {

	lsb, err := getLSB()
	if err != nil {
		lsb = &LSB{}
	}

	if common.PathExists(common.HostEtc("oracle-release")) {
		platform = "oracle"
		contents, err := common.ReadLines(common.HostEtc("oracle-release"))
		if err == nil {
			version = getRedhatishVersion(contents)
		}

	} else if common.PathExists(common.HostEtc("enterprise-release")) {
		platform = "oracle"
		contents, err := common.ReadLines(common.HostEtc("enterprise-release"))
		if err == nil {
			version = getRedhatishVersion(contents)
		}
	} else if common.PathExists(common.HostEtc("debian_version")) {
		if lsb.ID == "Ubuntu" {
			platform = "ubuntu"
			version = lsb.Release
		} else if lsb.ID == "LinuxMint" {
			platform = "linuxmint"
			version = lsb.Release
		} else {
			if common.PathExists("/usr/bin/raspi-config") {
				platform = "raspbian"
			} else {
				platform = "debian"
			}
			contents, err := common.ReadLines(common.HostEtc("debian_version"))
			if err == nil {
				version = contents[0]
			}
		}
	} else if common.PathExists(common.HostEtc("redhat-release")) {
		contents, err := common.ReadLines(common.HostEtc("redhat-release"))
		if err == nil {
			version = getRedhatishVersion(contents)
			platform = getRedhatishPlatform(contents)
		}
	} else if common.PathExists(common.HostEtc("system-release")) {
		contents, err := common.ReadLines(common.HostEtc("system-release"))
		if err == nil {
			version = getRedhatishVersion(contents)
			platform = getRedhatishPlatform(contents)
		}
	} else if common.PathExists(common.HostEtc("gentoo-release")) {
		platform = "gentoo"
		contents, err := common.ReadLines(common.HostEtc("gentoo-release"))
		if err == nil {
			version = getRedhatishVersion(contents)
		}
	} else if common.PathExists(common.HostEtc("SuSE-release")) {
		contents, err := common.ReadLines(common.HostEtc("SuSE-release"))
		if err == nil {
			version = getSuseVersion(contents)
			platform = getSusePlatform(contents)
		}
		// TODO: slackware detecion
	} else if common.PathExists(common.HostEtc("arch-release")) {
		platform = "arch"
		version = lsb.Release
	} else if common.PathExists(common.HostEtc("alpine-release")) {
		platform = "alpine"
		contents, err := common.ReadLines(common.HostEtc("alpine-release"))
		if err == nil && len(contents) > 0 {
			version = contents[0]
		}
	} else if lsb.ID == "RedHat" {
		platform = "redhat"
		version = lsb.Release
	} else if lsb.ID == "Amazon" {
		platform = "amazon"
		version = lsb.Release
	} else if lsb.ID == "ScientificSL" {
		platform = "scientific"
		version = lsb.Release
	} else if lsb.ID == "XenServer" {
		platform = "xenserver"
		version = lsb.Release
	} else if lsb.ID != "" {
		platform = strings.ToLower(lsb.ID)
		version = lsb.Release
	}

	switch platform {
	case "debian", "ubuntu", "linuxmint", "raspbian":
		family = "debian"
	case "fedora":
		family = "fedora"
	case "oracle", "centos", "redhat", "scientific", "enterpriseenterprise", "amazon", "xenserver", "cloudlinux", "ibm_powerkvm":
		family = "rhel"
	case "suse", "opensuse":
		family = "suse"
	case "gentoo":
		family = "gentoo"
	case "slackware":
		family = "slackware"
	case "arch":
		family = "arch"
	case "exherbo":
		family = "exherbo"
	case "alpine":
		family = "alpine"
	}

	return platform, family, version, nil

}

func getRedhatishVersion(contents []string) string {
	c := strings.ToLower(strings.Join(contents, ""))

	if strings.Contains(c, "rawhide") {
		return "rawhide"
	}
	if matches := regexp.MustCompile(`release (\d[\d.]*)`).FindStringSubmatch(c); matches != nil {
		return matches[1]
	}
	return ""
}

func getRedhatishPlatform(contents []string) string {
	c := strings.ToLower(strings.Join(contents, ""))

	if strings.Contains(c, "red hat") {
		return "redhat"
	}
	f := strings.Split(c, " ")

	return f[0]
}

func getSuseVersion(contents []string) string {
	version := ""
	for _, line := range contents {
		if matches := regexp.MustCompile(`VERSION = ([\d.]+)`).FindStringSubmatch(line); matches != nil {
			version = matches[1]
		} else if matches := regexp.MustCompile(`PATCHLEVEL = ([\d]+)`).FindStringSubmatch(line); matches != nil {
			version = version + "." + matches[1]
		}
	}
	return version
}

func getSusePlatform(contents []string) string {
	c := strings.ToLower(strings.Join(contents, ""))
	if strings.Contains(c, "opensuse") {
		return "opensuse"
	}
	return "suse"
}

func Virtualization() (string, string, error) {
	var system string
	var role string

	filename := common.HostProc("xen")
	if common.PathExists(filename) {
		system = "xen"
		role = "guest" // assume guest

		if common.PathExists(filename + "/capabilities") {
			contents, err := common.ReadLines(filename + "/capabilities")
			if err == nil {
				if common.StringsHas(contents, "control_d") {
					role = "host"
				}
			}
		}
	}

	filename = common.HostProc("modules")
	if common.PathExists(filename) {
		contents, err := common.ReadLines(filename)
		if err == nil {
			if common.StringsContains(contents, "kvm") {
				system = "kvm"
				role = "host"
			} else if common.StringsContains(contents, "vboxdrv") {
				system = "vbox"
				role = "host"
			} else if common.StringsContains(contents, "vboxguest") {
				system = "vbox"
				role = "guest"
			}
		}
	}

	filename = common.HostProc("cpuinfo")
	if common.PathExists(filename) {
		contents, err := common.ReadLines(filename)
		if err == nil {
			if common.StringsHas(contents, "QEMU Virtual CPU") ||
				common.StringsHas(contents, "Common KVM processor") ||
				common.StringsHas(contents, "Common 32-bit KVM processor") {
				system = "kvm"
				role = "guest"
			}
		}
	}

	filename = common.HostProc()
	if common.PathExists(filename + "/bc/0") {
		system = "openvz"
		role = "host"
	} else if common.PathExists(filename + "/vz") {
		system = "openvz"
		role = "guest"
	}

	// not use dmidecode because it requires root
	if common.PathExists(filename + "/self/status") {
		contents, err := common.ReadLines(filename + "/self/status")
		if err == nil {

			if common.StringsHas(contents, "s_context:") ||
				common.StringsHas(contents, "VxID:") {
				system = "linux-vserver"
			}
			// TODO: guest or host
		}
	}

	if common.PathExists(filename + "/self/cgroup") {
		contents, err := common.ReadLines(filename + "/self/cgroup")
		if err == nil {
			if common.StringsHas(contents, "lxc") ||
				common.StringsHas(contents, "docker") {
				system = "lxc"
				role = "guest"
			} else if common.PathExists("/usr/bin/lxc-version") { // TODO: which
				system = "lxc"
				role = "host"
			}
		}
	}

	return system, role, nil
}
<|endoftext|>


<commit_before>package cornell

import (
	"bytes"
	"database/sql"
	"encoding/json"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/MyHomeworkSpace/api-server/data"
	"github.com/MyHomeworkSpace/api-server/schools"
	"github.com/MyHomeworkSpace/api-server/util"
	"github.com/PuerkitoBio/goquery"
)

var apiKeyRegexp = regexp.MustCompile(`"key":"(?P<Key>[^"]*)`)
var nameRegexp = regexp.MustCompile(`login.*"name":"([^"]*)`) //kinda hacky, I should fix this

func (s *school) Enroll(tx *sql.Tx, user *data.User, params map[string]interface{}) (map[string]interface{}, error) {
	netidRaw, ok := params["username"]
	passwordRaw, ok2 := params["password"]

	if !ok || !ok2 {
		return nil, data.SchoolError{Code: "missing_params"}
	}

	netID, ok := netidRaw.(string)
	password, ok2 := passwordRaw.(string)

	if !ok || !ok2 || netID == "" || password == "" {
		return nil, data.SchoolError{Code: "invalid_params"}
	}

	cookieJar, _ := cookiejar.New(nil)
	c := &http.Client{
		Jar: cookieJar,
	}

	term := GetCurrentTerm().Code

	resp, err := c.Get("https://classes.cornell.edu/sascuwalogin/login/redirect?redirectUri=https%3A%2F%2Fclasses.cornell.edu%2Fscheduler%2Froster%2F" + term)
	if err != nil {
		return nil, err
	}

	doc, err := goquery.NewDocumentFromReader(resp.Body)
	if err != nil {
		return nil, err
	}

	loginForm, exists := doc.Find("form").First().Attr("action")
	if !exists {
		return nil, data.SchoolError{Code: "internal_server_error"}
	}

	values := url.Values{
		"j_username":       {netID},
		"j_password":       {password},
		"_eventId_proceed": {"Login"},
	}

	loginResp, err := c.PostForm(resp.Request.URL.Scheme+"://"+resp.Request.URL.Host+loginForm, values)
	if err != nil {
		return nil, err
	}

	confirmationDoc, err := goquery.NewDocumentFromReader(loginResp.Body)
	confirmationDocString, _ := confirmationDoc.Html()

	if strings.Contains(confirmationDocString, "Unable to log in") {
		return nil, data.SchoolError{Code: "creds_incorrect"}
	}

	// otherwise we're logged in! Unfortunately, CUWebLogin has this intermediate screen that says Your login credentials are being transmitted to the website via POST. This only works if you use javascript, so we need to manually transmit the token

	classesURL, exists := confirmationDoc.Find("form").First().Attr("action")
	if !exists {
		return nil, data.SchoolError{Code: "internal_server_error"}
	}

	relayState, exists := confirmationDoc.Find("input[name=RelayState]").First().Attr("value")
	if !exists {
		return nil, data.SchoolError{Code: "internal_server_error"}
	}

	samlResponse, exists := confirmationDoc.Find("input[name=SAMLResponse]").First().Attr("value")
	if !exists {
		return nil, data.SchoolError{Code: "internal_server_error"}
	}

	schedulerResp, err := c.PostForm(classesURL, url.Values{"RelayState": {relayState}, "SAMLResponse": {samlResponse}})
	if err != nil {
		return nil, err
	}

	schedulerRespBuf := new(bytes.Buffer)
	schedulerRespBuf.ReadFrom(schedulerResp.Body)
	schedulerRespContent := schedulerRespBuf.String()

	APIKey := apiKeyRegexp.FindStringSubmatch(schedulerRespContent)[1]
	name := nameRegexp.FindStringSubmatch(schedulerRespContent)[1]

	scheduleReq, err := http.NewRequest("GET", "https://classes.cornell.edu/api/3.0/scheduler/current-enrollment", nil)
	if err != nil {
		return nil, err
	}

	scheduleReq.Header.Add("Authorization", "ClassRoster "+APIKey)

	scheduleResp, err := c.Do(scheduleReq)
	if err != nil {
		return nil, err
	}

	scheduleRespBuf := new(bytes.Buffer)
	scheduleRespBuf.ReadFrom(scheduleResp.Body)
	classesBytes := scheduleRespBuf.Bytes()

	courses := []classItem{}

	err = json.Unmarshal(classesBytes, &courses)
	if err != nil {
		return nil, err
	}

	cpairs := coursePairs{}

	for _, class := range courses {
		cpairs.CoursePairs = append(cpairs.CoursePairs, strconv.Itoa(class.CourseID)+","+strconv.Itoa(class.CourseOfferNumber))
	}

	coursePairsJSON, err := json.Marshal(cpairs)

	if err != nil {
		return nil, err
	}

	courseDetailReq, err := http.NewRequest("POST", "https://classes.cornell.edu/api/3.0/scheduler/course-detail", bytes.NewBuffer(coursePairsJSON))
	if err != nil {
		return nil, err
	}
	courseDetailReq.Header.Set("Content-Type", "application/json")
	courseDetailReq.Header.Add("Authorization", "ClassRoster "+APIKey)

	courseDetailsResp, err := c.Do(courseDetailReq)
	if err != nil {
		return nil, err
	}

	courseDetailBuf := new(bytes.Buffer)
	courseDetailBuf.ReadFrom(courseDetailsResp.Body)

	courseDetails := []course{}
	err = json.Unmarshal(courseDetailBuf.Bytes(), &courseDetails)
	if err != nil {
		return nil, err
	}

	// ok, now we have all the course details and the classes that the student is taking. Now we need to put it all in the db.

	for _, course := range courses {
		_, err = tx.Exec("INSERT INTO cornell_courses (userId, subject, catalogNum, title, units, rosterId) VALUES (?, ?, ?, ?, ?, ?)", user.ID, course.Subject, course.CatalogNumber, course.Title, course.Units, course.CourseID)
		if err != nil {
			return nil, err
		}

		for _, details := range courseDetails {
			if details.CourseID == course.CourseID {
				// we matched the details with the course!
				for _, eg := range details.EnrollGroups {
					for _, section := range eg.ClassSections {
						if util.IntSliceContains(course.ClassNumbers, section.ClassNum) {
							for _, meeting := range section.Meetings {
								if meeting.StartTime == "" || meeting.EndTime == "" {
									// this class is missing a time
									// not super clear what these are, skip it
									continue
								}

								startDate, err := time.Parse("01/02/2006", meeting.StartDate)
								if err != nil {
									return nil, err
								}

								endDate, err := time.Parse("01/02/2006", meeting.EndDate)
								if err != nil {
									return nil, err
								}

								startTime, err := time.Parse("3:04PM", meeting.StartTime)
								if err != nil {
									return nil, err
								}

								endTime, err := time.Parse("3:04PM", meeting.EndTime)
								if err != nil {
									return nil, err
								}

								// kinda hacky but it works
								startTime = time.Date(1970, 01, 01, startTime.Hour(), startTime.Minute(), 0, 0, time.UTC)
								endTime = time.Date(1970, 01, 01, endTime.Hour(), endTime.Minute(), 0, 0, time.UTC)

								satMeeting := false
								sunMeeting := false

								if strings.Contains("SSu", meeting.Pattern) {
									satMeeting = true
									sunMeeting = true
								} else if strings.Contains("Su", meeting.Pattern) {
									sunMeeting = true
								} else if strings.Contains("S", meeting.Pattern) {
									satMeeting = true
								}

								_, err = tx.Exec(`INSERT INTO cornell_events (
		title,
		userId,
		subject,
		catalogNum,
		classNum,
		component,
		componentLong,
		section,
		campus,
		campusLong,
		location,
		locationLong,
		startDate,
		endDate,
		startTime,
		endTime,
		monday,
		tuesday,
		wednesday,
		thursday,
		friday,
		saturday,
		sunday,
		facility,
		facilityLong,
		building
	)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
									course.Title,
									user.ID,
									course.Subject,
									course.CatalogNumber,
									section.ClassNum,
									section.Component,
									section.ComponentLong,
									section.Section,
									section.Campus,
									section.CampusDesc,
									section.Location,
									section.LocationDesc,
									startDate.Format("2006-01-02"),
									endDate.Format("2006-01-02"),
									startTime.Unix(),
									endTime.Unix(),
									strings.Contains(meeting.Pattern, "M"),
									strings.Contains(meeting.Pattern, "T"),
									strings.Contains(meeting.Pattern, "W"),
									strings.Contains(meeting.Pattern, "R"),
									strings.Contains(meeting.Pattern, "F"),
									satMeeting,
									sunMeeting,
									meeting.FacilityDescShort,
									meeting.FacilityDesc,
									meeting.BuildingDesc,
								)

								if err != nil {
									return nil, err
								}
							}
						}
					}
				}
				break
			}
		}
	}

	return map[string]interface{}{
		"netid":  strings.ToLower(netID),
		"name":   name,
		"status": schools.ImportStatusOK,
	}, nil
}

func (s *school) NeedsUpdate(db *sql.DB) (bool, error) {
	return (s.importStatus != schools.ImportStatusOK), nil
}

func (s *school) Unenroll(tx *sql.Tx, user *data.User) error {
	return clearUserData(tx, user)
}

func clearUserData(tx *sql.Tx, user *data.User) error {
	// clear away anything that is in the db
	_, err := tx.Exec("DELETE FROM cornell_courses WHERE userId = ?", user.ID)
	if err != nil {
		return err
	}

	_, err = tx.Exec("DELETE FROM cornell_events WHERE userId = ?", user.ID)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>cornell: improve error handling<commit_after>package cornell

import (
	"bytes"
	"database/sql"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/MyHomeworkSpace/api-server/data"
	"github.com/MyHomeworkSpace/api-server/schools"
	"github.com/MyHomeworkSpace/api-server/util"
	"github.com/PuerkitoBio/goquery"
)

var ErrUnexpectedPageStructure = errors.New("cornell: Unexpected page structure")

var apiKeyRegexp = regexp.MustCompile(`"key":"(?P<Key>[^"]*)`)
var nameRegexp = regexp.MustCompile(`login.*"name":"([^"]*)`) //kinda hacky, I should fix this

func (s *school) Enroll(tx *sql.Tx, user *data.User, params map[string]interface{}) (map[string]interface{}, error) {
	netidRaw, ok := params["username"]
	passwordRaw, ok2 := params["password"]

	if !ok || !ok2 {
		return nil, data.SchoolError{Code: "missing_params"}
	}

	netID, ok := netidRaw.(string)
	password, ok2 := passwordRaw.(string)

	if !ok || !ok2 || netID == "" || password == "" {
		return nil, data.SchoolError{Code: "invalid_params"}
	}

	cookieJar, _ := cookiejar.New(nil)
	c := &http.Client{
		Jar: cookieJar,
	}

	term := GetCurrentTerm().Code

	resp, err := c.Get("https://classes.cornell.edu/sascuwalogin/login/redirect?redirectUri=https%3A%2F%2Fclasses.cornell.edu%2Fscheduler%2Froster%2F" + term)
	if err != nil {
		return nil, err
	}

	doc, err := goquery.NewDocumentFromReader(resp.Body)
	if err != nil {
		return nil, err
	}

	loginForm, exists := doc.Find("form").First().Attr("action")
	if !exists {
		return nil, ErrUnexpectedPageStructure
	}

	values := url.Values{
		"j_username":       {netID},
		"j_password":       {password},
		"_eventId_proceed": {"Login"},
	}

	loginResp, err := c.PostForm(resp.Request.URL.Scheme+"://"+resp.Request.URL.Host+loginForm, values)
	if err != nil {
		return nil, err
	}

	confirmationDoc, err := goquery.NewDocumentFromReader(loginResp.Body)
	confirmationDocString, _ := confirmationDoc.Html()

	if strings.Contains(confirmationDocString, "Unable to log in") {
		return nil, data.SchoolError{Code: "creds_incorrect"}
	}

	// otherwise we're logged in! Unfortunately, CUWebLogin has this intermediate screen that says Your login credentials are being transmitted to the website via POST. This only works if you use javascript, so we need to manually transmit the token

	classesURL, exists := confirmationDoc.Find("form").First().Attr("action")
	if !exists {
		return nil, ErrUnexpectedPageStructure
	}

	relayState, exists := confirmationDoc.Find("input[name=RelayState]").First().Attr("value")
	if !exists {
		return nil, ErrUnexpectedPageStructure

	samlResponse, exists := confirmationDoc.Find("input[name=SAMLResponse]").First().Attr("value")
	if !exists {
		return nil, ErrUnexpectedPageStructure
	}

	schedulerResp, err := c.PostForm(classesURL, url.Values{"RelayState": {relayState}, "SAMLResponse": {samlResponse}})
	if err != nil {
		return nil, err
	}

	schedulerRespBuf := new(bytes.Buffer)
	schedulerRespBuf.ReadFrom(schedulerResp.Body)
	schedulerRespContent := schedulerRespBuf.String()

	APIKey := apiKeyRegexp.FindStringSubmatch(schedulerRespContent)[1]
	name := nameRegexp.FindStringSubmatch(schedulerRespContent)[1]

	scheduleReq, err := http.NewRequest("GET", "https://classes.cornell.edu/api/3.0/scheduler/current-enrollment", nil)
	if err != nil {
		return nil, err
	}

	scheduleReq.Header.Add("Authorization", "ClassRoster "+APIKey)

	scheduleResp, err := c.Do(scheduleReq)
	if err != nil {
		return nil, err
	}

	scheduleRespBuf := new(bytes.Buffer)
	scheduleRespBuf.ReadFrom(scheduleResp.Body)
	classesBytes := scheduleRespBuf.Bytes()

	courses := []classItem{}

	err = json.Unmarshal(classesBytes, &courses)
	if err != nil {
		return nil, err
	}

	cpairs := coursePairs{}

	for _, class := range courses {
		cpairs.CoursePairs = append(cpairs.CoursePairs, strconv.Itoa(class.CourseID)+","+strconv.Itoa(class.CourseOfferNumber))
	}

	coursePairsJSON, err := json.Marshal(cpairs)

	if err != nil {
		return nil, err
	}

	courseDetailReq, err := http.NewRequest("POST", "https://classes.cornell.edu/api/3.0/scheduler/course-detail", bytes.NewBuffer(coursePairsJSON))
	if err != nil {
		return nil, err
	}
	courseDetailReq.Header.Set("Content-Type", "application/json")
	courseDetailReq.Header.Add("Authorization", "ClassRoster "+APIKey)

	courseDetailsResp, err := c.Do(courseDetailReq)
	if err != nil {
		return nil, err
	}

	courseDetailBuf := new(bytes.Buffer)
	courseDetailBuf.ReadFrom(courseDetailsResp.Body)

	courseDetails := []course{}
	err = json.Unmarshal(courseDetailBuf.Bytes(), &courseDetails)
	if err != nil {
		return nil, err
	}

	// ok, now we have all the course details and the classes that the student is taking. Now we need to put it all in the db.

	for _, course := range courses {
		_, err = tx.Exec("INSERT INTO cornell_courses (userId, subject, catalogNum, title, units, rosterId) VALUES (?, ?, ?, ?, ?, ?)", user.ID, course.Subject, course.CatalogNumber, course.Title, course.Units, course.CourseID)
		if err != nil {
			return nil, err
		}

		for _, details := range courseDetails {
			if details.CourseID == course.CourseID {
				// we matched the details with the course!
				for _, eg := range details.EnrollGroups {
					for _, section := range eg.ClassSections {
						if util.IntSliceContains(course.ClassNumbers, section.ClassNum) {
							for _, meeting := range section.Meetings {
								if meeting.StartTime == "" || meeting.EndTime == "" {
									// this class is missing a time
									// not super clear what these are, skip it
									continue
								}

								startDate, err := time.Parse("01/02/2006", meeting.StartDate)
								if err != nil {
									return nil, err
								}

								endDate, err := time.Parse("01/02/2006", meeting.EndDate)
								if err != nil {
									return nil, err
								}

								startTime, err := time.Parse("3:04PM", meeting.StartTime)
								if err != nil {
									return nil, err
								}

								endTime, err := time.Parse("3:04PM", meeting.EndTime)
								if err != nil {
									return nil, err
								}

								// kinda hacky but it works
								startTime = time.Date(1970, 01, 01, startTime.Hour(), startTime.Minute(), 0, 0, time.UTC)
								endTime = time.Date(1970, 01, 01, endTime.Hour(), endTime.Minute(), 0, 0, time.UTC)

								satMeeting := false
								sunMeeting := false

								if strings.Contains("SSu", meeting.Pattern) {
									satMeeting = true
									sunMeeting = true
								} else if strings.Contains("Su", meeting.Pattern) {
									sunMeeting = true
								} else if strings.Contains("S", meeting.Pattern) {
									satMeeting = true
								}

								_, err = tx.Exec(`INSERT INTO cornell_events (
		title,
		userId,
		subject,
		catalogNum,
		classNum,
		component,
		componentLong,
		section,
		campus,
		campusLong,
		location,
		locationLong,
		startDate,
		endDate,
		startTime,
		endTime,
		monday,
		tuesday,
		wednesday,
		thursday,
		friday,
		saturday,
		sunday,
		facility,
		facilityLong,
		building
	)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
									course.Title,
									user.ID,
									course.Subject,
									course.CatalogNumber,
									section.ClassNum,
									section.Component,
									section.ComponentLong,
									section.Section,
									section.Campus,
									section.CampusDesc,
									section.Location,
									section.LocationDesc,
									startDate.Format("2006-01-02"),
									endDate.Format("2006-01-02"),
									startTime.Unix(),
									endTime.Unix(),
									strings.Contains(meeting.Pattern, "M"),
									strings.Contains(meeting.Pattern, "T"),
									strings.Contains(meeting.Pattern, "W"),
									strings.Contains(meeting.Pattern, "R"),
									strings.Contains(meeting.Pattern, "F"),
									satMeeting,
									sunMeeting,
									meeting.FacilityDescShort,
									meeting.FacilityDesc,
									meeting.BuildingDesc,
								)

								if err != nil {
									return nil, err
								}
							}
						}
					}
				}
				break
			}
		}
	}

	return map[string]interface{}{
		"netid":  strings.ToLower(netID),
		"name":   name,
		"status": schools.ImportStatusOK,
	}, nil
}

func (s *school) NeedsUpdate(db *sql.DB) (bool, error) {
	return (s.importStatus != schools.ImportStatusOK), nil
}

func (s *school) Unenroll(tx *sql.Tx, user *data.User) error {
	return clearUserData(tx, user)
}

func clearUserData(tx *sql.Tx, user *data.User) error {
	// clear away anything that is in the db
	_, err := tx.Exec("DELETE FROM cornell_courses WHERE userId = ?", user.ID)
	if err != nil {
		return err
	}

	_, err = tx.Exec("DELETE FROM cornell_events WHERE userId = ?", user.ID)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package store

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/Dataman-Cloud/swan/src/utils"

	"github.com/Sirupsen/logrus"
	zookeeper "github.com/samuel/go-zookeeper/zk"
	"golang.org/x/net/context"
)

const (
	SWAN_ATOMIC_STORE_NODE_PATH = "%s/atomic-store"
	SWAN_SNAPSHOT_PATH          = "%s/snapshot"
)

var (
	ZK_DEFAULT_ACL = zookeeper.WorldACL(zookeeper.PermAll)
)

// represents atomic operation both apply to ZK and intertal
// storage structure
type StoreOp uint8

var (
	OP_ADD    StoreOp = 1
	OP_REMOVE StoreOp = 2
	OP_UPDATE StoreOp = 3
)

func (op StoreOp) String() string {
	switch op {
	case OP_ADD:
		return "OP_ADD"
	case OP_REMOVE:
		return "OP_REMOVE"
	case OP_UPDATE:
		return "OP_UPDATE"
	}

	return ""
}

// represents the object type been manipulation by
// any specfic operation
type StoreEntity uint8

var (
	ENTITY_APP                  StoreEntity = 1
	ENTITY_SLOT                 StoreEntity = 2
	ENTITY_VERSION              StoreEntity = 3
	ENTITY_CURRENT_TASK         StoreEntity = 4
	ENTITY_FRAMEWORKID          StoreEntity = 5
	ENTITY_OFFER_ALLOCATOR_ITEM StoreEntity = 6
)

func (entity StoreEntity) String() string {
	switch entity {
	case ENTITY_APP:
		return "ENTITY_APP"
	case ENTITY_SLOT:
		return "ENTITY_SLOT"
	case ENTITY_VERSION:
		return "ENTITY_VERSION"
	case ENTITY_CURRENT_TASK:
		return "ENTITY_CURRENT_TASK"
	case ENTITY_FRAMEWORKID:
		return "ENTITY_FRAMEWORKID"
	case ENTITY_OFFER_ALLOCATOR_ITEM:
		return "ENTITY_OFFER_ALLOCATOR_ITEM"
	}

	return ""
}

var (
	ErrAppNotFound          = errors.New("app not found")
	ErrAppAlreadyExists     = errors.New("app already exists")
	ErrSlotNotFound         = errors.New("slot not found")
	ErrSlotAlreadyExists    = errors.New("slot already exists")
	ErrVersionAlreadyExists = errors.New("version already exists")
)

type AtomicOp struct {
	// atomic operaiton type, ADD | REMOVE | UPDATE
	Op StoreOp
	// which object type been operating
	Entity StoreEntity
	// can be explained by any specfic operation & object type operating on  appId/slotId maybe
	Param1 string
	// same as Param1
	Param2 string
	// same as Param1
	Param3 string
	// contains the data that the operation care, mostly object itself like App/Slot/Version
	Payload interface{}
}

type appHolder struct {
	App      *Application        `json:"app"`
	Versions map[string]*Version `json:"versions"`
	Slots    map[string]*Slot    `json:"slots"`
}

type Storage struct {
	Apps           map[string]*appHolder          `json:"apps"`
	OfferAllocator map[string]*OfferAllocatorItem `json:"offerAllocator"`
	FrameworkId    string                         `json:"frameworkid"`
}

func NewStorage() *Storage {
	return &Storage{
		Apps:           make(map[string]*appHolder),
		OfferAllocator: make(map[string]*OfferAllocatorItem),
	}
}

var (
	zs   *ZkStore
	once sync.Once
)

type ZkStore struct {
	Storage                  *Storage
	lastSequentialZkNodePath string
	lastSnapshotRevision     string

	readyToSnapshot bool
	mu              sync.RWMutex
	conn            *zookeeper.Conn
	zkPath          *url.URL
}

func DB() *ZkStore {
	return zs
}

func InitZkStore(zkPath *url.URL) error {
	conn, _, err := zookeeper.Connect(strings.Split(zkPath.Host, ","), 5*time.Second)
	if err != nil {
		return err
	}
	// TODO seems lack of re-connecting logic

	once.Do(func() {
		zs = &ZkStore{
			conn:            conn,
			Storage:         NewStorage(),
			zkPath:          zkPath,
			readyToSnapshot: false,
		}
	})
	return nil
}

func (zk *ZkStore) Start(ctx context.Context) error {
	ticker := time.NewTicker(2 * time.Second)
	for {
		select {
		case <-ticker.C:
			if zk.readyToSnapshot {
				revisionSnapshotted, err := zk.snapshot()
				if err != nil {
					logrus.Error(err)
				}

				err = zk.removeStaleAtomicOp(revisionSnapshotted)
				if err != nil {
					logrus.Error(err)
				}
			}
		case <-ctx.Done():
			zk.readyToSnapshot = false
			logrus.Info("store shutdown snapshot goroutine by ctx cancel")
			return ctx.Err()
		}
	}
}

func (zk *ZkStore) Apply(op *AtomicOp, zkPersistNeeded bool) error {
	zk.mu.Lock()
	defer zk.mu.Unlock()
	logrus.Debugf("Appling %s %s", op.Op.String(), op.Entity.String())

	switch op.Entity {
	case ENTITY_FRAMEWORKID:
		zk.applyFrameworkId(op)
	case ENTITY_APP:
		zk.applyApp(op)
	case ENTITY_SLOT:
		zk.applySlot(op)
	case ENTITY_VERSION:
		zk.applyVersion(op)
	case ENTITY_CURRENT_TASK:
		zk.applyCurrentTask(op)
	case ENTITY_OFFER_ALLOCATOR_ITEM:
		zk.applyOfferAllocatorItem(op)
	default:
		panic("invalid entity type")
	}

	var buf bytes.Buffer
	enc := json.NewEncoder(&buf)
	err := enc.Encode(op)
	if err != nil {
		return err
	}

	nodePath := filepath.Join(fmt.Sprintf(SWAN_ATOMIC_STORE_NODE_PATH, zk.zkPath.Path), "prefix")
	zk.lastSequentialZkNodePath, err = zk.conn.Create(
		nodePath,
		buf.Bytes(),
		zookeeper.FlagSequence,
		ZK_DEFAULT_ACL)

	if err != nil {
		return err
	}

	logrus.Debugf("create sequence node path is %s", zk.lastSequentialZkNodePath)
	return nil
}

func (zk *ZkStore) applyOfferAllocatorItem(op *AtomicOp) {
	switch op.Op {
	case OP_ADD:
		zk.Storage.OfferAllocator[op.Param1] = op.Payload.(*OfferAllocatorItem)
	case OP_REMOVE:
		delete(zk.Storage.OfferAllocator, op.Param1)
	case OP_UPDATE:
		delete(zk.Storage.OfferAllocator, op.Param1)
		zk.Storage.OfferAllocator[op.Param1] = op.Payload.(*OfferAllocatorItem)
	default:
		panic("applyFrameworkId not supportted operation")
	}
}

func (zk *ZkStore) applyCurrentTask(op *AtomicOp) {
	switch op.Op {
	case OP_UPDATE:
		zk.Storage.Apps[op.Param1].Slots[op.Param2].CurrentTask = op.Payload.(*Task)
	default:
		panic("applyCurrentTask not supportted operation")
	}
}

func (zk *ZkStore) applySlot(op *AtomicOp) {
	switch op.Op {
	case OP_ADD:
		zk.Storage.Apps[op.Param1].Slots[op.Param2] = op.Payload.(*Slot)
	case OP_REMOVE:
		delete(zk.Storage.Apps[op.Param1].Slots, op.Param2)
	case OP_UPDATE:
		zk.Storage.Apps[op.Param1].Slots[op.Param2] = op.Payload.(*Slot)
	default:
		panic("applySlot not supported operation")
	}
}

func (zk *ZkStore) applyVersion(op *AtomicOp) {
	switch op.Op {
	case OP_ADD:
		zk.Storage.Apps[op.Param1].Versions[op.Param2] = op.Payload.(*Version)
	default:
		panic("applyVersion not supportted operation")
	}
}

func (zk *ZkStore) applyFrameworkId(op *AtomicOp) {
	switch op.Op {
	case OP_ADD:
		zk.Storage.FrameworkId = op.Payload.(string)
	case OP_REMOVE:
		zk.Storage.FrameworkId = ""
	case OP_UPDATE:
		zk.Storage.FrameworkId = op.Payload.(string)
	default:
		panic("applyFrameworkId not supportted operation")
	}
}

func (zk *ZkStore) applyApp(op *AtomicOp) {
	switch op.Op {
	case OP_ADD:
		zk.Storage.Apps[op.Param1] = &appHolder{
			App:      op.Payload.(*Application),
			Versions: make(map[string]*Version),
			Slots:    make(map[string]*Slot),
		}
	case OP_REMOVE:
		delete(zk.Storage.Apps, op.Param1)

	case OP_UPDATE:
		if _, found := zk.Storage.Apps[op.Param1]; found {
			zk.Storage.Apps[op.Param1].App = op.Payload.(*Application)
		}

	default:
		panic("applyApp not supportted operation")
	}
}

// remove any atomic op that are already snapshotted
func (zk *ZkStore) removeStaleAtomicOp(snapshotTo string) error {
	atomicStorePath := fmt.Sprintf(SWAN_ATOMIC_STORE_NODE_PATH, zk.zkPath.Path)
	children, _, err := zk.conn.Children(atomicStorePath)
	if err != nil {
		return err
	}

	for _, child := range children {
		if strings.Compare(snapshotTo, child) != 1 {
			logrus.Debugf("deleting %s now", child)
			err := zk.conn.Delete(filepath.Join(atomicStorePath, child), -1)
			if err != nil {
				return err
			}
		}
	}

	logrus.Debugf("revision before %s are cleared", snapshotTo)
	return nil
}

func (zk *ZkStore) snapshot() (string, error) {
	snapshotPath := fmt.Sprintf(SWAN_SNAPSHOT_PATH, zk.zkPath.Path)
	revision := zk.lastSequentialZkNodePath

	zk.mu.Lock()
	data, err := json.Marshal(zk.Storage)
	if err != nil {
		return "", err
	}
	zk.mu.Unlock()

	logrus.Debugf("snapshot storage to zk with data len: %d", len(data))
	logrus.Debugf(string(data))

	exists, _, err := zk.conn.Exists(snapshotPath)
	if err != nil {
		return "", err
	}

	if !exists {
		_, err = zk.conn.Create(snapshotPath, data, 0, ZK_DEFAULT_ACL)
	} else {
		_, err = zk.conn.Set(snapshotPath, data, -1)
	}
	if err != nil {
		return "", err
	}

	zk.lastSnapshotRevision = revision
	logrus.Debugf("snapshot %s to zk success", revision)

	return revision, nil
}

func (zk *ZkStore) Recover() error {
	zk.mu.Lock()
	defer zk.mu.Unlock()

	if err := zk.recoverFromSnapshot(); err != nil {
		return err
	}

	if err := zk.recoverFromAtomicSequentialSlice(); err != nil {
		return err
	}

	zk.readyToSnapshot = true

	return nil
}

func (zk *ZkStore) recoverFromSnapshot() error {
	logrus.Debugf("syncFromSnapshot now")

	snapshotPath := fmt.Sprintf(SWAN_SNAPSHOT_PATH, zk.zkPath.Path)
	exists, _, err := zk.conn.Exists(snapshotPath)
	if err != nil {
		return err
	}

	if !exists { // do nothing when fresh start
		return nil
	}

	data, _, err := zk.conn.Get(snapshotPath)
	if err != nil {
		return err
	}

	err = json.Unmarshal(data, zk.Storage)
	if err != nil {
		return err
	}

	return nil
}

func (zk *ZkStore) recoverFromAtomicSequentialSlice() error {
	atomicStorePath := fmt.Sprintf(SWAN_ATOMIC_STORE_NODE_PATH, zk.zkPath.Path)

	children, _, err := zk.conn.Children(atomicStorePath)
	if err != nil {
		return err
	}

	sortedPaths := utils.SortableNodePath(children)
	sort.Sort(sortedPaths)

	for _, child := range sortedPaths {
		data, _, err := zk.conn.Get(filepath.Join(atomicStorePath, child))
		if err != nil {
			return err
		}

		ao, err := zk.unmarshalAtomicOp(data)
		if err != nil {
			return err
		}

		logrus.Debugf("decode  %s got %+v", child, ao)
		zk.Apply(ao, false)
	}

	return nil
}

func (zk *ZkStore) unmarshalAtomicOp(data []byte) (*AtomicOp, error) {
	var tmpAo struct {
		Op     StoreOp
		Entity StoreEntity
		Param1 string
		Param2 string
		Param3 string

		Payload json.RawMessage
	}

	err := json.Unmarshal(data, &tmpAo)
	if err != nil {
		return nil, err
	}

	var ao AtomicOp
	ao.Op = tmpAo.Op
	ao.Entity = tmpAo.Entity
	ao.Param1 = tmpAo.Param1
	ao.Param2 = tmpAo.Param2
	ao.Param3 = tmpAo.Param3

	if tmpAo.Payload != nil {
		switch ao.Entity {
		case ENTITY_APP:
			var app Application
			err = json.Unmarshal(tmpAo.Payload, &app)
			if err != nil {
				return nil, err
			}
			ao.Payload = &app
		case ENTITY_SLOT:
			var slot Slot
			err = json.Unmarshal(tmpAo.Payload, &slot)
			if err != nil {
				return nil, err
			}
			ao.Payload = &slot
		case ENTITY_VERSION:
			var version Version
			err = json.Unmarshal(tmpAo.Payload, &version)
			if err != nil {
				return nil, err
			}
			ao.Payload = &version
		case ENTITY_CURRENT_TASK:
			var task Task
			err = json.Unmarshal(tmpAo.Payload, &task)
			if err != nil {
				return nil, err
			}
			ao.Payload = &task

		case ENTITY_FRAMEWORKID:
			var frameworkid string
			err = json.Unmarshal(tmpAo.Payload, &frameworkid)
			if err != nil {
				return nil, err
			}
			ao.Payload = frameworkid

		case ENTITY_OFFER_ALLOCATOR_ITEM:
			var item OfferAllocatorItem
			err = json.Unmarshal(tmpAo.Payload, &item)
			if err != nil {
				return nil, err
			}
			ao.Payload = &item
		}
	}
	return &ao, nil
}
<commit_msg>reduce snapshot frequency<commit_after>package store

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/Dataman-Cloud/swan/src/utils"

	"github.com/Sirupsen/logrus"
	zookeeper "github.com/samuel/go-zookeeper/zk"
	"golang.org/x/net/context"
)

const (
	SWAN_ATOMIC_STORE_NODE_PATH = "%s/atomic-store"
	SWAN_SNAPSHOT_PATH          = "%s/snapshot"
)

var (
	ZK_DEFAULT_ACL = zookeeper.WorldACL(zookeeper.PermAll)
)

// represents atomic operation both apply to ZK and intertal
// storage structure
type StoreOp uint8

var (
	OP_ADD    StoreOp = 1
	OP_REMOVE StoreOp = 2
	OP_UPDATE StoreOp = 3
)

func (op StoreOp) String() string {
	switch op {
	case OP_ADD:
		return "OP_ADD"
	case OP_REMOVE:
		return "OP_REMOVE"
	case OP_UPDATE:
		return "OP_UPDATE"
	}

	return ""
}

// represents the object type been manipulation by
// any specfic operation
type StoreEntity uint8

var (
	ENTITY_APP                  StoreEntity = 1
	ENTITY_SLOT                 StoreEntity = 2
	ENTITY_VERSION              StoreEntity = 3
	ENTITY_CURRENT_TASK         StoreEntity = 4
	ENTITY_FRAMEWORKID          StoreEntity = 5
	ENTITY_OFFER_ALLOCATOR_ITEM StoreEntity = 6
)

func (entity StoreEntity) String() string {
	switch entity {
	case ENTITY_APP:
		return "ENTITY_APP"
	case ENTITY_SLOT:
		return "ENTITY_SLOT"
	case ENTITY_VERSION:
		return "ENTITY_VERSION"
	case ENTITY_CURRENT_TASK:
		return "ENTITY_CURRENT_TASK"
	case ENTITY_FRAMEWORKID:
		return "ENTITY_FRAMEWORKID"
	case ENTITY_OFFER_ALLOCATOR_ITEM:
		return "ENTITY_OFFER_ALLOCATOR_ITEM"
	}

	return ""
}

var (
	ErrAppNotFound          = errors.New("app not found")
	ErrAppAlreadyExists     = errors.New("app already exists")
	ErrSlotNotFound         = errors.New("slot not found")
	ErrSlotAlreadyExists    = errors.New("slot already exists")
	ErrVersionAlreadyExists = errors.New("version already exists")
)

type AtomicOp struct {
	// atomic operaiton type, ADD | REMOVE | UPDATE
	Op StoreOp
	// which object type been operating
	Entity StoreEntity
	// can be explained by any specfic operation & object type operating on  appId/slotId maybe
	Param1 string
	// same as Param1
	Param2 string
	// same as Param1
	Param3 string
	// contains the data that the operation care, mostly object itself like App/Slot/Version
	Payload interface{}
}

type appHolder struct {
	App      *Application        `json:"app"`
	Versions map[string]*Version `json:"versions"`
	Slots    map[string]*Slot    `json:"slots"`
}

type Storage struct {
	Apps           map[string]*appHolder          `json:"apps"`
	OfferAllocator map[string]*OfferAllocatorItem `json:"offerAllocator"`
	FrameworkId    string                         `json:"frameworkid"`
}

func NewStorage() *Storage {
	return &Storage{
		Apps:           make(map[string]*appHolder),
		OfferAllocator: make(map[string]*OfferAllocatorItem),
	}
}

var (
	zs   *ZkStore
	once sync.Once
)

type ZkStore struct {
	Storage                  *Storage
	lastSequentialZkNodePath string
	lastSnapshotRevision     string

	readyToSnapshot bool
	mu              sync.RWMutex
	conn            *zookeeper.Conn
	zkPath          *url.URL
}

func DB() *ZkStore {
	return zs
}

func InitZkStore(zkPath *url.URL) error {
	conn, _, err := zookeeper.Connect(strings.Split(zkPath.Host, ","), 5*time.Second)
	if err != nil {
		return err
	}
	// TODO seems lack of re-connecting logic

	once.Do(func() {
		zs = &ZkStore{
			conn:            conn,
			Storage:         NewStorage(),
			zkPath:          zkPath,
			readyToSnapshot: false,
		}
	})
	return nil
}

func (zk *ZkStore) Start(ctx context.Context) error {
	ticker := time.NewTicker(2 * time.Second)
	for {
		select {
		case <-ticker.C:
			if zk.readyToSnapshot && zk.lastSnapshotRevision != zk.lastSequentialZkNodePath {
				revisionSnapshotted, err := zk.snapshot()
				if err != nil {
					logrus.Error(err)
				}

				err = zk.removeStaleAtomicOp(revisionSnapshotted)
				if err != nil {
					logrus.Error(err)
				}
			}
		case <-ctx.Done():
			zk.readyToSnapshot = false
			logrus.Info("store shutdown snapshot goroutine by ctx cancel")
			return ctx.Err()
		}
	}
}

func (zk *ZkStore) Apply(op *AtomicOp, zkPersistNeeded bool) error {
	zk.mu.Lock()
	defer zk.mu.Unlock()
	logrus.Debugf("Appling %s %s", op.Op.String(), op.Entity.String())

	switch op.Entity {
	case ENTITY_FRAMEWORKID:
		zk.applyFrameworkId(op)
	case ENTITY_APP:
		zk.applyApp(op)
	case ENTITY_SLOT:
		zk.applySlot(op)
	case ENTITY_VERSION:
		zk.applyVersion(op)
	case ENTITY_CURRENT_TASK:
		zk.applyCurrentTask(op)
	case ENTITY_OFFER_ALLOCATOR_ITEM:
		zk.applyOfferAllocatorItem(op)
	default:
		panic("invalid entity type")
	}

	var buf bytes.Buffer
	enc := json.NewEncoder(&buf)
	err := enc.Encode(op)
	if err != nil {
		return err
	}

	nodePath := filepath.Join(fmt.Sprintf(SWAN_ATOMIC_STORE_NODE_PATH, zk.zkPath.Path), "prefix")
	zk.lastSequentialZkNodePath, err = zk.conn.Create(
		nodePath,
		buf.Bytes(),
		zookeeper.FlagSequence,
		ZK_DEFAULT_ACL)

	if err != nil {
		return err
	}

	logrus.Debugf("create sequence node path is %s", zk.lastSequentialZkNodePath)
	return nil
}

func (zk *ZkStore) applyOfferAllocatorItem(op *AtomicOp) {
	switch op.Op {
	case OP_ADD:
		zk.Storage.OfferAllocator[op.Param1] = op.Payload.(*OfferAllocatorItem)
	case OP_REMOVE:
		delete(zk.Storage.OfferAllocator, op.Param1)
	case OP_UPDATE:
		delete(zk.Storage.OfferAllocator, op.Param1)
		zk.Storage.OfferAllocator[op.Param1] = op.Payload.(*OfferAllocatorItem)
	default:
		panic("applyFrameworkId not supportted operation")
	}
}

func (zk *ZkStore) applyCurrentTask(op *AtomicOp) {
	switch op.Op {
	case OP_UPDATE:
		zk.Storage.Apps[op.Param1].Slots[op.Param2].CurrentTask = op.Payload.(*Task)
	default:
		panic("applyCurrentTask not supportted operation")
	}
}

func (zk *ZkStore) applySlot(op *AtomicOp) {
	switch op.Op {
	case OP_ADD:
		zk.Storage.Apps[op.Param1].Slots[op.Param2] = op.Payload.(*Slot)
	case OP_REMOVE:
		delete(zk.Storage.Apps[op.Param1].Slots, op.Param2)
	case OP_UPDATE:
		zk.Storage.Apps[op.Param1].Slots[op.Param2] = op.Payload.(*Slot)
	default:
		panic("applySlot not supported operation")
	}
}

func (zk *ZkStore) applyVersion(op *AtomicOp) {
	switch op.Op {
	case OP_ADD:
		zk.Storage.Apps[op.Param1].Versions[op.Param2] = op.Payload.(*Version)
	default:
		panic("applyVersion not supportted operation")
	}
}

func (zk *ZkStore) applyFrameworkId(op *AtomicOp) {
	switch op.Op {
	case OP_ADD:
		zk.Storage.FrameworkId = op.Payload.(string)
	case OP_REMOVE:
		zk.Storage.FrameworkId = ""
	case OP_UPDATE:
		zk.Storage.FrameworkId = op.Payload.(string)
	default:
		panic("applyFrameworkId not supportted operation")
	}
}

func (zk *ZkStore) applyApp(op *AtomicOp) {
	switch op.Op {
	case OP_ADD:
		zk.Storage.Apps[op.Param1] = &appHolder{
			App:      op.Payload.(*Application),
			Versions: make(map[string]*Version),
			Slots:    make(map[string]*Slot),
		}
	case OP_REMOVE:
		delete(zk.Storage.Apps, op.Param1)

	case OP_UPDATE:
		if _, found := zk.Storage.Apps[op.Param1]; found {
			zk.Storage.Apps[op.Param1].App = op.Payload.(*Application)
		}

	default:
		panic("applyApp not supportted operation")
	}
}

// remove any atomic op that are already snapshotted
func (zk *ZkStore) removeStaleAtomicOp(snapshotTo string) error {
	atomicStorePath := fmt.Sprintf(SWAN_ATOMIC_STORE_NODE_PATH, zk.zkPath.Path)
	children, _, err := zk.conn.Children(atomicStorePath)
	if err != nil {
		return err
	}

	for _, child := range children {
		if strings.Compare(snapshotTo, child) != 1 {
			logrus.Debugf("deleting %s now", child)
			err := zk.conn.Delete(filepath.Join(atomicStorePath, child), -1)
			if err != nil {
				return err
			}
		}
	}

	logrus.Debugf("revision before %s are cleared", snapshotTo)
	return nil
}

func (zk *ZkStore) snapshot() (string, error) {
	snapshotPath := fmt.Sprintf(SWAN_SNAPSHOT_PATH, zk.zkPath.Path)
	revision := zk.lastSequentialZkNodePath

	zk.mu.Lock()
	data, err := json.Marshal(zk.Storage)
	if err != nil {
		return "", err
	}
	zk.mu.Unlock()

	logrus.Debugf("snapshot storage to zk with data len: %d", len(data))
	logrus.Debugf(string(data))

	exists, _, err := zk.conn.Exists(snapshotPath)
	if err != nil {
		return "", err
	}

	if !exists {
		_, err = zk.conn.Create(snapshotPath, data, 0, ZK_DEFAULT_ACL)
	} else {
		_, err = zk.conn.Set(snapshotPath, data, -1)
	}
	if err != nil {
		return "", err
	}

	zk.lastSnapshotRevision = revision
	logrus.Debugf("snapshot %s to zk success", revision)

	return revision, nil
}

func (zk *ZkStore) Recover() error {
	zk.mu.Lock()
	defer zk.mu.Unlock()

	if err := zk.recoverFromSnapshot(); err != nil {
		return err
	}

	if err := zk.recoverFromAtomicSequentialSlice(); err != nil {
		return err
	}

	zk.readyToSnapshot = true

	return nil
}

func (zk *ZkStore) recoverFromSnapshot() error {
	logrus.Debugf("syncFromSnapshot now")

	snapshotPath := fmt.Sprintf(SWAN_SNAPSHOT_PATH, zk.zkPath.Path)
	exists, _, err := zk.conn.Exists(snapshotPath)
	if err != nil {
		return err
	}

	if !exists { // do nothing when fresh start
		return nil
	}

	data, _, err := zk.conn.Get(snapshotPath)
	if err != nil {
		return err
	}

	err = json.Unmarshal(data, zk.Storage)
	if err != nil {
		return err
	}

	return nil
}

func (zk *ZkStore) recoverFromAtomicSequentialSlice() error {
	atomicStorePath := fmt.Sprintf(SWAN_ATOMIC_STORE_NODE_PATH, zk.zkPath.Path)

	children, _, err := zk.conn.Children(atomicStorePath)
	if err != nil {
		return err
	}

	sortedPaths := utils.SortableNodePath(children)
	sort.Sort(sortedPaths)

	for _, child := range sortedPaths {
		data, _, err := zk.conn.Get(filepath.Join(atomicStorePath, child))
		if err != nil {
			return err
		}

		ao, err := zk.unmarshalAtomicOp(data)
		if err != nil {
			return err
		}

		logrus.Debugf("decode  %s got %+v", child, ao)
		zk.Apply(ao, false)
	}

	return nil
}

func (zk *ZkStore) unmarshalAtomicOp(data []byte) (*AtomicOp, error) {
	var tmpAo struct {
		Op     StoreOp
		Entity StoreEntity
		Param1 string
		Param2 string
		Param3 string

		Payload json.RawMessage
	}

	err := json.Unmarshal(data, &tmpAo)
	if err != nil {
		return nil, err
	}

	var ao AtomicOp
	ao.Op = tmpAo.Op
	ao.Entity = tmpAo.Entity
	ao.Param1 = tmpAo.Param1
	ao.Param2 = tmpAo.Param2
	ao.Param3 = tmpAo.Param3

	if tmpAo.Payload != nil {
		switch ao.Entity {
		case ENTITY_APP:
			var app Application
			err = json.Unmarshal(tmpAo.Payload, &app)
			if err != nil {
				return nil, err
			}
			ao.Payload = &app
		case ENTITY_SLOT:
			var slot Slot
			err = json.Unmarshal(tmpAo.Payload, &slot)
			if err != nil {
				return nil, err
			}
			ao.Payload = &slot
		case ENTITY_VERSION:
			var version Version
			err = json.Unmarshal(tmpAo.Payload, &version)
			if err != nil {
				return nil, err
			}
			ao.Payload = &version
		case ENTITY_CURRENT_TASK:
			var task Task
			err = json.Unmarshal(tmpAo.Payload, &task)
			if err != nil {
				return nil, err
			}
			ao.Payload = &task

		case ENTITY_FRAMEWORKID:
			var frameworkid string
			err = json.Unmarshal(tmpAo.Payload, &frameworkid)
			if err != nil {
				return nil, err
			}
			ao.Payload = frameworkid

		case ENTITY_OFFER_ALLOCATOR_ITEM:
			var item OfferAllocatorItem
			err = json.Unmarshal(tmpAo.Payload, &item)
			if err != nil {
				return nil, err
			}
			ao.Payload = &item
		}
	}
	return &ao, nil
}
<|endoftext|>


<commit_before>/*-
 * Copyright (c) 2016, 1&1 Internet SE
 * Copyright (c) 2016, Jörg Pernfuß
 *
 * Use of this source code is governed by a 2-clause BSD license
 * that can be found in the LICENSE file.
 */

package adm

import (
	"fmt"

	"github.com/1and1/soma/lib/proto"
	resty "gopkg.in/resty.v0"
)

// LookupOncallId looks up the UUID for an oncall duty on the
// server with name s. Error is set if no such oncall duty was
// found or an error occured.
// If s is already a UUID, then s is immediately returned.
func LookupOncallId(s string) (string, error) {
	if isUUID(s) {
		return s, nil
	}
	return oncallIdByName(s)
}

// oncallIdByName implements the actual serverside lookup of the
// oncall duty UUID
func oncallIdByName(oncall string) (string, error) {
	req := proto.Request{}
	req.Filter = &proto.Filter{}
	req.Filter.Oncall = &proto.OncallFilter{Name: oncall}

	var (
		err  error
		resp *resty.Response
		res  *proto.Result
	)
	if resp, err = PostReqBody(req, `/filter/oncall/`); err != nil {
		// transport errors
		goto abort
	}

	if res, err = decodeResponse(resp); err != nil {
		// http code errors
		goto abort
	}

	if res.StatusCode >= 300 {
		// application errors
		s := fmt.Sprintf("Request failed: %d - %s",
			res.StatusCode, res.StatusText)
		m := []string{s}

		if res.Errors != nil {
			m = append(m, *res.Errors...)
		}

		err = fmt.Errorf(combineStrings(m...))
		goto abort
	}

	// check the received record against the input
	if oncall != (*res.Oncalls)[0].Name {
		err = fmt.Errorf("Name mismatch: %s vs %s",
			oncall, (*res.Oncalls)[0].Name)
		goto abort
	}
	return (*res.Oncalls)[0].Id, nil

abort:
	return ``, fmt.Errorf("OncallId lookup failed: %s", err.Error())
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<commit_msg>adm: add LookupUserId()<commit_after>/*-
 * Copyright (c) 2016, 1&1 Internet SE
 * Copyright (c) 2016, Jörg Pernfuß
 *
 * Use of this source code is governed by a 2-clause BSD license
 * that can be found in the LICENSE file.
 */

package adm

import (
	"fmt"

	"github.com/1and1/soma/lib/proto"
	resty "gopkg.in/resty.v0"
)

// LookupOncallId looks up the UUID for an oncall duty on the
// server with name s. Error is set if no such oncall duty was
// found or an error occured.
// If s is already a UUID, then s is immediately returned.
func LookupOncallId(s string) (string, error) {
	if isUUID(s) {
		return s, nil
	}
	return oncallIdByName(s)
}

// LookupOncallId looks up the UUID for a user on the server
// with username s. Error is set if no such user was found
// or an error occured.
// If s is already a UUID, then s is immediately returned.
func LookupUserId(s string) (string, error) {
	if isUUID(s) {
		return s, nil
	}
	return userIdByUserName(s)
}

// oncallIdByName implements the actual serverside lookup of the
// oncall duty UUID
func oncallIdByName(oncall string) (string, error) {
	req := proto.NewOncallFilter()
	req.Filter.Oncall = &proto.OncallFilter{Name: oncall}

	var (
		err  error
		resp *resty.Response
		res  *proto.Result
	)
	if resp, err = PostReqBody(req, `/filter/oncall/`); err != nil {
		// transport errors
		goto abort
	}

	if res, err = decodeResponse(resp); err != nil {
		// http code errors
		goto abort
	}

	if err = checkApplicationError(res); err != nil {
		goto abort
	}

	// check the received record against the input
	if oncall != (*res.Oncalls)[0].Name {
		err = fmt.Errorf("Name mismatch: %s vs %s",
			oncall, (*res.Oncalls)[0].Name)
		goto abort
	}
	return (*res.Oncalls)[0].Id, nil

abort:
	return ``, fmt.Errorf("OncallId lookup failed: %s", err.Error())
}

// userIdByUserName implements the actual serverside lookup of the
// user's UUID
func userIdByUserName(user string) (string, error) {
	req := proto.NewUserFilter()
	req.Filter.User.UserName = user

	var (
		err  error
		resp *resty.Response
		res  *proto.Result
	)
	if resp, err = PostReqBody(req, `/filter/users/`); err != nil {
		// transport errors
		goto abort
	}

	if res, err = decodeResponse(resp); err != nil {
		// http code errors
		goto abort
	}

	if err = checkApplicationError(res); err != nil {
		goto abort
	}

	// check the received record against the input
	if user != (*res.Users)[0].UserName {
		err = fmt.Errorf("Name mismatch: %s vs %s",
			user, (*res.Users)[0].UserName)
		goto abort
	}
	return (*res.Users)[0].Id, nil

abort:
	return ``, fmt.Errorf("UserId lookup failed: %s", err.Error())
}

// checkApplicationError tests the server result for
// application errors
func checkApplicationError(result *proto.Result) error {
	if result.StatusCode >= 300 {
		// application errors
		s := fmt.Sprintf("Request failed: %d - %s",
			result.StatusCode, result.StatusText)
		m := []string{s}

		if result.Errors != nil {
			m = append(m, *result.Errors...)
		}

		return fmt.Errorf(combineStrings(m...))
	}
	return nil
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<|endoftext|>


<commit_before>// Copyright 2019 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 fetch provides a way to fetch modules from a proxy.
package fetch

import (
	"archive/zip"
	"bytes"
	"context"
	"errors"
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"io"
	"io/ioutil"
	"math"
	"net/http"
	"os"
	"path"
	"strings"

	"go.opencensus.io/trace"
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/config"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/godoc"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/source"
	"golang.org/x/pkgsite/internal/stdlib"
)

// BadPackageError represents an error loading a package
// because its contents do not make up a valid package.
//
// This can happen, for example, if the .go files fail
// to parse or declare different package names.
type BadPackageError struct {
	Err error // Not nil.
}

func (bpe *BadPackageError) Error() string { return bpe.Err.Error() }

// loadPackage loads a Go package by calling loadPackageWithBuildContext, trying
// several build contexts in turn. It returns a goPackage with documentation
// information for each build context that results in a valid package, in the
// same order that the build contexts are listed. If none of them result in a
// package, then loadPackage returns nil, nil.
//
// If a package is fine except that its documentation is too large, loadPackage
// returns a goPackage whose err field is a non-nil error with godoc.ErrTooLarge in its chain.
func loadPackage(ctx context.Context, zipGoFiles []*zip.File, innerPath string, sourceInfo *source.Info, modInfo *godoc.ModuleInfo) (_ *goPackage, err error) {
	defer derrors.Wrap(&err, "loadPackage(ctx, zipGoFiles, %q, sourceInfo, modInfo)", innerPath)
	ctx, span := trace.StartSpan(ctx, "fetch.loadPackage")
	defer span.End()

	// Make a map with all the zip file contents.
	files := make(map[string][]byte)
	for _, f := range zipGoFiles {
		_, name := path.Split(f.Name)
		b, err := readZipFile(f, MaxFileSize)
		if err != nil {
			return nil, err
		}
		files[name] = b
	}

	modulePath := modInfo.ModulePath
	importPath := path.Join(modulePath, innerPath)
	if modulePath == stdlib.ModulePath {
		importPath = innerPath
	}
	v1path := internal.V1Path(importPath, modulePath)

	var pkg *goPackage
	for _, bc := range internal.BuildContexts {
		mfiles, err := matchingFiles(bc.GOOS, bc.GOARCH, files)
		if err != nil {
			return nil, err
		}
		name, imports, synopsis, source, err := loadPackageWithBuildContext(ctx, bc.GOOS, bc.GOARCH, mfiles, innerPath, sourceInfo, modInfo)
		switch {
		case errors.Is(err, derrors.NotFound):
			// No package for this build context.
			continue
		case errors.Is(err, godoc.ErrTooLarge):
			// The doc for this build context is too large. Remember that and
			// return the package for this build context; ignore the others.
			return &goPackage{
				err:     err,
				path:    importPath,
				v1path:  v1path,
				name:    name,
				imports: imports,
				docs: []*internal.Documentation{{
					GOOS:     bc.GOOS,
					GOARCH:   bc.GOARCH,
					Synopsis: synopsis,
					Source:   source,
				}},
			}, nil
		case err != nil:
			// Serious error. Fail.
			return nil, err
		default:
			// No error.
			if pkg == nil {
				pkg = &goPackage{
					path:    importPath,
					v1path:  v1path,
					name:    name,
					imports: imports, // Use the imports from the first successful build context.
				}
			}
			// All the build contexts should use the same package name. Although
			// it's technically legal for different build tags to result in different
			// package names, it's not something we support.
			if name != pkg.name {
				return nil, &BadPackageError{
					Err: fmt.Errorf("more than one package name (%q and %q)", pkg.name, name),
				}
			}
			pkg.docs = append(pkg.docs, &internal.Documentation{
				GOOS:     bc.GOOS,
				GOARCH:   bc.GOARCH,
				Synopsis: synopsis,
				Source:   source,
			})
		}
	}
	return pkg, nil
}

// httpPost allows package fetch tests to stub out playground URL fetches.
var httpPost = http.Post

// loadPackageWithBuildContext loads a Go package made of .go files in
// files using a build context constructed from the given GOOS and GOARCH
// values. modulePath is stdlib.ModulePath for the Go standard library and the
// module path for all other modules. innerPath is the path of the Go package
// directory relative to the module root. The files argument must contain only
// .go files that have been verified to be of reasonable size and that match
// the build context.
//
// It returns the package name, list of imports, the package synopsis, and the
// serialized source (AST) for the package.
//
// It returns an error with NotFound in its chain if the directory doesn't
// contain a Go package or all .go files have been excluded by constraints. A
// *BadPackageError error is returned if the directory contains .go files but do
// not make up a valid package.
//
// If it returns an error with ErrTooLarge in its chain, the other return values
// are still valid.
func loadPackageWithBuildContext(ctx context.Context, goos, goarch string, files map[string][]byte, innerPath string, sourceInfo *source.Info, modInfo *godoc.ModuleInfo) (name string, imports []string, synopsis string, source []byte, err error) {
	modulePath := modInfo.ModulePath
	defer derrors.Wrap(&err, "loadPackageWithBuildContext(%q, %q, files, %q, %q, %+v)",
		goos, goarch, innerPath, modulePath, sourceInfo)

	packageName, goFiles, fset, err := loadFilesWithBuildContext(innerPath, files)
	if err != nil {
		return "", nil, "", nil, err
	}
	docPkg := godoc.NewPackage(fset, goos, goarch, modInfo.ModulePackages)
	for _, pf := range goFiles {
		removeNodes := true
		// Don't strip the seemingly unexported functions from the builtin package;
		// they are actually Go builtins like make, new, etc.
		if modulePath == stdlib.ModulePath && innerPath == "builtin" {
			removeNodes = false
		}
		docPkg.AddFile(pf, removeNodes)
	}

	// Encode first, because Render messes with the AST.
	src, err := docPkg.Encode(ctx)
	if err != nil {
		return "", nil, "", nil, err
	}

	synopsis, imports, _, err = docPkg.Render(ctx, innerPath, sourceInfo, modInfo, goos, goarch)
	if err != nil && !errors.Is(err, godoc.ErrTooLarge) {
		return "", nil, "", nil, err
	}
	return packageName, imports, synopsis, src, err
}

// loadFilesWithBuildContext loads all the given Go files at innerPath. It
// returns the package name as it occurs in the source, a map of the ASTs of all
// the Go files, and the token.FileSet used for parsing.
// If there are no non-test Go files, it returns a NotFound error.
func loadFilesWithBuildContext(innerPath string, files map[string][]byte) (pkgName string, fileMap map[string]*ast.File, _ *token.FileSet, _ error) {
	// Parse .go files and add them to the goFiles slice.
	var (
		fset            = token.NewFileSet()
		goFiles         = make(map[string]*ast.File)
		numNonTestFiles int
		packageName     string
		packageNameFile string // Name of file where packageName came from.
	)
	for name, b := range files {
		pf, err := parser.ParseFile(fset, name, b, parser.ParseComments)
		if err != nil {
			if pf == nil {
				return "", nil, nil, fmt.Errorf("internal error: the source couldn't be read: %v", err)
			}
			return "", nil, nil, &BadPackageError{Err: err}
		}
		// Remember all files, including test files for their examples.
		goFiles[name] = pf
		if strings.HasSuffix(name, "_test.go") {
			continue
		}
		// Keep track of the number of non-test files to check that the package name is the same.
		// and also because a directory with only test files doesn't count as a
		// Go package.
		numNonTestFiles++
		if numNonTestFiles == 1 {
			packageName = pf.Name.Name
			packageNameFile = name
		} else if pf.Name.Name != packageName {
			return "", nil, nil, &BadPackageError{Err: &build.MultiplePackageError{
				Dir:      innerPath,
				Packages: []string{packageName, pf.Name.Name},
				Files:    []string{packageNameFile, name},
			}}
		}
	}
	if numNonTestFiles == 0 {
		// This directory doesn't contain a package, or at least not one
		// that matches this build context.
		return "", nil, nil, derrors.NotFound
	}
	return packageName, goFiles, fset, nil
}

// matchingFiles returns a map from file names to their contents, read from zipGoFiles.
// It includes only those files that match the build context determined by goos and goarch.
func matchingFiles(goos, goarch string, allFiles map[string][]byte) (matchedFiles map[string][]byte, err error) {
	defer derrors.Wrap(&err, "matchingFiles(%q, %q, zipGoFiles)", goos, goarch)

	// bctx is used to make decisions about which of the .go files are included
	// by build constraints.
	bctx := &build.Context{
		GOOS:        goos,
		GOARCH:      goarch,
		CgoEnabled:  true,
		Compiler:    build.Default.Compiler,
		ReleaseTags: build.Default.ReleaseTags,

		JoinPath: path.Join,
		OpenFile: func(name string) (io.ReadCloser, error) {
			return ioutil.NopCloser(bytes.NewReader(allFiles[name])), nil
		},

		// If left nil, the default implementations of these read from disk,
		// which we do not want. None of these functions should be used
		// inside this function; it would be an internal error if they are.
		// Set them to non-nil values to catch if that happens.
		SplitPathList: func(string) []string { panic("internal error: unexpected call to SplitPathList") },
		IsAbsPath:     func(string) bool { panic("internal error: unexpected call to IsAbsPath") },
		IsDir:         func(string) bool { panic("internal error: unexpected call to IsDir") },
		HasSubdir:     func(string, string) (string, bool) { panic("internal error: unexpected call to HasSubdir") },
		ReadDir:       func(string) ([]os.FileInfo, error) { panic("internal error: unexpected call to ReadDir") },
	}

	// Copy the input map so we don't modify it.
	matchedFiles = map[string][]byte{}
	for n, c := range allFiles {
		matchedFiles[n] = c
	}
	for name := range allFiles {
		match, err := bctx.MatchFile(".", name) // This will access the file we just added to files map above.
		if err != nil {
			return nil, &BadPackageError{Err: fmt.Errorf(`bctx.MatchFile(".", %q): %w`, name, err)}
		}
		if !match {
			delete(matchedFiles, name)
		}
	}
	return matchedFiles, nil
}

// readZipFile decompresses zip file f and returns its uncompressed contents.
// The caller can check f.UncompressedSize64 before calling readZipFile to
// get the expected uncompressed size of f.
//
// limit is the maximum number of bytes to read.
func readZipFile(f *zip.File, limit int64) (_ []byte, err error) {
	defer derrors.Add(&err, "readZipFile(%q)", f.Name)

	r, err := f.Open()
	if err != nil {
		return nil, fmt.Errorf("f.Open(): %v", err)
	}
	b, err := ioutil.ReadAll(io.LimitReader(r, limit))
	if err != nil {
		r.Close()
		return nil, fmt.Errorf("ioutil.ReadAll(r): %v", err)
	}
	if err := r.Close(); err != nil {
		return nil, fmt.Errorf("closing: %v", err)
	}
	return b, nil
}

// mib is the number of bytes in a mebibyte (Mi).
const mib = 1024 * 1024

// The largest module zip size we can comfortably process.
// We probably will OOM if we process a module whose zip is larger.
var maxModuleZipSize int64 = math.MaxInt64

func init() {
	v := config.GetEnvInt("GO_DISCOVERY_MAX_MODULE_ZIP_MI", -1)
	if v > 0 {
		maxModuleZipSize = int64(v) * mib
	}
}

var zipLoadShedder *loadShedder

func init() {
	ctx := context.Background()
	mebis := config.GetEnvInt("GO_DISCOVERY_MAX_IN_FLIGHT_ZIP_MI", -1)
	if mebis > 0 {
		log.Infof(ctx, "shedding load over %dMi", mebis)
		zipLoadShedder = &loadShedder{maxSizeInFlight: uint64(mebis) * mib}
	}
}

// ZipLoadShedStats returns a snapshot of the current LoadShedStats for zip files.
func ZipLoadShedStats() LoadShedStats {
	if zipLoadShedder != nil {
		return zipLoadShedder.stats()
	}
	return LoadShedStats{}
}
<commit_msg>internal/fetch: track build context file sets<commit_after>// Copyright 2019 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 fetch provides a way to fetch modules from a proxy.
package fetch

import (
	"archive/zip"
	"bytes"
	"context"
	"errors"
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"io"
	"io/ioutil"
	"math"
	"net/http"
	"os"
	"path"
	"sort"
	"strings"

	"go.opencensus.io/trace"
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/config"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/godoc"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/source"
	"golang.org/x/pkgsite/internal/stdlib"
)

// BadPackageError represents an error loading a package
// because its contents do not make up a valid package.
//
// This can happen, for example, if the .go files fail
// to parse or declare different package names.
type BadPackageError struct {
	Err error // Not nil.
}

func (bpe *BadPackageError) Error() string { return bpe.Err.Error() }

// loadPackage loads a Go package by calling loadPackageWithBuildContext, trying
// several build contexts in turn. It returns a goPackage with documentation
// information for each build context that results in a valid package, in the
// same order that the build contexts are listed. If none of them result in a
// package, then loadPackage returns nil, nil.
//
// If a package is fine except that its documentation is too large, loadPackage
// returns a goPackage whose err field is a non-nil error with godoc.ErrTooLarge in its chain.
func loadPackage(ctx context.Context, zipGoFiles []*zip.File, innerPath string, sourceInfo *source.Info, modInfo *godoc.ModuleInfo) (_ *goPackage, err error) {
	defer derrors.Wrap(&err, "loadPackage(ctx, zipGoFiles, %q, sourceInfo, modInfo)", innerPath)
	ctx, span := trace.StartSpan(ctx, "fetch.loadPackage")
	defer span.End()

	// Make a map with all the zip file contents.
	files := make(map[string][]byte)
	for _, f := range zipGoFiles {
		_, name := path.Split(f.Name)
		b, err := readZipFile(f, MaxFileSize)
		if err != nil {
			return nil, err
		}
		files[name] = b
	}

	modulePath := modInfo.ModulePath
	importPath := path.Join(modulePath, innerPath)
	if modulePath == stdlib.ModulePath {
		importPath = innerPath
	}
	v1path := internal.V1Path(importPath, modulePath)

	var pkg *goPackage
	// Parse the package for each build context.
	// The documentation is determined by the set of matching files, so keep
	// track of those to avoid duplication.
	docsByFiles := map[string]*internal.Documentation{}
	for _, bc := range internal.BuildContexts {
		mfiles, err := matchingFiles(bc.GOOS, bc.GOARCH, files)
		if err != nil {
			return nil, err
		}
		filesKey := mapKeyForFiles(mfiles)
		if doc := docsByFiles[filesKey]; doc != nil {
			// We have seen this set of files before.
			// loadPackageWithBuildContext will produce the same outputs,
			// so don't bother calling it. Just copy the doc.
			doc2 := *doc
			doc2.GOOS = bc.GOOS
			doc2.GOARCH = bc.GOARCH
			pkg.docs = append(pkg.docs, &doc2)
			continue
		}
		name, imports, synopsis, source, err := loadPackageWithBuildContext(ctx, bc.GOOS, bc.GOARCH, mfiles, innerPath, sourceInfo, modInfo)
		switch {
		case errors.Is(err, derrors.NotFound):
			// No package for this build context.
			continue
		case errors.Is(err, godoc.ErrTooLarge):
			// The doc for this build context is too large. Remember that and
			// return the package for this build context; ignore the others.
			return &goPackage{
				err:     err,
				path:    importPath,
				v1path:  v1path,
				name:    name,
				imports: imports,
				docs: []*internal.Documentation{{
					GOOS:     bc.GOOS,
					GOARCH:   bc.GOARCH,
					Synopsis: synopsis,
					Source:   source,
				}},
			}, nil
		case err != nil:
			// Serious error. Fail.
			return nil, err
		default:
			// No error.
			if pkg == nil {
				pkg = &goPackage{
					path:    importPath,
					v1path:  v1path,
					name:    name,
					imports: imports, // Use the imports from the first successful build context.
				}
			}
			// All the build contexts should use the same package name. Although
			// it's technically legal for different build tags to result in different
			// package names, it's not something we support.
			if name != pkg.name {
				return nil, &BadPackageError{
					Err: fmt.Errorf("more than one package name (%q and %q)", pkg.name, name),
				}
			}
			doc := &internal.Documentation{
				GOOS:     bc.GOOS,
				GOARCH:   bc.GOARCH,
				Synopsis: synopsis,
				Source:   source,
			}
			docsByFiles[filesKey] = doc
			pkg.docs = append(pkg.docs, doc)
		}
	}
	return pkg, nil
}

// mapKeyForFiles generates a value that corresponds to the given set of file
// names and can be used as a map key.
func mapKeyForFiles(files map[string][]byte) string {
	var names []string
	for n := range files {
		names = append(names, n)
	}
	sort.Strings(names)
	return strings.Join(names, " ")
}

// httpPost allows package fetch tests to stub out playground URL fetches.
var httpPost = http.Post

// loadPackageWithBuildContext loads a Go package made of .go files in
// files using a build context constructed from the given GOOS and GOARCH
// values. modulePath is stdlib.ModulePath for the Go standard library and the
// module path for all other modules. innerPath is the path of the Go package
// directory relative to the module root. The files argument must contain only
// .go files that have been verified to be of reasonable size and that match
// the build context.
//
// It returns the package name, list of imports, the package synopsis, and the
// serialized source (AST) for the package.
//
// It returns an error with NotFound in its chain if the directory doesn't
// contain a Go package or all .go files have been excluded by constraints. A
// *BadPackageError error is returned if the directory contains .go files but do
// not make up a valid package.
//
// If it returns an error with ErrTooLarge in its chain, the other return values
// are still valid.
func loadPackageWithBuildContext(ctx context.Context, goos, goarch string, files map[string][]byte, innerPath string, sourceInfo *source.Info, modInfo *godoc.ModuleInfo) (name string, imports []string, synopsis string, source []byte, err error) {
	modulePath := modInfo.ModulePath
	defer derrors.Wrap(&err, "loadPackageWithBuildContext(%q, %q, files, %q, %q, %+v)",
		goos, goarch, innerPath, modulePath, sourceInfo)

	packageName, goFiles, fset, err := loadFilesWithBuildContext(innerPath, files)
	if err != nil {
		return "", nil, "", nil, err
	}
	docPkg := godoc.NewPackage(fset, goos, goarch, modInfo.ModulePackages)
	for _, pf := range goFiles {
		removeNodes := true
		// Don't strip the seemingly unexported functions from the builtin package;
		// they are actually Go builtins like make, new, etc.
		if modulePath == stdlib.ModulePath && innerPath == "builtin" {
			removeNodes = false
		}
		docPkg.AddFile(pf, removeNodes)
	}

	// Encode first, because Render messes with the AST.
	src, err := docPkg.Encode(ctx)
	if err != nil {
		return "", nil, "", nil, err
	}

	synopsis, imports, _, err = docPkg.Render(ctx, innerPath, sourceInfo, modInfo, goos, goarch)
	if err != nil && !errors.Is(err, godoc.ErrTooLarge) {
		return "", nil, "", nil, err
	}
	return packageName, imports, synopsis, src, err
}

// loadFilesWithBuildContext loads all the given Go files at innerPath. It
// returns the package name as it occurs in the source, a map of the ASTs of all
// the Go files, and the token.FileSet used for parsing.
// If there are no non-test Go files, it returns a NotFound error.
func loadFilesWithBuildContext(innerPath string, files map[string][]byte) (pkgName string, fileMap map[string]*ast.File, _ *token.FileSet, _ error) {
	// Parse .go files and add them to the goFiles slice.
	var (
		fset            = token.NewFileSet()
		goFiles         = make(map[string]*ast.File)
		numNonTestFiles int
		packageName     string
		packageNameFile string // Name of file where packageName came from.
	)
	for name, b := range files {
		pf, err := parser.ParseFile(fset, name, b, parser.ParseComments)
		if err != nil {
			if pf == nil {
				return "", nil, nil, fmt.Errorf("internal error: the source couldn't be read: %v", err)
			}
			return "", nil, nil, &BadPackageError{Err: err}
		}
		// Remember all files, including test files for their examples.
		goFiles[name] = pf
		if strings.HasSuffix(name, "_test.go") {
			continue
		}
		// Keep track of the number of non-test files to check that the package name is the same.
		// and also because a directory with only test files doesn't count as a
		// Go package.
		numNonTestFiles++
		if numNonTestFiles == 1 {
			packageName = pf.Name.Name
			packageNameFile = name
		} else if pf.Name.Name != packageName {
			return "", nil, nil, &BadPackageError{Err: &build.MultiplePackageError{
				Dir:      innerPath,
				Packages: []string{packageName, pf.Name.Name},
				Files:    []string{packageNameFile, name},
			}}
		}
	}
	if numNonTestFiles == 0 {
		// This directory doesn't contain a package, or at least not one
		// that matches this build context.
		return "", nil, nil, derrors.NotFound
	}
	return packageName, goFiles, fset, nil
}

// matchingFiles returns a map from file names to their contents, read from zipGoFiles.
// It includes only those files that match the build context determined by goos and goarch.
func matchingFiles(goos, goarch string, allFiles map[string][]byte) (matchedFiles map[string][]byte, err error) {
	defer derrors.Wrap(&err, "matchingFiles(%q, %q, zipGoFiles)", goos, goarch)

	// bctx is used to make decisions about which of the .go files are included
	// by build constraints.
	bctx := &build.Context{
		GOOS:        goos,
		GOARCH:      goarch,
		CgoEnabled:  true,
		Compiler:    build.Default.Compiler,
		ReleaseTags: build.Default.ReleaseTags,

		JoinPath: path.Join,
		OpenFile: func(name string) (io.ReadCloser, error) {
			return ioutil.NopCloser(bytes.NewReader(allFiles[name])), nil
		},

		// If left nil, the default implementations of these read from disk,
		// which we do not want. None of these functions should be used
		// inside this function; it would be an internal error if they are.
		// Set them to non-nil values to catch if that happens.
		SplitPathList: func(string) []string { panic("internal error: unexpected call to SplitPathList") },
		IsAbsPath:     func(string) bool { panic("internal error: unexpected call to IsAbsPath") },
		IsDir:         func(string) bool { panic("internal error: unexpected call to IsDir") },
		HasSubdir:     func(string, string) (string, bool) { panic("internal error: unexpected call to HasSubdir") },
		ReadDir:       func(string) ([]os.FileInfo, error) { panic("internal error: unexpected call to ReadDir") },
	}

	// Copy the input map so we don't modify it.
	matchedFiles = map[string][]byte{}
	for n, c := range allFiles {
		matchedFiles[n] = c
	}
	for name := range allFiles {
		match, err := bctx.MatchFile(".", name) // This will access the file we just added to files map above.
		if err != nil {
			return nil, &BadPackageError{Err: fmt.Errorf(`bctx.MatchFile(".", %q): %w`, name, err)}
		}
		if !match {
			delete(matchedFiles, name)
		}
	}
	return matchedFiles, nil
}

// readZipFile decompresses zip file f and returns its uncompressed contents.
// The caller can check f.UncompressedSize64 before calling readZipFile to
// get the expected uncompressed size of f.
//
// limit is the maximum number of bytes to read.
func readZipFile(f *zip.File, limit int64) (_ []byte, err error) {
	defer derrors.Add(&err, "readZipFile(%q)", f.Name)

	r, err := f.Open()
	if err != nil {
		return nil, fmt.Errorf("f.Open(): %v", err)
	}
	b, err := ioutil.ReadAll(io.LimitReader(r, limit))
	if err != nil {
		r.Close()
		return nil, fmt.Errorf("ioutil.ReadAll(r): %v", err)
	}
	if err := r.Close(); err != nil {
		return nil, fmt.Errorf("closing: %v", err)
	}
	return b, nil
}

// mib is the number of bytes in a mebibyte (Mi).
const mib = 1024 * 1024

// The largest module zip size we can comfortably process.
// We probably will OOM if we process a module whose zip is larger.
var maxModuleZipSize int64 = math.MaxInt64

func init() {
	v := config.GetEnvInt("GO_DISCOVERY_MAX_MODULE_ZIP_MI", -1)
	if v > 0 {
		maxModuleZipSize = int64(v) * mib
	}
}

var zipLoadShedder *loadShedder

func init() {
	ctx := context.Background()
	mebis := config.GetEnvInt("GO_DISCOVERY_MAX_IN_FLIGHT_ZIP_MI", -1)
	if mebis > 0 {
		log.Infof(ctx, "shedding load over %dMi", mebis)
		zipLoadShedder = &loadShedder{maxSizeInFlight: uint64(mebis) * mib}
	}
}

// ZipLoadShedStats returns a snapshot of the current LoadShedStats for zip files.
func ZipLoadShedStats() LoadShedStats {
	if zipLoadShedder != nil {
		return zipLoadShedder.stats()
	}
	return LoadShedStats{}
}
<|endoftext|>


<commit_before>package deepcopier

import (
	"reflect"
	"strings"
)

// GetTagOptions parses deepcopier tag field and returns options.
func GetTagOptions(value string) map[string]string {
	options := map[string]string{}

	for _, opt := range strings.Split(value, ";") {
		o := strings.Split(opt, ":")

		// deepcopier:"keyword; without; value;"
		if len(o) == 1 {
			k := o[0]
			options[k] = ""
		}

		// deepcopier:"key:value; anotherkey:anothervalue"
		if len(o) == 2 {
			k, v := o[0], o[1]
			k = strings.TrimSpace(k)
			v = strings.TrimSpace(v)
			options[k] = v
		}
	}

	return options
}

// GetTypeMethods returns type methods.
func GetTypeMethods(t reflect.Type) []string {
	var methods []string
	for i := 0; i < t.NumMethod(); i++ {
		methods = append(methods, t.Method(i).Name)
	}
	return methods
}

// InStringSlice checks if the given string is in the given slice of string.
func InStringSlice(haystack []string, needle string) bool {
	for _, str := range haystack {
		if needle == str {
			return true
		}
	}
	return false
}
<commit_msg>Refacto: fix GetMethods() util.<commit_after>package deepcopier

import (
	"reflect"
	"strings"
)

// GetTagOptions parses deepcopier tag field and returns options.
func GetTagOptions(value string) map[string]string {
	options := map[string]string{}

	for _, opt := range strings.Split(value, ";") {
		o := strings.Split(opt, ":")

		// deepcopier:"keyword; without; value;"
		if len(o) == 1 {
			k := o[0]
			options[k] = ""
		}

		// deepcopier:"key:value; anotherkey:anothervalue"
		if len(o) == 2 {
			k, v := o[0], o[1]
			k = strings.TrimSpace(k)
			v = strings.TrimSpace(v)
			options[k] = v
		}
	}

	return options
}

// GetMethods returns type methods.
func GetMethods(instance interface{}) []string {
	t := reflect.TypeOf(instance)

	var methods []string
	for i := 0; i < t.NumMethod(); i++ {
		methods = append(methods, t.Method(i).Name)
	}

	return methods
}

// InStringSlice checks if the given string is in the given slice of string.
func InStringSlice(haystack []string, needle string) bool {
	for _, str := range haystack {
		if needle == str {
			return true
		}
	}

	return false
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/opencontainers/runtime-spec/specs-go"

	"github.com/sirupsen/logrus"
	"github.com/urfave/cli"
)

const (
	exactArgs = iota
	minArgs
	maxArgs
)

func checkArgs(context *cli.Context, expected, checkType int) error {
	var err error
	cmdName := context.Command.Name
	switch checkType {
	case exactArgs:
		if context.NArg() != expected {
			err = fmt.Errorf("%s: %q requires exactly %d argument(s)", os.Args[0], cmdName, expected)
		}
	case minArgs:
		if context.NArg() < expected {
			err = fmt.Errorf("%s: %q requires a minimum of %d argument(s)", os.Args[0], cmdName, expected)
		}
	case maxArgs:
		if context.NArg() > expected {
			err = fmt.Errorf("%s: %q requires a maximum of %d argument(s)", os.Args[0], cmdName, expected)
		}
	}

	if err != nil {
		fmt.Printf("Incorrect Usage.\n\n")
		cli.ShowCommandHelp(context, cmdName)
		return err
	}
	return nil
}

func logrusToStderr() bool {
	l, ok := logrus.StandardLogger().Out.(*os.File)
	return ok && l.Fd() == os.Stderr.Fd()
}

// fatal prints the error's details if it is a libcontainer specific error type
// then exits the program with an exit status of 1.
func fatal(err error) {
	// make sure the error is written to the logger
	logrus.Error(err)
	if !logrusToStderr() {
		fmt.Fprintln(os.Stderr, err)
	}

	os.Exit(1)
}

// setupSpec performs initial setup based on the cli.Context for the container
func setupSpec(context *cli.Context) (*specs.Spec, error) {
	bundle := context.String("bundle")
	if bundle != "" {
		if err := os.Chdir(bundle); err != nil {
			return nil, err
		}
	}
	spec, err := loadSpec(specConfig)
	if err != nil {
		return nil, err
	}
	return spec, nil
}

func revisePidFile(context *cli.Context) error {
	pidFile := context.String("pid-file")
	if pidFile == "" {
		return nil
	}

	// convert pid-file to an absolute path so we can write to the right
	// file after chdir to bundle
	pidFile, err := filepath.Abs(pidFile)
	if err != nil {
		return err
	}
	return context.Set("pid-file", pidFile)
}

// reviseRootDir convert the root to absolute path
func reviseRootDir(context *cli.Context) error {
	root := context.GlobalString("root")
	if root == "" {
		return nil
	}

	root, err := filepath.Abs(root)
	if err != nil {
		return err
	}

	return context.GlobalSet("root", root)
}

// parseBoolOrAuto returns (nil, nil) if s is empty or "auto"
func parseBoolOrAuto(s string) (*bool, error) {
	if s == "" || strings.ToLower(s) == "auto" {
		return nil, nil
	}
	b, err := strconv.ParseBool(s)
	return &b, err
}
<commit_msg>Show error stack trace if --debug is set<commit_after>package main

import (
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/opencontainers/runtime-spec/specs-go"

	"github.com/sirupsen/logrus"
	"github.com/urfave/cli"
)

const (
	exactArgs = iota
	minArgs
	maxArgs
)

func checkArgs(context *cli.Context, expected, checkType int) error {
	var err error
	cmdName := context.Command.Name
	switch checkType {
	case exactArgs:
		if context.NArg() != expected {
			err = fmt.Errorf("%s: %q requires exactly %d argument(s)", os.Args[0], cmdName, expected)
		}
	case minArgs:
		if context.NArg() < expected {
			err = fmt.Errorf("%s: %q requires a minimum of %d argument(s)", os.Args[0], cmdName, expected)
		}
	case maxArgs:
		if context.NArg() > expected {
			err = fmt.Errorf("%s: %q requires a maximum of %d argument(s)", os.Args[0], cmdName, expected)
		}
	}

	if err != nil {
		fmt.Printf("Incorrect Usage.\n\n")
		cli.ShowCommandHelp(context, cmdName)
		return err
	}
	return nil
}

func logrusToStderr() bool {
	l, ok := logrus.StandardLogger().Out.(*os.File)
	return ok && l.Fd() == os.Stderr.Fd()
}

// fatal prints the error's details if it is a libcontainer specific error type
// then exits the program with an exit status of 1.
func fatal(err error) {
	// make sure the error is written to the logger
	logrus.Error(err)
	// If debug is enabled and pkg/errors was used, show its stack trace.
	logrus.Debugf("%+v", err)
	if !logrusToStderr() {
		fmt.Fprintln(os.Stderr, err)
	}

	os.Exit(1)
}

// setupSpec performs initial setup based on the cli.Context for the container
func setupSpec(context *cli.Context) (*specs.Spec, error) {
	bundle := context.String("bundle")
	if bundle != "" {
		if err := os.Chdir(bundle); err != nil {
			return nil, err
		}
	}
	spec, err := loadSpec(specConfig)
	if err != nil {
		return nil, err
	}
	return spec, nil
}

func revisePidFile(context *cli.Context) error {
	pidFile := context.String("pid-file")
	if pidFile == "" {
		return nil
	}

	// convert pid-file to an absolute path so we can write to the right
	// file after chdir to bundle
	pidFile, err := filepath.Abs(pidFile)
	if err != nil {
		return err
	}
	return context.Set("pid-file", pidFile)
}

// reviseRootDir convert the root to absolute path
func reviseRootDir(context *cli.Context) error {
	root := context.GlobalString("root")
	if root == "" {
		return nil
	}

	root, err := filepath.Abs(root)
	if err != nil {
		return err
	}

	return context.GlobalSet("root", root)
}

// parseBoolOrAuto returns (nil, nil) if s is empty or "auto"
func parseBoolOrAuto(s string) (*bool, error) {
	if s == "" || strings.ToLower(s) == "auto" {
		return nil, nil
	}
	b, err := strconv.ParseBool(s)
	return &b, err
}
<|endoftext|>