<commit_before>// Copyright 2017-2019 the u-root 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 main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"time"

	"github.com/u-root/u-root/pkg/boot/systembooter"
	"github.com/u-root/u-root/pkg/ipmi"
	"github.com/u-root/u-root/pkg/ipmi/ocp"
	"github.com/u-root/u-root/pkg/smbios"
	"github.com/u-root/u-root/pkg/vpd"
)

var (
	allowInteractive = flag.Bool("i", true, "Allow user to interrupt boot process and run commands")
	doQuiet          = flag.Bool("q", false, fmt.Sprintf("Disable verbose output. If not specified, read it from VPD var '%s'. Default false", vpdSystembootLogLevel))
	interval         = flag.Int("I", 1, "Interval in seconds before looping to the next boot command")
	noDefaultBoot    = flag.Bool("nodefault", false, "Do not attempt default boot entries if regular ones fail")
)

const (
	// vpdSystembootLogLevel is the name of the VPD variable used to set the log level.
	vpdSystembootLogLevel = "systemboot_log_level"
)

// isFlagPassed checks whether a flag was explicitly passed on the command line
func isFlagPassed(name string) bool {
	found := false
	flag.Visit(func(f *flag.Flag) {
		if f.Name == name {
			found = true
		}
	})
	return found
}

var defaultBootsequence = [][]string{
	{"fbnetboot", "-userclass", "linuxboot"},
	{"localboot", "-grub"},
}

// VPD variable for enabling IPMI BMC overriding boot order, default is not set
const VpdBmcBootOrderOverride = "bmc_bootorder_override"

var bmcBootOverride bool

// Product list for running IPMI OEM commands
var productList = [4]string{"Tioga Pass", "Mono Lake", "Delta Lake", "CraterLake"}

var selRecorded bool

func isMatched(productName string) bool {
	for _, v := range productList {
		if strings.HasPrefix(productName, v) {
			return true
		}
	}
	return false
}

func getBaseboardProductName(si *smbios.Info) (string, error) {
	t2, err := si.GetBaseboardInfo()
	if err != nil {
		log.Printf("Error getting Baseboard Information: %v", err)
		return "", err
	}
	return t2[0].Product, nil
}

func getSystemFWVersion(si *smbios.Info) (string, error) {
	t0, err := si.GetBIOSInfo()
	if err != nil {
		log.Printf("Error getting BIOS Information: %v", err)
		return "", err
	}
	return t0.Version, nil
}

func checkCMOSClear(ipmi *ipmi.IPMI) error {
	if cmosclear, bootorder, err := ocp.IsCMOSClearSet(ipmi); cmosclear {
		log.Printf("CMOS clear starts")
		if err = cmosClear(); err != nil {
			return err
		}
		if err = vpd.ClearRwVpd(); err != nil {
			return err
		}

		if err = ocp.ClearCMOSClearValidBits(ipmi, bootorder); err != nil {
			return err
		}
		addSEL("cmosclear")
		if err = reboot(); err != nil {
			return err
		}
	} else if err != nil {
		return err
	}

	return nil
}

func runIPMICommands() {
	i, err := ipmi.Open(0)
	if err != nil {
		log.Printf("Failed to open ipmi device %v, watchdog may still be running", err)
		return
	}
	defer i.Close()

	if err = i.ShutoffWatchdog(); err != nil {
		log.Printf("Failed to stop watchdog %v.", err)
	} else {
		log.Printf("Watchdog is stopped.")
	}
	// Try RW_VPD first
	value, err := systembooter.Get(VpdBmcBootOrderOverride, false)
	if err != nil {
		// Try RO_VPD
		value, err = systembooter.Get(VpdBmcBootOrderOverride, true)
	}
	if err == nil && string(value) == "1" {
		bmcBootOverride = true
	}
	log.Printf("VPD %s is %v", VpdBmcBootOrderOverride, string(value))
	// Below IPMI commands would require SMBIOS data
	si, err := smbios.FromSysfs()
	if err != nil {
		log.Printf("Error reading SMBIOS info: %v", err)
		return
	}

	if fwVersion, err := getSystemFWVersion(si); err == nil {
		log.Printf("System firmware version: %s", fwVersion)
		if err = i.SetSystemFWVersion(fwVersion); err != nil {
			log.Printf("Failed to set system firmware version to BMC %v.", err)
		}
	}

	if productName, err := getBaseboardProductName(si); err == nil {
		if isMatched(productName) {
			log.Printf("Running OEM IPMI commands.")
			if err = checkCMOSClear(i); err != nil {
				log.Printf("IPMI CMOS clear err: %v", err)
			}
			if err = ocp.CheckBMCBootOrder(i, bmcBootOverride); err != nil {
				log.Printf("Failed to sync BMC Boot Order %v.", err)
			}
			dimmInfo, err := ocp.GetOemIpmiDimmInfo(si)
			if err == nil {
				if err = ocp.SendOemIpmiDimmInfo(i, dimmInfo); err == nil {
					log.Printf("Send the information of DIMMs to BMC.")
				} else {
					log.Printf("Failed to send the information of DIMMs to BMC: %v.", err)
				}
			} else {
				log.Printf("Failed to get the information of DIMMs: %v.", err)
			}

			processorInfo, err := ocp.GetOemIpmiProcessorInfo(si)
			if err == nil {
				if err = ocp.SendOemIpmiProcessorInfo(i, processorInfo); err == nil {
					log.Printf("Send the information of processors to BMC.")
				} else {
					log.Printf("Failed to send the information of processors to BMC: %v.", err)
				}
			} else {
				log.Printf("Failed to get the information of Processors: %v.", err)
			}

			if err = ocp.SetOemIpmiPostEnd(i); err == nil {
				log.Printf("Send IPMI POST end to BMC")
			} else {
				log.Printf("Failed to send IPMI POST end to BMC: %v.", err)
			}

		} else {
			log.Printf("No product name is matched for OEM commands.")
		}
	}
}

func addSEL(sequence string) {
	var bootErr ipmi.Event

	i, err := ipmi.Open(0)
	if err != nil {
		log.Printf("Failed to open ipmi device to send SEL %v", err)
		return
	}
	defer i.Close()

	switch sequence {
	case "netboot":
		fallthrough
	case "fbnetboot":
		bootErr.RecordID = 0
		bootErr.RecordType = ipmi.OEM_NTS_TYPE
		bootErr.OEMNontsDefinedData[0] = 0x28
		bootErr.OEMNontsDefinedData[5] = 0xf0
		for idx := 6; idx < 13; idx++ {
			bootErr.OEMNontsDefinedData[idx] = 0xff
		}
		if err := i.LogSystemEvent(&bootErr); err != nil {
			log.Printf("SEL recorded: %s fail\n", sequence)
		}
	case "cmosclear":
		bootErr.RecordID = 0
		bootErr.RecordType = ipmi.OEM_NTS_TYPE
		bootErr.OEMNontsDefinedData[0] = 0x28
		bootErr.OEMNontsDefinedData[5] = 0xf1
		for idx := 6; idx < 13; idx++ {
			bootErr.OEMNontsDefinedData[idx] = 0xff
		}
		if err := i.LogSystemEvent(&bootErr); err != nil {
			log.Printf("SEL recorded: %s fail\n", sequence)
		}
	default:
	}
}

// getDebugEnabled checks whether debug output is requested, either via command line or via VPD
// variables.
// If -q was explicitly passed on the command line, will use that value, otherwise will look for
// the VPD variable "systemboot_log_level".
// Valid values are coreboot loglevels https://review.coreboot.org/cgit/coreboot.git/tree/src/commonlib/include/commonlib/loglevel.h,
// either as integer (1, 2) or string ("debug").
// If the VPD variable is missing or it is set to an invalid value, it will use the default.
func getDebugEnabled() bool {
	if isFlagPassed("q") {
		return !*doQuiet
	}

	// -q was not passed, so `doQuiet` contains the default value
	defaultDebugEnabled := !*doQuiet
	// check for the VPD variable "systemboot_log_level". First the read-write, then the read-only
	v, err := vpd.Get(vpdSystembootLogLevel, false)
	if err != nil {
		// TODO do not print warning if file is not found
		log.Printf("Warning: failed to read read-write VPD variable '%s', will try the read-only one. Error was: %v", vpdSystembootLogLevel, err)
		v, err = vpd.Get(vpdSystembootLogLevel, true)
		if err != nil {
			// TODO do not print warning if file is not found
			log.Printf("Warning: failed to read read-only VPD variable '%s', will use the default value. Error was: %v", vpdSystembootLogLevel, err)
			return defaultDebugEnabled
		}
	}
	level := strings.ToLower(strings.TrimSpace(string(v)))
	switch level {
	case "0", "emerg", "1", "alert", "2", "crit", "3", "err", "4", "warning", "9", "never":
		// treat as quiet
		return false
	case "5", "notice", "6", "info", "7", "debug", "8", "spew":
		// treat as debug
		return true
	default:
		log.Printf("Invalid value '%s' for VPD variable '%s', using default", level, vpdSystembootLogLevel)
		return defaultDebugEnabled
	}
}

func main() {
	flag.Parse()

	debugEnabled := getDebugEnabled()

	log.Print(`
                     ____            _                 _                 _
                    / ___| _   _ ___| |_ ___ _ __ ___ | |__   ___   ___ | |_
                    \___ \| | | / __| __/ _ \ '_ ` + "`" + ` _ \| '_ \ / _ \ / _ \| __|
                     ___) | |_| \__ \ ||  __/ | | | | | |_) | (_) | (_) | |_
                    |____/ \__, |___/\__\___|_| |_| |_|_.__/ \___/ \___/ \__|
                           |___/
`)
	runIPMICommands()
	sleepInterval := time.Duration(*interval) * time.Second
	if *allowInteractive {
		log.Printf("**************************************************************************")
		log.Print("Starting boot sequence, press CTRL-C within 5 seconds to drop into a shell")
		log.Printf("**************************************************************************")
		time.Sleep(5 * time.Second)
	} else {
		signal.Ignore()
	}

	// Get and show boot entries
	var bootEntries []systembooter.BootEntry
	if bmcBootOverride && ocp.BmcUpdatedBootorder {
		bootEntries = ocp.BootEntries
	} else {
		bootEntries = systembooter.GetBootEntries()
	}
	log.Printf("BOOT ENTRIES:")
	for _, entry := range bootEntries {
		log.Printf("    %v) %+v", entry.Name, string(entry.Config))
	}
	for _, entry := range bootEntries {
		log.Printf("Trying boot entry %s: %s", entry.Name, string(entry.Config))
		if err := entry.Booter.Boot(debugEnabled); err != nil {
			log.Printf("Warning: failed to boot with configuration: %+v", entry)
			addSEL(entry.Booter.TypeName())
		}
		if debugEnabled {
			log.Printf("Sleeping %v before attempting next boot command", sleepInterval)
		}
		time.Sleep(sleepInterval)
	}

	// if boot entries failed, use the default boot sequence
	log.Printf("Boot entries failed")

	if !*noDefaultBoot {
		log.Print("Falling back to the default boot sequence")
		for {
			for _, bootcmd := range defaultBootsequence {
				if debugEnabled {
					bootcmd = append(bootcmd, "-d")
				}
				log.Printf("Running boot command: %v", bootcmd)
				cmd := exec.Command(bootcmd[0], bootcmd[1:]...)
				cmd.Stdout = os.Stdout
				cmd.Stderr = os.Stderr
				if err := cmd.Run(); err != nil {
					log.Printf("Error executing %v: %v", cmd, err)
					if !selRecorded {
						addSEL(bootcmd[0])
					}
				}
			}
			selRecorded = true

			if debugEnabled {
				log.Printf("Sleeping %v before attempting next boot command", sleepInterval)
			}
			time.Sleep(sleepInterval)
		}
	}
}
<commit_msg>systemboot: add Quanta S9S to supported OEM product list<commit_after>// Copyright 2017-2019 the u-root 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 main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"time"

	"github.com/u-root/u-root/pkg/boot/systembooter"
	"github.com/u-root/u-root/pkg/ipmi"
	"github.com/u-root/u-root/pkg/ipmi/ocp"
	"github.com/u-root/u-root/pkg/smbios"
	"github.com/u-root/u-root/pkg/vpd"
)

var (
	allowInteractive = flag.Bool("i", true, "Allow user to interrupt boot process and run commands")
	doQuiet          = flag.Bool("q", false, fmt.Sprintf("Disable verbose output. If not specified, read it from VPD var '%s'. Default false", vpdSystembootLogLevel))
	interval         = flag.Int("I", 1, "Interval in seconds before looping to the next boot command")
	noDefaultBoot    = flag.Bool("nodefault", false, "Do not attempt default boot entries if regular ones fail")
)

const (
	// vpdSystembootLogLevel is the name of the VPD variable used to set the log level.
	vpdSystembootLogLevel = "systemboot_log_level"
)

// isFlagPassed checks whether a flag was explicitly passed on the command line
func isFlagPassed(name string) bool {
	found := false
	flag.Visit(func(f *flag.Flag) {
		if f.Name == name {
			found = true
		}
	})
	return found
}

var defaultBootsequence = [][]string{
	{"fbnetboot", "-userclass", "linuxboot"},
	{"localboot", "-grub"},
}

// VPD variable for enabling IPMI BMC overriding boot order, default is not set
const VpdBmcBootOrderOverride = "bmc_bootorder_override"

var bmcBootOverride bool

// Product list for running IPMI OEM commands
var productList = [5]string{"Tioga Pass", "Mono Lake", "Delta Lake", "CraterLake", "S9S"}

var selRecorded bool

func isMatched(productName string) bool {
	for _, v := range productList {
		if strings.HasPrefix(productName, v) {
			return true
		}
	}
	return false
}

func getBaseboardProductName(si *smbios.Info) (string, error) {
	t2, err := si.GetBaseboardInfo()
	if err != nil {
		log.Printf("Error getting Baseboard Information: %v", err)
		return "", err
	}
	return t2[0].Product, nil
}

func getSystemFWVersion(si *smbios.Info) (string, error) {
	t0, err := si.GetBIOSInfo()
	if err != nil {
		log.Printf("Error getting BIOS Information: %v", err)
		return "", err
	}
	return t0.Version, nil
}

func checkCMOSClear(ipmi *ipmi.IPMI) error {
	if cmosclear, bootorder, err := ocp.IsCMOSClearSet(ipmi); cmosclear {
		log.Printf("CMOS clear starts")
		if err = cmosClear(); err != nil {
			return err
		}
		if err = vpd.ClearRwVpd(); err != nil {
			return err
		}

		if err = ocp.ClearCMOSClearValidBits(ipmi, bootorder); err != nil {
			return err
		}
		addSEL("cmosclear")
		if err = reboot(); err != nil {
			return err
		}
	} else if err != nil {
		return err
	}

	return nil
}

func runIPMICommands() {
	i, err := ipmi.Open(0)
	if err != nil {
		log.Printf("Failed to open ipmi device %v, watchdog may still be running", err)
		return
	}
	defer i.Close()

	if err = i.ShutoffWatchdog(); err != nil {
		log.Printf("Failed to stop watchdog %v.", err)
	} else {
		log.Printf("Watchdog is stopped.")
	}
	// Try RW_VPD first
	value, err := systembooter.Get(VpdBmcBootOrderOverride, false)
	if err != nil {
		// Try RO_VPD
		value, err = systembooter.Get(VpdBmcBootOrderOverride, true)
	}
	if err == nil && string(value) == "1" {
		bmcBootOverride = true
	}
	log.Printf("VPD %s is %v", VpdBmcBootOrderOverride, string(value))
	// Below IPMI commands would require SMBIOS data
	si, err := smbios.FromSysfs()
	if err != nil {
		log.Printf("Error reading SMBIOS info: %v", err)
		return
	}

	if fwVersion, err := getSystemFWVersion(si); err == nil {
		log.Printf("System firmware version: %s", fwVersion)
		if err = i.SetSystemFWVersion(fwVersion); err != nil {
			log.Printf("Failed to set system firmware version to BMC %v.", err)
		}
	}

	if productName, err := getBaseboardProductName(si); err == nil {
		if isMatched(productName) {
			log.Printf("Running OEM IPMI commands.")
			if err = checkCMOSClear(i); err != nil {
				log.Printf("IPMI CMOS clear err: %v", err)
			}
			if err = ocp.CheckBMCBootOrder(i, bmcBootOverride); err != nil {
				log.Printf("Failed to sync BMC Boot Order %v.", err)
			}
			dimmInfo, err := ocp.GetOemIpmiDimmInfo(si)
			if err == nil {
				if err = ocp.SendOemIpmiDimmInfo(i, dimmInfo); err == nil {
					log.Printf("Send the information of DIMMs to BMC.")
				} else {
					log.Printf("Failed to send the information of DIMMs to BMC: %v.", err)
				}
			} else {
				log.Printf("Failed to get the information of DIMMs: %v.", err)
			}

			processorInfo, err := ocp.GetOemIpmiProcessorInfo(si)
			if err == nil {
				if err = ocp.SendOemIpmiProcessorInfo(i, processorInfo); err == nil {
					log.Printf("Send the information of processors to BMC.")
				} else {
					log.Printf("Failed to send the information of processors to BMC: %v.", err)
				}
			} else {
				log.Printf("Failed to get the information of Processors: %v.", err)
			}

			if err = ocp.SetOemIpmiPostEnd(i); err == nil {
				log.Printf("Send IPMI POST end to BMC")
			} else {
				log.Printf("Failed to send IPMI POST end to BMC: %v.", err)
			}

		} else {
			log.Printf("No product name is matched for OEM commands.")
		}
	}
}

func addSEL(sequence string) {
	var bootErr ipmi.Event

	i, err := ipmi.Open(0)
	if err != nil {
		log.Printf("Failed to open ipmi device to send SEL %v", err)
		return
	}
	defer i.Close()

	switch sequence {
	case "netboot":
		fallthrough
	case "fbnetboot":
		bootErr.RecordID = 0
		bootErr.RecordType = ipmi.OEM_NTS_TYPE
		bootErr.OEMNontsDefinedData[0] = 0x28
		bootErr.OEMNontsDefinedData[5] = 0xf0
		for idx := 6; idx < 13; idx++ {
			bootErr.OEMNontsDefinedData[idx] = 0xff
		}
		if err := i.LogSystemEvent(&bootErr); err != nil {
			log.Printf("SEL recorded: %s fail\n", sequence)
		}
	case "cmosclear":
		bootErr.RecordID = 0
		bootErr.RecordType = ipmi.OEM_NTS_TYPE
		bootErr.OEMNontsDefinedData[0] = 0x28
		bootErr.OEMNontsDefinedData[5] = 0xf1
		for idx := 6; idx < 13; idx++ {
			bootErr.OEMNontsDefinedData[idx] = 0xff
		}
		if err := i.LogSystemEvent(&bootErr); err != nil {
			log.Printf("SEL recorded: %s fail\n", sequence)
		}
	default:
	}
}

// getDebugEnabled checks whether debug output is requested, either via command line or via VPD
// variables.
// If -q was explicitly passed on the command line, will use that value, otherwise will look for
// the VPD variable "systemboot_log_level".
// Valid values are coreboot loglevels https://review.coreboot.org/cgit/coreboot.git/tree/src/commonlib/include/commonlib/loglevel.h,
// either as integer (1, 2) or string ("debug").
// If the VPD variable is missing or it is set to an invalid value, it will use the default.
func getDebugEnabled() bool {
	if isFlagPassed("q") {
		return !*doQuiet
	}

	// -q was not passed, so `doQuiet` contains the default value
	defaultDebugEnabled := !*doQuiet
	// check for the VPD variable "systemboot_log_level". First the read-write, then the read-only
	v, err := vpd.Get(vpdSystembootLogLevel, false)
	if err != nil {
		// TODO do not print warning if file is not found
		log.Printf("Warning: failed to read read-write VPD variable '%s', will try the read-only one. Error was: %v", vpdSystembootLogLevel, err)
		v, err = vpd.Get(vpdSystembootLogLevel, true)
		if err != nil {
			// TODO do not print warning if file is not found
			log.Printf("Warning: failed to read read-only VPD variable '%s', will use the default value. Error was: %v", vpdSystembootLogLevel, err)
			return defaultDebugEnabled
		}
	}
	level := strings.ToLower(strings.TrimSpace(string(v)))
	switch level {
	case "0", "emerg", "1", "alert", "2", "crit", "3", "err", "4", "warning", "9", "never":
		// treat as quiet
		return false
	case "5", "notice", "6", "info", "7", "debug", "8", "spew":
		// treat as debug
		return true
	default:
		log.Printf("Invalid value '%s' for VPD variable '%s', using default", level, vpdSystembootLogLevel)
		return defaultDebugEnabled
	}
}

func main() {
	flag.Parse()

	debugEnabled := getDebugEnabled()

	log.Print(`
                     ____            _                 _                 _
                    / ___| _   _ ___| |_ ___ _ __ ___ | |__   ___   ___ | |_
                    \___ \| | | / __| __/ _ \ '_ ` + "`" + ` _ \| '_ \ / _ \ / _ \| __|
                     ___) | |_| \__ \ ||  __/ | | | | | |_) | (_) | (_) | |_
                    |____/ \__, |___/\__\___|_| |_| |_|_.__/ \___/ \___/ \__|
                           |___/
`)
	runIPMICommands()
	sleepInterval := time.Duration(*interval) * time.Second
	if *allowInteractive {
		log.Printf("**************************************************************************")
		log.Print("Starting boot sequence, press CTRL-C within 5 seconds to drop into a shell")
		log.Printf("**************************************************************************")
		time.Sleep(5 * time.Second)
	} else {
		signal.Ignore()
	}

	// Get and show boot entries
	var bootEntries []systembooter.BootEntry
	if bmcBootOverride && ocp.BmcUpdatedBootorder {
		bootEntries = ocp.BootEntries
	} else {
		bootEntries = systembooter.GetBootEntries()
	}
	log.Printf("BOOT ENTRIES:")
	for _, entry := range bootEntries {
		log.Printf("    %v) %+v", entry.Name, string(entry.Config))
	}
	for _, entry := range bootEntries {
		log.Printf("Trying boot entry %s: %s", entry.Name, string(entry.Config))
		if err := entry.Booter.Boot(debugEnabled); err != nil {
			log.Printf("Warning: failed to boot with configuration: %+v", entry)
			addSEL(entry.Booter.TypeName())
		}
		if debugEnabled {
			log.Printf("Sleeping %v before attempting next boot command", sleepInterval)
		}
		time.Sleep(sleepInterval)
	}

	// if boot entries failed, use the default boot sequence
	log.Printf("Boot entries failed")

	if !*noDefaultBoot {
		log.Print("Falling back to the default boot sequence")
		for {
			for _, bootcmd := range defaultBootsequence {
				if debugEnabled {
					bootcmd = append(bootcmd, "-d")
				}
				log.Printf("Running boot command: %v", bootcmd)
				cmd := exec.Command(bootcmd[0], bootcmd[1:]...)
				cmd.Stdout = os.Stdout
				cmd.Stderr = os.Stderr
				if err := cmd.Run(); err != nil {
					log.Printf("Error executing %v: %v", cmd, err)
					if !selRecorded {
						addSEL(bootcmd[0])
					}
				}
			}
			selRecorded = true

			if debugEnabled {
				log.Printf("Sleeping %v before attempting next boot command", sleepInterval)
			}
			time.Sleep(sleepInterval)
		}
	}
}
<|endoftext|>


<commit_before>package model

import (
	"database/sql"
	_ "github.com/mattn/go-sqlite3"
)

var metrics Metrics

func InitMetrics() {
	metrics.init()
}

func GetMetrics() (metrics *Metrics) {
	return &metrics
}

const (
	SQL_SELECT_METRICS         = `SELECT id, name, type, detector, md5 FROM metrics;`
	SQL_SELECT_METRICS_DEFAULT = `SELECT id, name, params, interval
	                              FROM default_metrics AS d, metrics AS m
	                              WHERE d.id = m.id;`
	SQL_INSERT_METRIC      = `INSERT INTO metrics (name, type, detector, md5) VALUES(?,?,?,?);`
	SQL_UPDATE_METRIC_NAME = `UPDATE metrics SET name = ? WHERE id = ?;`
	SQL_UPDATE_METRIC      = `UPDATE metrics SET type = ?, detector = ?, md5 = ? WHERE id = ?;`
)

type Metric struct {
	ID       int64
	Name     string
	Type     string
	Detector string
	MD5      string
}

type DefaultMetric struct {
	ID       int64
	Name     string
	Params   string
	Interval int
}

type Metrics struct {
	metrics  map[int64]Metric
	defaults map[int64]DefaultMetric
	db       *sql.DB
}

func (self *Metrics) init() {

	self.metrics = make(map[int64]Metric)
	self.defaults = make(map[int64]DefaultMetric)

	rows, err := self.db.Query(SQL_SELECT_METRICS)

	if err != nil {
		panic(err)
	}

	defer rows.Close()

	for rows.Next() {

		var metric Metric

		err = rows.Scan(&metric.ID, &metric.Name, &metric.Type, &metric.Detector, &metric.MD5)

		if err != nil {
			panic(err)
		}

		self.metrics[metric.ID] = metric
	}

	err = rows.Err()

	if err != nil {
		panic(err)
	}

	rows, err = self.db.Query(SQL_SELECT_METRICS_DEFAULT)

	if err != nil {
		panic(err)
	}

	defer rows.Close()

	for rows.Next() {

		var metric DefaultMetric

		err = rows.Scan(&metric.ID, &metric.Name, &metric.Params, &metric.Interval)

		if err != nil {
			panic(err)
		}

		self.defaults[metric.ID] = metric
	}

	err = rows.Err()

	if err != nil {
		panic(err)
	}
}

func (self *Metrics) All() (metrics map[int64]Metric) {
	return self.metrics
}

func (self *Metrics) AllDefault() (dfts map[int64]DefaultMetric) {
	return self.defaults
}

func (self *Metrics) AddMetric(name, t, detector, md5 string) (metric Metric, err error) {

	metric = Metric{Name: name, Type: t, Detector: detector, MD5: md5}

	result, err := self.db.Exec(SQL_INSERT_METRIC, name, t, detector, md5)

	if err != nil {
		return
	}

	metric.ID, err = result.LastInsertId()
	return
}

func (self *Metrics) RenameMetric(id int64, name string) (err error) {

	result, err := self.db.Exec(SQL_UPDATE_METRIC_NAME, name, id)

	if err != nil {
		return
	}

	c, err := r.RowsAffected()

	if err != nil {
		return
	}

	if c < 1 {
		err = errors.New("No such metric: " + strconv.FormatInt(id, 10))
	}

	return
}

func (self *Metrics) UpdateMetric(id int64, t, detector, md5 string) (err error) {

	result, err := self.db.Exec(SQL_UPDATE_METRIC, t, detector, md5, id)

	if err != nil {
		return
	}

	c, err := r.RowsAffected()

	if err != nil {
		return
	}

	if c < 1 {
		err = errors.New("No such metric: " + strconv.FormatInt(id, 10))
	}

	return
}
<commit_msg>metrics model<commit_after>package model

import (
	"database/sql"
	_ "github.com/mattn/go-sqlite3"
)

var metrics Metrics

func GetMetrics() (metrics *Metrics) {
	return &metrics
}

const (
	SQL_SELECT_METRICS      = `SELECT id, name, type, detector, md5 FROM metrics;`
	SQL_SELECT_METRIC_BY_ID = `SELECT id, name, type, detector, md5 FROM metrics WHERE id = ?;`
	SQL_SELECT_METRICS_DFT  = `SELECT id, name, params, interval
	                            FROM default_metrics AS d, metrics AS m
	                           WHERE d.id = m.id;`
	SQL_INSERT_METRIC      = `INSERT INTO metrics (name, type, detector, md5) VALUES(?,?,?,?);`
	SQL_UPDATE_METRIC_NAME = `UPDATE metrics SET name = ? WHERE id = ?;`
	SQL_UPDATE_METRIC      = `UPDATE metrics SET type = ?, detector = ?, md5 = ? WHERE id = ?;`
	SQL_DELETE_METRIC      = `DELETE FROM metrics WHERE id = ?;`
)

type Metric struct {
	ID       int64
	Name     string
	Type     string
	Detector string
	MD5      string
}

type DefaultMetric struct {
	ID       int64
	Name     string
	Params   string
	Interval int
}

type Metrics struct {
}

func (self *Metrics) All() (metrics map[int64]Metric, err error) {

	metrics = make(map[int64]Metric)

	rows, err := db.Query(SQL_SELECT_METRICS)

	if err != nil {
		return
	}

	defer rows.Close()

	for rows.Next() {

		var metric Metric

		err = rows.Scan(&metric.ID, &metric.Name, &metric.Type, &metric.Detector, &metric.MD5)

		if err != nil {
			return
		}

		metrics[metric.ID] = metric
	}

	err = rows.Err()

	return
}

func (self *Metrics) AllDefault() (dfts map[int64]DefaultMetric, err error) {

	defaults = make(map[int64]DefaultMetric)

	rows, err = db.Query(SQL_SELECT_METRICS_DEFAULT)

	if err != nil {
		return
	}

	defer rows.Close()

	for rows.Next() {

		var metric DefaultMetric

		err = rows.Scan(&metric.ID, &metric.Name, &metric.Params, &metric.Interval)

		if err != nil {
			return
		}

		self.defaults[metric.ID] = metric
	}

	err = rows.Err()

	return
}

func (self *Metrics) AddMetric(name, t, detector, md5 string) (metric Metric, err error) {

	metric = Metric{Name: name, Type: t, Detector: detector, MD5: md5}

	result, err := db.Exec(SQL_INSERT_METRIC, name, t, detector, md5)

	if err != nil {
		return
	}

	metric.ID, err = result.LastInsertId()
	return
}

func (self *Metrics) RenameMetric(id int64, name string) (err error) {

	result, err := db.Exec(SQL_UPDATE_METRIC_NAME, name, id)

	if err != nil {
		return
	}

	c, err := result.RowsAffected()

	if err != nil {
		return
	}

	if c < 1 {
		err = errors.New("No such metric: " + strconv.FormatInt(id, 10))
	}

	return
}

func (self *Metrics) UpdateMetric(id int64, t, detector, md5 string) (err error) {

	result, err := db.Exec(SQL_UPDATE_METRIC, t, detector, md5, id)

	if err != nil {
		return
	}

	c, err := result.RowsAffected()

	if err != nil {
		return
	}

	if c < 1 {
		err = errors.New("No such metric: " + strconv.FormatInt(id, 10))
	}

	return
}

func (self *Metrics) DelMetric(id int64) (err error) {

	result, err := db.Exec(SQL_DELETE_METRIC, id)

	if err != nil {
		return
	}

	c, err := result.RowsAffected()

	if err != nil {
		return
	}

	if c < 1 {
		err = errors.New("No such metric: " + strconv.FormatInt(id, 10))
	}

	return
}

func (self *Metrics) GetMetric(id int64) (metric Metric, err error) {

	err = db.QueryRow(SQL_SELECT_METRIC_BY_ID, id).Scan(&metric.ID, &metric.Name,
		&metric.Type, &metric.Detector, &metric.MD5)

	return
}
<|endoftext|>


<commit_before>// +build ebpf

/*
 * Copyright (C) 2017 Red Hat, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy ofthe License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specificlanguage governing permissions and
 * limitations under the License.
 *
 */

package ebpf

import (
	"bytes"
	"errors"
	"fmt"
	"math"
	"runtime"
	"sync"
	"syscall"
	"time"
	"unsafe"

	"github.com/newtools/ebpf"

	"github.com/skydive-project/skydive/api/types"
	"github.com/skydive-project/skydive/ebpf/statics"
	"github.com/skydive-project/skydive/flow"
	"github.com/skydive-project/skydive/flow/probes"
	"github.com/skydive-project/skydive/graffiti/graph"
	"github.com/skydive-project/skydive/probe"
	"github.com/skydive-project/skydive/rawsocket"
	"github.com/skydive-project/skydive/topology"
)

/*
#cgo CFLAGS: -I../../../ebpf
#include "flow.h"
#include <string.h>
#include <sys/resource.h>
#include <sys/socket.h>

int probe_bpf_attach_socket(int sock, int fd)
{
	return setsockopt(sock, SOL_SOCKET, SO_ATTACH_BPF, &fd, sizeof(fd));
}

int probe_bpf_detach_socket(int sock, int fd)
{
	return setsockopt(sock, SOL_SOCKET, SO_DETACH_BPF, &fd, sizeof(fd));
}
*/
import "C"

// Probe the eBPF probe
type Probe struct {
	Ctx          probes.Context
	probeNodeTID string
	fd           int
	flowTable    *flow.Table
	module       *ebpf.Collection
	fmap         []*ebpf.Map
	cmap         *ebpf.Map
	expire       time.Duration
	quit         chan bool
	flowPage     int
}

// ProbesHandler creates new eBPF probes
type ProbesHandler struct {
	Ctx probes.Context
	wg  sync.WaitGroup
}

func (p *Probe) swapPage() {
	key := uint32(C.FLOW_PAGE)
	writeInPage := int64(p.flowPage)

	p.cmap.Put(key, writeInPage)

	if p.flowPage == 1 {
		p.flowPage = 0
	} else {
		p.flowPage = 1
	}
}

func (p *Probe) run() {
	var info syscall.Sysinfo_t
	syscall.Sysinfo(&info)

	_, extFlowChan, statsChan := p.flowTable.Start(nil)
	defer p.flowTable.Stop()

	ebpfPollingRate := time.Second / time.Duration(p.Ctx.Config.GetInt("agent.flow.ebpf.polling_rate"))
	const ebpfMaxPollDelay = 10 * time.Second

	var startKTimeNs int64
	var start time.Time

	updateNow := time.NewTicker(ebpfMaxPollDelay)
	defer updateNow.Stop()

	flowPoolSize := 2 * cap(extFlowChan)
	kernFlows := make([]C.struct_flow, flowPoolSize)

	extFlows := make([]flow.ExtFlow, flowPoolSize)
	for i := range extFlows {
		extFlows[i] = flow.ExtFlow{
			Type: flow.EBPFExtFlowType,
			Obj:  &flow.EBPFFlow{},
		}
	}

	prevKey := make([]byte, 8)
	key := make([]byte, 8)
	var nextAvailablePtr int
	now := time.Now()
	getFirstKey := true
	for {
		select {
		case <-p.quit:
			return
		case now = <-updateNow.C:
		default:
			if statsMap := p.module.Maps["stats_map"]; statsMap != nil {
				var dropKey uint32
				var dropValue int64

				if found, err := statsMap.Get(dropKey, &dropValue); err == nil && found {
					if dropValue > 0 {
						statsChan <- flow.Stats{KernelFlowDropped: dropValue}
					}
					statsMap.Delete(dropKey)
				}
			}
			// try to get start monotonic time
			if startKTimeNs == 0 {
				key := uint32(C.START_TIME_NS)
				var sns int64

				if found, err := p.cmap.Get(key, &sns); err == nil && found && sns != 0 {
					startKTimeNs = sns
					start = now
				}
				continue
			}

			tCancel := now.Add(ebpfMaxPollDelay)

			for {
				var err error
				var found bool
				if getFirstKey {
					if found, err = p.fmap[p.flowPage].NextKey(nil, &key); !found {
						/* map empty */
						p.swapPage()
						time.Sleep(time.Second)
						break
					}
					getFirstKey = false
				} else {
					found, err = p.fmap[p.flowPage].NextKey(prevKey, &key)
				}
				if !found || err != nil {
					getFirstKey = true
					break
				}

				kernFlow := unsafe.Pointer(&kernFlows[nextAvailablePtr])
				if _, err = p.fmap[p.flowPage].GetBytes(key, kernFlow); err != nil {
					getFirstKey = true
					break
				}
				p.fmap[p.flowPage].Delete(key)
				prevKey = key

				lastK := int64(kernFlows[nextAvailablePtr].last)
				last := start.Add(time.Duration(lastK - startKTimeNs))
				startK := int64(kernFlows[nextAvailablePtr].start)
				startFlow := start.Add(time.Duration(startK - startKTimeNs))
				if startFlow.After(now) {
					startFlow = now
				}

				extFlow := extFlows[nextAvailablePtr]

				ebpfFlow := extFlow.Obj.(*flow.EBPFFlow)
				ebpfFlow.Start = startFlow
				ebpfFlow.Last = last
				ebpfFlow.StartKTimeNs = startKTimeNs
				flow.SetEBPFKernFlow(ebpfFlow, kernFlow)

				extFlowChan <- &extFlow

				nextAvailablePtr = (nextAvailablePtr + 1) % flowPoolSize

				if now.After(tCancel) {
					break
				}
				time.Sleep(ebpfPollingRate)
				now = now.Add(ebpfPollingRate)
			}
		}
	}
}

func (p *Probe) stop() {
	p.quit <- true
}

// RegisterProbe registers an eBPF probe on an interface
func (p *ProbesHandler) RegisterProbe(n *graph.Node, capture *types.Capture, e probes.ProbeEventHandler) (probes.Probe, error) {
	ifName, _ := n.GetFieldString("Name")
	if ifName == "" {
		return nil, fmt.Errorf("No name for node %s", n.ID)
	}

	tid, _ := n.GetFieldString("TID")
	if tid == "" {
		return nil, fmt.Errorf("No tid for node %s", n.ID)
	}

	_, nsPath, err := topology.NamespaceFromNode(p.Ctx.Graph, n)
	if err != nil {
		return nil, err
	}

	module, err := p.LoadModule()
	if err != nil {
		return nil, err
	}

	fmap1 := module.Maps["flow_table_p1"]
	if fmap1 == nil {
		module.Close()
		return nil, fmt.Errorf("Unable to find flow_table_p1 map")
	}

	fmap2 := module.Maps["flow_table_p2"]
	if fmap2 == nil {
		module.Close()
		return nil, fmt.Errorf("Unable to find flow_table_p2 map")
	}

	cmap := module.Maps["u64_config_values"]
	if cmap == nil {
		module.Close()
		return nil, fmt.Errorf("Unable to find u64_config_values map")
	}

	socketFilter := module.Programs["bpf_flow_table"]
	if socketFilter == nil {
		module.Close()
		return nil, errors.New("No flow_table socket filter")
	}

	var rs *rawsocket.RawSocket
	if nsPath != "" {
		rs, err = rawsocket.NewRawSocketInNs(nsPath, ifName, syscall.ETH_P_ALL)
	} else {
		rs, err = rawsocket.NewRawSocket(ifName, syscall.ETH_P_ALL)
	}
	if err != nil {
		module.Close()
		return nil, err
	}
	fd := rs.GetFd()

	if ret := C.probe_bpf_attach_socket(C.int(fd), C.int(socketFilter.FD())); ret != 0 {
		rs.Close()
		module.Close()
		return nil, fmt.Errorf("Unable to attach socket filter to node: %s", n.ID)
	}

	uuids := flow.UUIDs{NodeTID: tid, CaptureID: capture.UUID}
	ft := p.Ctx.FTA.Alloc(uuids, flow.TableOpts{})

	probe := &Probe{
		Ctx:          p.Ctx,
		probeNodeTID: tid,
		fd:           rs.GetFd(),
		flowTable:    ft,
		module:       module,
		fmap:         []*ebpf.Map{fmap1, fmap2},
		cmap:         cmap,
		flowPage:     1,
		expire:       p.Ctx.FTA.ExpireAfter(),
		quit:         make(chan bool),
	}

	p.wg.Add(1)
	go func() {
		defer p.wg.Done()

		e.OnStarted(&probes.CaptureMetadata{})

		probe.run()

		if ret := C.probe_bpf_detach_socket(C.int(fd), C.int(socketFilter.FD())); ret != 0 {
			p.Ctx.Logger.Errorf("Unable to detach eBPF probe: %s", err)
		}
		rs.Close()
		module.Close()

		e.OnStopped()
	}()

	return probe, nil
}

// UnregisterProbe stops an eBPF probe on an interface
func (p *ProbesHandler) UnregisterProbe(n *graph.Node, e probes.ProbeEventHandler, fp probes.Probe) error {
	fp.(*Probe).stop()
	return nil
}

func (p *ProbesHandler) Start() error {
	return nil
}

func (p *ProbesHandler) Stop() {
	p.wg.Wait()
}

func LoadJumpMap(module *ebpf.Collection) error {
	var jmpTable []string = []string{"network_layer"}

	jmpTableMap := module.Maps["jmp_map"]
	if jmpTableMap == nil {
		return fmt.Errorf("Map: jmp_map not found")
	}
	for i, sym := range jmpTable {
		entry := module.Programs[sym]
		if entry == nil {
			return fmt.Errorf("Symbol %s not found", sym)
		}

		fd := uint32(entry.FD())
		err := jmpTableMap.Put(uint32(i), &fd)
		if err != nil {
			return err
		}
	}
	return nil
}

func (p *ProbesHandler) loadModuleFromAsset(path string) (*ebpf.Collection, error) {
	data, err := statics.Asset(path)
	if err != nil {
		return nil, fmt.Errorf("Unable to find eBPF elf binary in bindata")
	}

	collspec, err := ebpf.LoadCollectionSpecFromReader(bytes.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("Can't load %s: %v\n", path, err)
	}

	module, err := ebpf.NewCollection(collspec)
	if err != nil {
		return nil, fmt.Errorf("Can't create collection %s: %v\n", path, err)
	}

	if err == nil {
		p.Ctx.Logger.Infof("Loaded eBPF module %s", path)
	}
	return module, err
}

func (p *ProbesHandler) LoadModule() (*ebpf.Collection, error) {
	err := syscall.Setrlimit(C.RLIMIT_MEMLOCK, &syscall.Rlimit{
		Cur: math.MaxUint64,
		Max: math.MaxUint64,
	})
	if err != nil {
		return nil, fmt.Errorf("Failed to adjust rlimit map lock")
	}

	module, err := p.loadModuleFromAsset("ebpf/flow-gre.o")
	if err != nil {
		p.Ctx.Logger.Errorf("Unable to load eBPF elf binary (host %s) from bindata: %s, trying to fallback", runtime.GOARCH, err)

		module, err = p.loadModuleFromAsset("ebpf/flow.o")
		if err != nil {
			return nil, fmt.Errorf("Unable to load fallback eBPF elf binary (host %s) from bindata: %s", runtime.GOARCH, err)
		}
		p.Ctx.Logger.Info("Using fallback eBPF program")

		return module, nil
	}
	if err = LoadJumpMap(module); err != nil {
		return nil, err
	}
	return module, nil
}

// CaptureTypes supported
func (p *ProbesHandler) CaptureTypes() []string {
	return []string{"ebpf"}
}

// NewProbe returns a new eBPF probe
func NewProbe(ctx probes.Context, bundle *probe.Bundle) (probes.FlowProbeHandler, error) {
	return &ProbesHandler{Ctx: ctx}, nil
}
<commit_msg>flow: Delete prevKey to optimize NextKey functionality and add sleep to avoid busy loop<commit_after>// +build ebpf

/*
 * Copyright (C) 2017 Red Hat, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy ofthe License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specificlanguage governing permissions and
 * limitations under the License.
 *
 */

package ebpf

import (
	"bytes"
	"errors"
	"fmt"
	"math"
	"runtime"
	"sync"
	"syscall"
	"time"
	"unsafe"

	"github.com/newtools/ebpf"

	"github.com/skydive-project/skydive/api/types"
	"github.com/skydive-project/skydive/ebpf/statics"
	"github.com/skydive-project/skydive/flow"
	"github.com/skydive-project/skydive/flow/probes"
	"github.com/skydive-project/skydive/graffiti/graph"
	"github.com/skydive-project/skydive/probe"
	"github.com/skydive-project/skydive/rawsocket"
	"github.com/skydive-project/skydive/topology"
)

/*
#cgo CFLAGS: -I../../../ebpf
#include "flow.h"
#include <string.h>
#include <sys/resource.h>
#include <sys/socket.h>

int probe_bpf_attach_socket(int sock, int fd)
{
	return setsockopt(sock, SOL_SOCKET, SO_ATTACH_BPF, &fd, sizeof(fd));
}

int probe_bpf_detach_socket(int sock, int fd)
{
	return setsockopt(sock, SOL_SOCKET, SO_DETACH_BPF, &fd, sizeof(fd));
}
*/
import "C"

// Probe the eBPF probe
type Probe struct {
	Ctx          probes.Context
	probeNodeTID string
	fd           int
	flowTable    *flow.Table
	module       *ebpf.Collection
	fmap         []*ebpf.Map
	cmap         *ebpf.Map
	expire       time.Duration
	quit         chan bool
	flowPage     int
}

// ProbesHandler creates new eBPF probes
type ProbesHandler struct {
	Ctx probes.Context
	wg  sync.WaitGroup
}

func (p *Probe) swapPage() {
	key := uint32(C.FLOW_PAGE)
	writeInPage := int64(p.flowPage)

	p.cmap.Put(key, writeInPage)

	if p.flowPage == 1 {
		p.flowPage = 0
	} else {
		p.flowPage = 1
	}
}

func (p *Probe) run() {
	var info syscall.Sysinfo_t
	syscall.Sysinfo(&info)

	_, extFlowChan, statsChan := p.flowTable.Start(nil)
	defer p.flowTable.Stop()

	ebpfPollingRate := time.Second / time.Duration(p.Ctx.Config.GetInt("agent.flow.ebpf.polling_rate"))
	const ebpfMaxPollDelay = 10 * time.Second

	var startKTimeNs int64
	var start time.Time

	updateNow := time.NewTicker(ebpfMaxPollDelay)
	defer updateNow.Stop()

	flowPoolSize := 2 * cap(extFlowChan)
	kernFlows := make([]C.struct_flow, flowPoolSize)

	extFlows := make([]flow.ExtFlow, flowPoolSize)
	for i := range extFlows {
		extFlows[i] = flow.ExtFlow{
			Type: flow.EBPFExtFlowType,
			Obj:  &flow.EBPFFlow{},
		}
	}

	zeroKey := make([]byte, 8)
	prevKey := make([]byte, 8)
	delPrevKey := func() {
		if string(prevKey) != string(zeroKey) {
			p.fmap[p.flowPage].Delete(prevKey)
		}
	}
	key := make([]byte, 8)
	var nextAvailablePtr int
	now := time.Now()
	getFirstKey := true
	for {
		select {
		case <-p.quit:
			return
		case now = <-updateNow.C:
		default:
			if statsMap := p.module.Maps["stats_map"]; statsMap != nil {
				var dropKey uint32
				var dropValue int64

				if found, err := statsMap.Get(dropKey, &dropValue); err == nil && found {
					if dropValue > 0 {
						statsChan <- flow.Stats{KernelFlowDropped: dropValue}
					}
					statsMap.Delete(dropKey)
				}
			}
			// try to get start monotonic time
			if startKTimeNs == 0 {
				key := uint32(C.START_TIME_NS)
				var sns int64

				if found, err := p.cmap.Get(key, &sns); err == nil && found && sns != 0 {
					startKTimeNs = sns
					start = now
				}
				time.Sleep(time.Second)
				continue
			}

			tCancel := now.Add(ebpfMaxPollDelay)

			for {
				var err error
				var found bool
				if getFirstKey {
					if found, err = p.fmap[p.flowPage].NextKey(nil, &key); !found {
						/* map empty */
						p.swapPage()
						time.Sleep(time.Second)
						break
					}
					getFirstKey = false
				} else {
					found, err = p.fmap[p.flowPage].NextKey(prevKey, &key)
				}
				if !found || err != nil {
					delPrevKey()
					getFirstKey = true
					prevKey = zeroKey
					break
				}

				kernFlow := unsafe.Pointer(&kernFlows[nextAvailablePtr])
				if _, err = p.fmap[p.flowPage].GetBytes(key, kernFlow); err != nil {
					delPrevKey()
					getFirstKey = true
					prevKey = zeroKey
					break
				}
				delPrevKey()
				prevKey = key

				lastK := int64(kernFlows[nextAvailablePtr].last)
				last := start.Add(time.Duration(lastK - startKTimeNs))
				startK := int64(kernFlows[nextAvailablePtr].start)
				startFlow := start.Add(time.Duration(startK - startKTimeNs))
				if startFlow.After(now) {
					startFlow = now
				}

				extFlow := extFlows[nextAvailablePtr]

				ebpfFlow := extFlow.Obj.(*flow.EBPFFlow)
				ebpfFlow.Start = startFlow
				ebpfFlow.Last = last
				ebpfFlow.StartKTimeNs = startKTimeNs
				flow.SetEBPFKernFlow(ebpfFlow, kernFlow)

				extFlowChan <- &extFlow

				nextAvailablePtr = (nextAvailablePtr + 1) % flowPoolSize

				if now.After(tCancel) {
					break
				}
				time.Sleep(ebpfPollingRate)
				now = now.Add(ebpfPollingRate)
			}
		}
	}
}

func (p *Probe) stop() {
	p.quit <- true
}

// RegisterProbe registers an eBPF probe on an interface
func (p *ProbesHandler) RegisterProbe(n *graph.Node, capture *types.Capture, e probes.ProbeEventHandler) (probes.Probe, error) {
	ifName, _ := n.GetFieldString("Name")
	if ifName == "" {
		return nil, fmt.Errorf("No name for node %s", n.ID)
	}

	tid, _ := n.GetFieldString("TID")
	if tid == "" {
		return nil, fmt.Errorf("No tid for node %s", n.ID)
	}

	_, nsPath, err := topology.NamespaceFromNode(p.Ctx.Graph, n)
	if err != nil {
		return nil, err
	}

	module, err := p.LoadModule()
	if err != nil {
		return nil, err
	}

	fmap1 := module.Maps["flow_table_p1"]
	if fmap1 == nil {
		module.Close()
		return nil, fmt.Errorf("Unable to find flow_table_p1 map")
	}

	fmap2 := module.Maps["flow_table_p2"]
	if fmap2 == nil {
		module.Close()
		return nil, fmt.Errorf("Unable to find flow_table_p2 map")
	}

	cmap := module.Maps["u64_config_values"]
	if cmap == nil {
		module.Close()
		return nil, fmt.Errorf("Unable to find u64_config_values map")
	}

	socketFilter := module.Programs["bpf_flow_table"]
	if socketFilter == nil {
		module.Close()
		return nil, errors.New("No flow_table socket filter")
	}

	var rs *rawsocket.RawSocket
	if nsPath != "" {
		rs, err = rawsocket.NewRawSocketInNs(nsPath, ifName, syscall.ETH_P_ALL)
	} else {
		rs, err = rawsocket.NewRawSocket(ifName, syscall.ETH_P_ALL)
	}
	if err != nil {
		module.Close()
		return nil, err
	}
	fd := rs.GetFd()

	if ret := C.probe_bpf_attach_socket(C.int(fd), C.int(socketFilter.FD())); ret != 0 {
		rs.Close()
		module.Close()
		return nil, fmt.Errorf("Unable to attach socket filter to node: %s", n.ID)
	}

	uuids := flow.UUIDs{NodeTID: tid, CaptureID: capture.UUID}
	ft := p.Ctx.FTA.Alloc(uuids, flow.TableOpts{})

	probe := &Probe{
		Ctx:          p.Ctx,
		probeNodeTID: tid,
		fd:           rs.GetFd(),
		flowTable:    ft,
		module:       module,
		fmap:         []*ebpf.Map{fmap1, fmap2},
		cmap:         cmap,
		flowPage:     1,
		expire:       p.Ctx.FTA.ExpireAfter(),
		quit:         make(chan bool),
	}

	p.wg.Add(1)
	go func() {
		defer p.wg.Done()

		e.OnStarted(&probes.CaptureMetadata{})

		probe.run()

		if ret := C.probe_bpf_detach_socket(C.int(fd), C.int(socketFilter.FD())); ret != 0 {
			p.Ctx.Logger.Errorf("Unable to detach eBPF probe: %s", err)
		}
		rs.Close()
		module.Close()

		e.OnStopped()
	}()

	return probe, nil
}

// UnregisterProbe stops an eBPF probe on an interface
func (p *ProbesHandler) UnregisterProbe(n *graph.Node, e probes.ProbeEventHandler, fp probes.Probe) error {
	fp.(*Probe).stop()
	return nil
}

func (p *ProbesHandler) Start() error {
	return nil
}

func (p *ProbesHandler) Stop() {
	p.wg.Wait()
}

func LoadJumpMap(module *ebpf.Collection) error {
	var jmpTable []string = []string{"network_layer"}

	jmpTableMap := module.Maps["jmp_map"]
	if jmpTableMap == nil {
		return fmt.Errorf("Map: jmp_map not found")
	}
	for i, sym := range jmpTable {
		entry := module.Programs[sym]
		if entry == nil {
			return fmt.Errorf("Symbol %s not found", sym)
		}

		fd := uint32(entry.FD())
		err := jmpTableMap.Put(uint32(i), &fd)
		if err != nil {
			return err
		}
	}
	return nil
}

func (p *ProbesHandler) loadModuleFromAsset(path string) (*ebpf.Collection, error) {
	data, err := statics.Asset(path)
	if err != nil {
		return nil, fmt.Errorf("Unable to find eBPF elf binary in bindata")
	}

	collspec, err := ebpf.LoadCollectionSpecFromReader(bytes.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("Can't load %s: %v\n", path, err)
	}

	module, err := ebpf.NewCollection(collspec)
	if err != nil {
		return nil, fmt.Errorf("Can't create collection %s: %v\n", path, err)
	}

	if err == nil {
		p.Ctx.Logger.Infof("Loaded eBPF module %s", path)
	}
	return module, err
}

func (p *ProbesHandler) LoadModule() (*ebpf.Collection, error) {
	err := syscall.Setrlimit(C.RLIMIT_MEMLOCK, &syscall.Rlimit{
		Cur: math.MaxUint64,
		Max: math.MaxUint64,
	})
	if err != nil {
		return nil, fmt.Errorf("Failed to adjust rlimit map lock")
	}

	module, err := p.loadModuleFromAsset("ebpf/flow-gre.o")
	if err != nil {
		p.Ctx.Logger.Errorf("Unable to load eBPF elf binary (host %s) from bindata: %s, trying to fallback", runtime.GOARCH, err)

		module, err = p.loadModuleFromAsset("ebpf/flow.o")
		if err != nil {
			return nil, fmt.Errorf("Unable to load fallback eBPF elf binary (host %s) from bindata: %s", runtime.GOARCH, err)
		}
		p.Ctx.Logger.Info("Using fallback eBPF program")

		return module, nil
	}
	if err = LoadJumpMap(module); err != nil {
		return nil, err
	}
	return module, nil
}

// CaptureTypes supported
func (p *ProbesHandler) CaptureTypes() []string {
	return []string{"ebpf"}
}

// NewProbe returns a new eBPF probe
func NewProbe(ctx probes.Context, bundle *probe.Bundle) (probes.FlowProbeHandler, error) {
	return &ProbesHandler{Ctx: ctx}, nil
}
<|endoftext|>


<commit_before>package formfilter

import (
	"strings"

	"github.com/webx-top/echo"
)

type (
	Filter  func(*Data)
	Filters map[string][]Filter
	Options func() (key string, filter Filter)
	Data    struct {
		Key           string
		Value         []string
		normalizedKey string
	}
)

const All = `*`

func (d *Data) NormalizedKey() string {
	return d.normalizedKey
}

func Build(options ...Options) echo.FormDataFilter {
	filterMap := Filters{}
	for _, opt := range options {
		key, filter := opt()
		key = strings.Title(key)
		if _, ok := filterMap[key]; !ok {
			filterMap[key] = []Filter{}
		}
		filterMap[key] = append(filterMap[key], filter)
	}
	return echo.FormDataFilter(func(k string, v []string) (string, []string) {
		if k == All {
			return ``, nil
		}
		key := strings.Title(k)
		filters, ok := filterMap[key]
		if !ok {
			filters, ok = filterMap[All]
			if !ok {
				return k, v
			}
		}
		data := &Data{Key: k, Value: v, normalizedKey: key}
		for _, filter := range filters {
			filter(data)
			if len(data.Key) == 0 {
				break
			}
		}
		return data.Key, data.Value
	})
}
<commit_msg>update<commit_after>package formfilter

import (
	"strings"

	"github.com/webx-top/echo"
)

type (
	Filter  func(*Data)
	Filters map[string][]Filter
	Options func() (key string, filter Filter)
	Data    struct {
		Key           string
		Value         []string
		normalizedKey string
	}
)

const All = `*`

func (d *Data) NormalizedKey() string {
	return d.normalizedKey
}

func Build(options ...Options) echo.FormDataFilter {
	filterMap := Filters{}
	for _, opt := range options {
		key, filter := opt()
		key = strings.Title(key)
		if _, ok := filterMap[key]; !ok {
			filterMap[key] = []Filter{}
		}
		filterMap[key] = append(filterMap[key], filter)
	}
	return echo.FormDataFilter(func(k string, v []string) (string, []string) {
		if k == All {
			return ``, nil
		}
		key := strings.Title(k)
		filters, ok := filterMap[key]
		filtersAll, okAll := filterMap[All]
		if !ok && !okAll {
			return k, v
		}
		data := &Data{Key: k, Value: v, normalizedKey: key}
		if ok {
			for _, filter := range filters {
				filter(data)
				if len(data.Key) == 0 {
					return data.Key, data.Value
				}
			}
		}
		if okAll {
			for _, filter := range filtersAll {
				filter(data)
				if len(data.Key) == 0 {
					return data.Key, data.Value
				}
			}
		}
		return data.Key, data.Value
	})
}
<|endoftext|>


<commit_before>package gateway

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
	"time"

	"github.com/docker/distribution/reference"
	"github.com/moby/buildkit/cache"
	"github.com/moby/buildkit/client/llb"
	"github.com/moby/buildkit/executor"
	"github.com/moby/buildkit/frontend"
	pb "github.com/moby/buildkit/frontend/gateway/pb"
	"github.com/moby/buildkit/identity"
	"github.com/moby/buildkit/session"
	"github.com/moby/buildkit/solver"
	"github.com/moby/buildkit/util/tracing"
	"github.com/moby/buildkit/worker"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/net/http2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/health"
	"google.golang.org/grpc/health/grpc_health_v1"
)

const (
	keySource           = "source"
	keyDevel            = "gateway-devel"
	exporterImageConfig = "containerimage.config"
)

func NewGatewayFrontend() frontend.Frontend {
	return &gatewayFrontend{}
}

type gatewayFrontend struct {
}

func filterPrefix(opts map[string]string, pfx string) map[string]string {
	m := map[string]string{}
	for k, v := range opts {
		if strings.HasPrefix(k, pfx) {
			m[strings.TrimPrefix(k, pfx)] = v
		}
	}
	return m
}

func (gf *gatewayFrontend) Solve(ctx context.Context, llbBridge frontend.FrontendLLBBridge, opts map[string]string) (retRef solver.CachedResult, exporterAttr map[string][]byte, retErr error) {
	source, ok := opts[keySource]
	if !ok {
		return nil, nil, errors.Errorf("no source specified for gateway")
	}

	sid := session.FromContext(ctx)

	_, isDevel := opts[keyDevel]
	var img ocispec.Image
	var rootFS cache.ImmutableRef
	var readonly bool // TODO: try to switch to read-only by default.

	if isDevel {
		ref, exp, err := llbBridge.Solve(session.NewContext(ctx, "gateway:"+sid),
			frontend.SolveRequest{
				Frontend:    source,
				FrontendOpt: filterPrefix(opts, "gateway-"),
			})
		if err != nil {
			return nil, nil, err
		}
		defer ref.Release(context.TODO())

		workerRef, ok := ref.Sys().(*worker.WorkerRef)
		if !ok {
			return nil, nil, errors.Errorf("invalid ref: %T", ref.Sys())
		}
		rootFS = workerRef.ImmutableRef
		config, ok := exp[exporterImageConfig]
		if ok {
			if err := json.Unmarshal(config, &img); err != nil {
				return nil, nil, err
			}
		}
	} else {
		sourceRef, err := reference.ParseNormalizedNamed(source)
		if err != nil {
			return nil, nil, err
		}

		dgst, config, err := llbBridge.ResolveImageConfig(ctx, reference.TagNameOnly(sourceRef).String())
		if err != nil {
			return nil, nil, err
		}

		if err := json.Unmarshal(config, &img); err != nil {
			return nil, nil, err
		}

		sourceRef, err = reference.WithDigest(sourceRef, dgst)
		if err != nil {
			return nil, nil, err
		}

		src := llb.Image(sourceRef.String())

		def, err := src.Marshal()
		if err != nil {
			return nil, nil, err
		}

		ref, _, err := llbBridge.Solve(ctx, frontend.SolveRequest{
			Definition: def.ToPB(),
		})
		if err != nil {
			return nil, nil, err
		}
		defer ref.Release(context.TODO())
		workerRef, ok := ref.Sys().(*worker.WorkerRef)
		if !ok {
			return nil, nil, errors.Errorf("invalid ref: %T", ref.Sys())
		}
		rootFS = workerRef.ImmutableRef
	}

	lbf, err := newLLBBridgeForwarder(ctx, llbBridge)
	defer lbf.conn.Close()
	if err != nil {
		return nil, nil, err
	}

	args := []string{"/run"}
	env := []string{}
	cwd := "/"
	if img.Config.Env != nil {
		env = img.Config.Env
	}
	if img.Config.Entrypoint != nil {
		args = img.Config.Entrypoint
	}
	if img.Config.WorkingDir != "" {
		cwd = img.Config.WorkingDir
	}
	i := 0
	for k, v := range opts {
		env = append(env, fmt.Sprintf("BUILDKIT_FRONTEND_OPT_%d", i)+"="+k+"="+v)
		i++
	}

	env = append(env, "BUILDKIT_SESSION_ID="+sid)

	defer func() {
		for _, r := range lbf.refs {
			if r != nil && (lbf.lastRef != r || retErr != nil) {
				r.Release(context.TODO())
			}
		}
	}()

	err = llbBridge.Exec(ctx, executor.Meta{
		Env:            env,
		Args:           args,
		Cwd:            cwd,
		ReadonlyRootFS: readonly,
	}, rootFS, lbf.Stdin, lbf.Stdout, os.Stderr)

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

	return lbf.lastRef, lbf.exporterAttr, nil
}

func newLLBBridgeForwarder(ctx context.Context, llbBridge frontend.FrontendLLBBridge) (*llbBridgeForwarder, error) {
	lbf := &llbBridgeForwarder{
		callCtx:   ctx,
		llbBridge: llbBridge,
		refs:      map[string]solver.Result{},
		pipe:      newPipe(),
	}

	server := grpc.NewServer()
	grpc_health_v1.RegisterHealthServer(server, health.NewServer())
	pb.RegisterLLBBridgeServer(server, lbf)

	go serve(ctx, server, lbf.conn)

	return lbf, nil
}

type pipe struct {
	Stdin  io.ReadCloser
	Stdout io.WriteCloser
	conn   net.Conn
}

func newPipe() *pipe {
	pr1, pw1, _ := os.Pipe()
	pr2, pw2, _ := os.Pipe()
	return &pipe{
		Stdin:  pr1,
		Stdout: pw2,
		conn: &conn{
			Reader: pr2,
			Writer: pw1,
			Closer: pw2,
		},
	}
}

type conn struct {
	io.Reader
	io.Writer
	io.Closer
}

func (s *conn) LocalAddr() net.Addr {
	return dummyAddr{}
}
func (s *conn) RemoteAddr() net.Addr {
	return dummyAddr{}
}
func (s *conn) SetDeadline(t time.Time) error {
	return nil
}
func (s *conn) SetReadDeadline(t time.Time) error {
	return nil
}
func (s *conn) SetWriteDeadline(t time.Time) error {
	return nil
}

type dummyAddr struct {
}

func (d dummyAddr) Network() string {
	return "pipe"
}

func (d dummyAddr) String() string {
	return "localhost"
}

type llbBridgeForwarder struct {
	callCtx      context.Context
	llbBridge    frontend.FrontendLLBBridge
	refs         map[string]solver.Result
	lastRef      solver.CachedResult
	exporterAttr map[string][]byte
	*pipe
}

func (lbf *llbBridgeForwarder) ResolveImageConfig(ctx context.Context, req *pb.ResolveImageConfigRequest) (*pb.ResolveImageConfigResponse, error) {
	ctx = tracing.ContextWithSpanFromContext(ctx, lbf.callCtx)
	dgst, dt, err := lbf.llbBridge.ResolveImageConfig(ctx, req.Ref)
	if err != nil {
		return nil, err
	}
	return &pb.ResolveImageConfigResponse{
		Digest: dgst,
		Config: dt,
	}, nil
}

func (lbf *llbBridgeForwarder) Solve(ctx context.Context, req *pb.SolveRequest) (*pb.SolveResponse, error) {
	ctx = tracing.ContextWithSpanFromContext(ctx, lbf.callCtx)
	ref, expResp, err := lbf.llbBridge.Solve(ctx, frontend.SolveRequest{
		Definition:      req.Definition,
		Frontend:        req.Frontend,
		FrontendOpt:     req.FrontendOpt,
		ImportCacheRefs: req.ImportCacheRefs,
	})
	if err != nil {
		return nil, err
	}

	exp := map[string][]byte{}
	if err := json.Unmarshal(req.ExporterAttr, &exp); err != nil {
		return nil, err
	}

	if expResp != nil {
		for k, v := range expResp {
			exp[k] = v
		}
	}

	id := identity.NewID()
	lbf.refs[id] = ref
	if req.Final {
		lbf.lastRef = ref
		lbf.exporterAttr = exp
	}
	if ref == nil {
		id = ""
	}
	return &pb.SolveResponse{Ref: id}, nil
}
func (lbf *llbBridgeForwarder) ReadFile(ctx context.Context, req *pb.ReadFileRequest) (*pb.ReadFileResponse, error) {
	ctx = tracing.ContextWithSpanFromContext(ctx, lbf.callCtx)
	ref, ok := lbf.refs[req.Ref]
	if !ok {
		return nil, errors.Errorf("no such ref: %v", req.Ref)
	}
	if ref == nil {
		return nil, errors.Wrapf(os.ErrNotExist, "%s no found", req.FilePath)
	}
	workerRef, ok := ref.Sys().(*worker.WorkerRef)
	if !ok {
		return nil, errors.Errorf("invalid ref: %T", ref.Sys())
	}

	newReq := cache.ReadRequest{
		Filename: req.FilePath,
	}
	if r := req.Range; r != nil {
		newReq.Range = &cache.FileRange{
			Offset: int(r.Offset),
			Length: int(r.Length),
		}
	}

	dt, err := cache.ReadFile(ctx, workerRef.ImmutableRef, newReq)
	if err != nil {
		return nil, err
	}

	return &pb.ReadFileResponse{Data: dt}, nil
}

func (lbf *llbBridgeForwarder) Ping(context.Context, *pb.PingRequest) (*pb.PongResponse, error) {
	return &pb.PongResponse{}, nil
}

func serve(ctx context.Context, grpcServer *grpc.Server, conn net.Conn) {
	go func() {
		<-ctx.Done()
		conn.Close()
	}()
	logrus.Debugf("serving grpc connection")
	(&http2.Server{}).ServeConn(conn, &http2.ServeConnOpts{Handler: grpcServer})
}
<commit_msg>gateway: protect gateway forwarder map access<commit_after>package gateway

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/docker/distribution/reference"
	"github.com/moby/buildkit/cache"
	"github.com/moby/buildkit/client/llb"
	"github.com/moby/buildkit/executor"
	"github.com/moby/buildkit/frontend"
	pb "github.com/moby/buildkit/frontend/gateway/pb"
	"github.com/moby/buildkit/identity"
	"github.com/moby/buildkit/session"
	"github.com/moby/buildkit/solver"
	"github.com/moby/buildkit/util/tracing"
	"github.com/moby/buildkit/worker"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/net/http2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/health"
	"google.golang.org/grpc/health/grpc_health_v1"
)

const (
	keySource           = "source"
	keyDevel            = "gateway-devel"
	exporterImageConfig = "containerimage.config"
)

func NewGatewayFrontend() frontend.Frontend {
	return &gatewayFrontend{}
}

type gatewayFrontend struct {
}

func filterPrefix(opts map[string]string, pfx string) map[string]string {
	m := map[string]string{}
	for k, v := range opts {
		if strings.HasPrefix(k, pfx) {
			m[strings.TrimPrefix(k, pfx)] = v
		}
	}
	return m
}

func (gf *gatewayFrontend) Solve(ctx context.Context, llbBridge frontend.FrontendLLBBridge, opts map[string]string) (retRef solver.CachedResult, exporterAttr map[string][]byte, retErr error) {
	source, ok := opts[keySource]
	if !ok {
		return nil, nil, errors.Errorf("no source specified for gateway")
	}

	sid := session.FromContext(ctx)

	_, isDevel := opts[keyDevel]
	var img ocispec.Image
	var rootFS cache.ImmutableRef
	var readonly bool // TODO: try to switch to read-only by default.

	if isDevel {
		ref, exp, err := llbBridge.Solve(session.NewContext(ctx, "gateway:"+sid),
			frontend.SolveRequest{
				Frontend:    source,
				FrontendOpt: filterPrefix(opts, "gateway-"),
			})
		if err != nil {
			return nil, nil, err
		}
		defer ref.Release(context.TODO())

		workerRef, ok := ref.Sys().(*worker.WorkerRef)
		if !ok {
			return nil, nil, errors.Errorf("invalid ref: %T", ref.Sys())
		}
		rootFS = workerRef.ImmutableRef
		config, ok := exp[exporterImageConfig]
		if ok {
			if err := json.Unmarshal(config, &img); err != nil {
				return nil, nil, err
			}
		}
	} else {
		sourceRef, err := reference.ParseNormalizedNamed(source)
		if err != nil {
			return nil, nil, err
		}

		dgst, config, err := llbBridge.ResolveImageConfig(ctx, reference.TagNameOnly(sourceRef).String())
		if err != nil {
			return nil, nil, err
		}

		if err := json.Unmarshal(config, &img); err != nil {
			return nil, nil, err
		}

		sourceRef, err = reference.WithDigest(sourceRef, dgst)
		if err != nil {
			return nil, nil, err
		}

		src := llb.Image(sourceRef.String())

		def, err := src.Marshal()
		if err != nil {
			return nil, nil, err
		}

		ref, _, err := llbBridge.Solve(ctx, frontend.SolveRequest{
			Definition: def.ToPB(),
		})
		if err != nil {
			return nil, nil, err
		}
		defer ref.Release(context.TODO())
		workerRef, ok := ref.Sys().(*worker.WorkerRef)
		if !ok {
			return nil, nil, errors.Errorf("invalid ref: %T", ref.Sys())
		}
		rootFS = workerRef.ImmutableRef
	}

	lbf, err := newLLBBridgeForwarder(ctx, llbBridge)
	defer lbf.conn.Close()
	if err != nil {
		return nil, nil, err
	}

	args := []string{"/run"}
	env := []string{}
	cwd := "/"
	if img.Config.Env != nil {
		env = img.Config.Env
	}
	if img.Config.Entrypoint != nil {
		args = img.Config.Entrypoint
	}
	if img.Config.WorkingDir != "" {
		cwd = img.Config.WorkingDir
	}
	i := 0
	for k, v := range opts {
		env = append(env, fmt.Sprintf("BUILDKIT_FRONTEND_OPT_%d", i)+"="+k+"="+v)
		i++
	}

	env = append(env, "BUILDKIT_SESSION_ID="+sid)

	defer func() {
		for _, r := range lbf.refs {
			if r != nil && (lbf.lastRef != r || retErr != nil) {
				r.Release(context.TODO())
			}
		}
	}()

	err = llbBridge.Exec(ctx, executor.Meta{
		Env:            env,
		Args:           args,
		Cwd:            cwd,
		ReadonlyRootFS: readonly,
	}, rootFS, lbf.Stdin, lbf.Stdout, os.Stderr)

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

	return lbf.lastRef, lbf.exporterAttr, nil
}

func newLLBBridgeForwarder(ctx context.Context, llbBridge frontend.FrontendLLBBridge) (*llbBridgeForwarder, error) {
	lbf := &llbBridgeForwarder{
		callCtx:   ctx,
		llbBridge: llbBridge,
		refs:      map[string]solver.Result{},
		pipe:      newPipe(),
	}

	server := grpc.NewServer()
	grpc_health_v1.RegisterHealthServer(server, health.NewServer())
	pb.RegisterLLBBridgeServer(server, lbf)

	go serve(ctx, server, lbf.conn)

	return lbf, nil
}

type pipe struct {
	Stdin  io.ReadCloser
	Stdout io.WriteCloser
	conn   net.Conn
}

func newPipe() *pipe {
	pr1, pw1, _ := os.Pipe()
	pr2, pw2, _ := os.Pipe()
	return &pipe{
		Stdin:  pr1,
		Stdout: pw2,
		conn: &conn{
			Reader: pr2,
			Writer: pw1,
			Closer: pw2,
		},
	}
}

type conn struct {
	io.Reader
	io.Writer
	io.Closer
}

func (s *conn) LocalAddr() net.Addr {
	return dummyAddr{}
}
func (s *conn) RemoteAddr() net.Addr {
	return dummyAddr{}
}
func (s *conn) SetDeadline(t time.Time) error {
	return nil
}
func (s *conn) SetReadDeadline(t time.Time) error {
	return nil
}
func (s *conn) SetWriteDeadline(t time.Time) error {
	return nil
}

type dummyAddr struct {
}

func (d dummyAddr) Network() string {
	return "pipe"
}

func (d dummyAddr) String() string {
	return "localhost"
}

type llbBridgeForwarder struct {
	mu           sync.Mutex
	callCtx      context.Context
	llbBridge    frontend.FrontendLLBBridge
	refs         map[string]solver.Result
	lastRef      solver.CachedResult
	exporterAttr map[string][]byte
	*pipe
}

func (lbf *llbBridgeForwarder) ResolveImageConfig(ctx context.Context, req *pb.ResolveImageConfigRequest) (*pb.ResolveImageConfigResponse, error) {
	ctx = tracing.ContextWithSpanFromContext(ctx, lbf.callCtx)
	dgst, dt, err := lbf.llbBridge.ResolveImageConfig(ctx, req.Ref)
	if err != nil {
		return nil, err
	}
	return &pb.ResolveImageConfigResponse{
		Digest: dgst,
		Config: dt,
	}, nil
}

func (lbf *llbBridgeForwarder) Solve(ctx context.Context, req *pb.SolveRequest) (*pb.SolveResponse, error) {
	ctx = tracing.ContextWithSpanFromContext(ctx, lbf.callCtx)
	ref, expResp, err := lbf.llbBridge.Solve(ctx, frontend.SolveRequest{
		Definition:      req.Definition,
		Frontend:        req.Frontend,
		FrontendOpt:     req.FrontendOpt,
		ImportCacheRefs: req.ImportCacheRefs,
	})
	if err != nil {
		return nil, err
	}

	exp := map[string][]byte{}
	if err := json.Unmarshal(req.ExporterAttr, &exp); err != nil {
		return nil, err
	}

	if expResp != nil {
		for k, v := range expResp {
			exp[k] = v
		}
	}

	id := identity.NewID()
	lbf.mu.Lock()
	lbf.refs[id] = ref
	lbf.mu.Unlock()
	if req.Final {
		lbf.lastRef = ref
		lbf.exporterAttr = exp
	}
	if ref == nil {
		id = ""
	}
	return &pb.SolveResponse{Ref: id}, nil
}
func (lbf *llbBridgeForwarder) ReadFile(ctx context.Context, req *pb.ReadFileRequest) (*pb.ReadFileResponse, error) {
	ctx = tracing.ContextWithSpanFromContext(ctx, lbf.callCtx)
	lbf.mu.Lock()
	ref, ok := lbf.refs[req.Ref]
	lbf.mu.Unlock()
	if !ok {
		return nil, errors.Errorf("no such ref: %v", req.Ref)
	}
	if ref == nil {
		return nil, errors.Wrapf(os.ErrNotExist, "%s no found", req.FilePath)
	}
	workerRef, ok := ref.Sys().(*worker.WorkerRef)
	if !ok {
		return nil, errors.Errorf("invalid ref: %T", ref.Sys())
	}

	newReq := cache.ReadRequest{
		Filename: req.FilePath,
	}
	if r := req.Range; r != nil {
		newReq.Range = &cache.FileRange{
			Offset: int(r.Offset),
			Length: int(r.Length),
		}
	}

	dt, err := cache.ReadFile(ctx, workerRef.ImmutableRef, newReq)
	if err != nil {
		return nil, err
	}

	return &pb.ReadFileResponse{Data: dt}, nil
}

func (lbf *llbBridgeForwarder) Ping(context.Context, *pb.PingRequest) (*pb.PongResponse, error) {
	return &pb.PongResponse{}, nil
}

func serve(ctx context.Context, grpcServer *grpc.Server, conn net.Conn) {
	go func() {
		<-ctx.Done()
		conn.Close()
	}()
	logrus.Debugf("serving grpc connection")
	(&http2.Server{}).ServeConn(conn, &http2.ServeConnOpts{Handler: grpcServer})
}
<|endoftext|>


<commit_before>package namespace

import (
	"time"

	"v.io/v23"
	"v.io/v23/context"
	"v.io/v23/ipc"
	"v.io/v23/naming"
	"v.io/v23/naming/ns"
	"v.io/v23/services/security/access"
	"v.io/v23/verror"
)

type Server struct {
	ns ns.Namespace
}

func New(ctx *context.T) *Server {
	return &Server{v23.GetNamespace(ctx)}
}

func (s *Server) Glob(call *NamespaceGlobServerCallStub, pattern string) error {
	// Call Glob on the namespace client instance
	ch, err := s.ns.Glob(call.Context(), pattern)
	if err != nil {
		return err
	}

	stream := call.SendStream()

	for mp := range ch {
		var reply naming.VDLGlobReply
		switch v := mp.(type) {
		case *naming.GlobError:
			reply = naming.VDLGlobReplyError{*v}
		case *naming.MountEntry:
			reply = naming.VDLGlobReplyEntry{convertToVDLEntry(*v)}
		}
		if err = stream.Send(reply); err != nil {
			return err
		}
	}
	return nil
}

func (s *Server) Mount(call ipc.ServerCall, name, server string, ttl time.Duration, replace bool) error {
	rmOpt := naming.ReplaceMountOpt(replace)
	err := s.ns.Mount(call.Context(), name, server, ttl, rmOpt)
	if err != nil {
		err = verror.Convert(verror.ErrInternal, call.Context(), err)
	}
	return err
}

func (s *Server) Unmount(call ipc.ServerCall, name, server string) error {
	return s.ns.Unmount(call.Context(), name, server)
}

func (s *Server) Resolve(call ipc.ServerCall, name string) ([]string, error) {
	me, err := s.ns.Resolve(call.Context(), name)
	if err != nil {
		return nil, verror.Convert(verror.ErrInternal, call.Context(), err)
	}
	return me.Names(), nil
}

func (s *Server) ResolveToMT(call ipc.ServerCall, name string) ([]string, error) {
	me, err := s.ns.ResolveToMountTable(call.Context(), name)
	if err != nil {
		return nil, verror.Convert(verror.ErrInternal, call.Context(), err)
	}
	return me.Names(), nil
}

func (s *Server) FlushCacheEntry(call ipc.ServerCall, name string) (bool, error) {
	return s.ns.FlushCacheEntry(name), nil
}

func (s *Server) DisableCache(call ipc.ServerCall, disable bool) error {
	disableCacheCtl := naming.DisableCache(disable)
	_ = s.ns.CacheCtl(disableCacheCtl)
	return nil
}

func (s *Server) Roots(call ipc.ServerCall) ([]string, error) {
	return s.ns.Roots(), nil
}

func (s *Server) SetRoots(call ipc.ServerCall, roots []string) error {
	if err := s.ns.SetRoots(roots...); err != nil {
		return verror.Convert(verror.ErrInternal, call.Context(), err)
	}
	return nil
}

func (s *Server) SetACL(call ipc.ServerCall, name string, acl access.TaggedACLMap, etag string) error {
	return s.ns.SetACL(call.Context(), name, acl, etag)
}

func (s *Server) GetACL(call ipc.ServerCall, name string) (access.TaggedACLMap, string, error) {
	return s.ns.GetACL(call.Context(), name)
}

func (s *Server) Delete(call ipc.ServerCall, name string, deleteSubtree bool) error {
	return s.ns.Delete(call.Context(), name, deleteSubtree)
}

func convertToVDLEntry(value naming.MountEntry) naming.VDLMountEntry {
	result := naming.VDLMountEntry{
		Name: value.Name,
		MT:   value.ServesMountTable(),
	}
	for _, s := range value.Servers {
		result.Servers = append(result.Servers,
			naming.VDLMountedServer{
				Server: s.Server,
				TTL:    uint32(s.Expires.Sub(time.Now())),
			})
	}
	return result
}
<commit_msg>TBR: Include BlessingPatterns in Glob result.<commit_after>package namespace

import (
	"time"

	"v.io/v23"
	"v.io/v23/context"
	"v.io/v23/ipc"
	"v.io/v23/naming"
	"v.io/v23/naming/ns"
	"v.io/v23/services/security/access"
	"v.io/v23/verror"
)

type Server struct {
	ns ns.Namespace
}

func New(ctx *context.T) *Server {
	return &Server{v23.GetNamespace(ctx)}
}

func (s *Server) Glob(call *NamespaceGlobServerCallStub, pattern string) error {
	// Call Glob on the namespace client instance
	ch, err := s.ns.Glob(call.Context(), pattern)
	if err != nil {
		return err
	}

	stream := call.SendStream()

	for mp := range ch {
		var reply naming.VDLGlobReply
		switch v := mp.(type) {
		case *naming.GlobError:
			reply = naming.VDLGlobReplyError{*v}
		case *naming.MountEntry:
			reply = naming.VDLGlobReplyEntry{convertToVDLEntry(*v)}
		}
		if err = stream.Send(reply); err != nil {
			return err
		}
	}
	return nil
}

func (s *Server) Mount(call ipc.ServerCall, name, server string, ttl time.Duration, replace bool) error {
	rmOpt := naming.ReplaceMountOpt(replace)
	err := s.ns.Mount(call.Context(), name, server, ttl, rmOpt)
	if err != nil {
		err = verror.Convert(verror.ErrInternal, call.Context(), err)
	}
	return err
}

func (s *Server) Unmount(call ipc.ServerCall, name, server string) error {
	return s.ns.Unmount(call.Context(), name, server)
}

func (s *Server) Resolve(call ipc.ServerCall, name string) ([]string, error) {
	me, err := s.ns.Resolve(call.Context(), name)
	if err != nil {
		return nil, verror.Convert(verror.ErrInternal, call.Context(), err)
	}
	return me.Names(), nil
}

func (s *Server) ResolveToMT(call ipc.ServerCall, name string) ([]string, error) {
	me, err := s.ns.ResolveToMountTable(call.Context(), name)
	if err != nil {
		return nil, verror.Convert(verror.ErrInternal, call.Context(), err)
	}
	return me.Names(), nil
}

func (s *Server) FlushCacheEntry(call ipc.ServerCall, name string) (bool, error) {
	return s.ns.FlushCacheEntry(name), nil
}

func (s *Server) DisableCache(call ipc.ServerCall, disable bool) error {
	disableCacheCtl := naming.DisableCache(disable)
	_ = s.ns.CacheCtl(disableCacheCtl)
	return nil
}

func (s *Server) Roots(call ipc.ServerCall) ([]string, error) {
	return s.ns.Roots(), nil
}

func (s *Server) SetRoots(call ipc.ServerCall, roots []string) error {
	if err := s.ns.SetRoots(roots...); err != nil {
		return verror.Convert(verror.ErrInternal, call.Context(), err)
	}
	return nil
}

func (s *Server) SetACL(call ipc.ServerCall, name string, acl access.TaggedACLMap, etag string) error {
	return s.ns.SetACL(call.Context(), name, acl, etag)
}

func (s *Server) GetACL(call ipc.ServerCall, name string) (access.TaggedACLMap, string, error) {
	return s.ns.GetACL(call.Context(), name)
}

func (s *Server) Delete(call ipc.ServerCall, name string, deleteSubtree bool) error {
	return s.ns.Delete(call.Context(), name, deleteSubtree)
}

func convertToVDLEntry(value naming.MountEntry) naming.VDLMountEntry {
	result := naming.VDLMountEntry{
		Name: value.Name,
		MT:   value.ServesMountTable(),
	}
	for _, s := range value.Servers {
		result.Servers = append(result.Servers,
			naming.VDLMountedServer{
				Server:           s.Server,
				TTL:              uint32(s.Expires.Sub(time.Now())),
				BlessingPatterns: s.BlessingPatterns,
			})
	}
	return result
}
<|endoftext|>


<commit_before>// Copyright 2016 The LUCI Authors. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

package monitor

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"

	"github.com/golang/protobuf/jsonpb"
	"golang.org/x/net/context"

	"github.com/luci/luci-go/common/clock"
	"github.com/luci/luci-go/common/lhttp"
	"github.com/luci/luci-go/common/logging"
	pb "github.com/luci/luci-go/common/tsmon/ts_mon_proto_v2"
	"github.com/luci/luci-go/common/tsmon/types"
)

var (
	// ProdxmonScopes is the list of oauth2 scopes needed on the http client
	// given to NewHTTPMonitor.
	ProdxmonScopes = []string{"https://www.googleapis.com/auth/prodxmon"}
)

type httpMonitor struct {
	client   *http.Client
	endpoint *url.URL
}

// NewHTTPMonitor creates a new Monitor object that sends metric to an HTTP
// (or HTTPS) endpoint.  The http client should be authenticated as required.
func NewHTTPMonitor(ctx context.Context, client *http.Client, endpoint *url.URL) (Monitor, error) {
	return &httpMonitor{
		client:   client,
		endpoint: endpoint,
	}, nil
}

func (m *httpMonitor) ChunkSize() int {
	return 1000
}

func (m *httpMonitor) Send(ctx context.Context, cells []types.Cell) error {
	// Serialize the tsmon cells into protobufs.
	req := &pb.Request{
		Payload: &pb.MetricsPayload{
			MetricsCollection: SerializeCells(cells, clock.Now(ctx)),
		},
	}

	// JSON encode the request.
	encoded := bytes.Buffer{}
	marshaller := jsonpb.Marshaler{}
	if err := marshaller.Marshal(&encoded, req); err != nil {
		return err
	}

	// Make the request.
	status, err := lhttp.NewRequest(ctx, m.client, nil, func() (*http.Request, error) {
		return http.NewRequest("POST", m.endpoint.String(), bytes.NewReader(encoded.Bytes()))
	}, func(resp *http.Response) error {
		return resp.Body.Close()
	}, func(resp *http.Response, oErr error) error {
		if resp != nil {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				logging.WithError(err).Errorf(ctx, "Failed to read error response body")
			} else {
				logging.Errorf(ctx, "Monitoring push failed.  Response body: %s", body)
			}
			resp.Body.Close()
		}
		return oErr
	})()
	if err != nil {
		return err
	}
	if status != http.StatusOK {
		return fmt.Errorf("bad response status %d from endpoint %s", status, m.endpoint)
	}

	logging.Debugf(ctx, "Sent %d tsmon cells to %s", len(cells), m.endpoint)
	return nil
}

func (m *httpMonitor) Close() error {
	return nil
}
<commit_msg>Reduce the number of cells in each HTTP request from 1000 to 500.<commit_after>// Copyright 2016 The LUCI Authors. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

package monitor

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"

	"github.com/golang/protobuf/jsonpb"
	"golang.org/x/net/context"

	"github.com/luci/luci-go/common/clock"
	"github.com/luci/luci-go/common/lhttp"
	"github.com/luci/luci-go/common/logging"
	pb "github.com/luci/luci-go/common/tsmon/ts_mon_proto_v2"
	"github.com/luci/luci-go/common/tsmon/types"
)

var (
	// ProdxmonScopes is the list of oauth2 scopes needed on the http client
	// given to NewHTTPMonitor.
	ProdxmonScopes = []string{"https://www.googleapis.com/auth/prodxmon"}
)

type httpMonitor struct {
	client   *http.Client
	endpoint *url.URL
}

// NewHTTPMonitor creates a new Monitor object that sends metric to an HTTP
// (or HTTPS) endpoint.  The http client should be authenticated as required.
func NewHTTPMonitor(ctx context.Context, client *http.Client, endpoint *url.URL) (Monitor, error) {
	return &httpMonitor{
		client:   client,
		endpoint: endpoint,
	}, nil
}

func (m *httpMonitor) ChunkSize() int {
	return 500
}

func (m *httpMonitor) Send(ctx context.Context, cells []types.Cell) error {
	// Serialize the tsmon cells into protobufs.
	req := &pb.Request{
		Payload: &pb.MetricsPayload{
			MetricsCollection: SerializeCells(cells, clock.Now(ctx)),
		},
	}

	// JSON encode the request.
	encoded := bytes.Buffer{}
	marshaller := jsonpb.Marshaler{}
	if err := marshaller.Marshal(&encoded, req); err != nil {
		return err
	}

	// Make the request.
	status, err := lhttp.NewRequest(ctx, m.client, nil, func() (*http.Request, error) {
		return http.NewRequest("POST", m.endpoint.String(), bytes.NewReader(encoded.Bytes()))
	}, func(resp *http.Response) error {
		return resp.Body.Close()
	}, func(resp *http.Response, oErr error) error {
		if resp != nil {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				logging.WithError(err).Errorf(ctx, "Failed to read error response body")
			} else {
				logging.Errorf(ctx, "Monitoring push failed.  Response body: %s", body)
			}
			resp.Body.Close()
		}
		return oErr
	})()
	if err != nil {
		return err
	}
	if status != http.StatusOK {
		return fmt.Errorf("bad response status %d from endpoint %s", status, m.endpoint)
	}

	logging.Debugf(ctx, "Sent %d tsmon cells to %s", len(cells), m.endpoint)
	return nil
}

func (m *httpMonitor) Close() error {
	return nil
}
<|endoftext|>


<commit_before>package objc

/*
#cgo LDFLAGS: -lobjc
#include <objc/runtime.h>

void *GoObjc_RegisterSelector(char *name) {
	return (void *) sel_registerName(name);
}

char *GoObjc_TypeInfoForMethod(void *cls, void *sel) {
	Method m = class_getInstanceMethod(cls, sel);
	return (char *) method_getTypeEncoding(m);
}
*/
import "C"
import "unsafe"

// selectorWithName looks up a selector by name.
func selectorWithName(name string) unsafe.Pointer {
	return C.GoObjc_RegisterSelector(C.CString(name))
}

// typeInfoForMethod returns the type encoding string for
// selector on obj's Class.
func typeInfoForMethod(obj Object, selector string) string {
	sel := selectorWithName(selector)
	cls := getObjectClass(obj)
	return C.GoString(C.GoObjc_TypeInfoForMethod(unsafe.Pointer(cls.Pointer()), sel))
}

// simpleTypeInfoForMethod returns a simplified typeInfo representation
// with C specifiers and stack information stripped out.
func simpleTypeInfoForMethod(obj Object, selector string) string {
	ti := typeInfoForMethod(obj, selector)
	sti := []byte{}
	for i := 0; i < len(ti); i++ {
		if ti[i] >= '0' && ti[i] <= '9' {
			continue
		}
		if string(ti[i]) == encConst {
			continue
		}
		// fixme(mkrautz): What is V? The NSObject release method uses V.
		if ti[i] == 'V' {
			continue
		}
		sti = append(sti, ti[i])
	}
	return string(sti)
}
<commit_msg>selector: implement a Selector interface.<commit_after>package objc

/*
#cgo LDFLAGS: -lobjc
#include <objc/runtime.h>

void *GoObjc_RegisterSelector(char *name) {
	return (void *) sel_registerName(name);
}

char *GoObjc_TypeInfoForMethod(void *cls, void *sel) {
	Method m = class_getInstanceMethod(cls, sel);
	return (char *) method_getTypeEncoding(m);
}
*/
import "C"
import "unsafe"

// A Selector represents an Objective-C method selector.
type Selector interface {
	// Selector returns a string representation of
	// a selector.
	Selector() string

	// String returns the same string as Selector does.
	// It is only implemented to implement the Stringer
	// interface.
	String() string
}

// Type selector is the underlying implementation
// of the Selector interface. It is represented as
// a Go string.
type selector string

// Selector implements the Selector method of the
// Selector interface.
func (sel selector) Selector() string {
	return string(sel)
}

// String implements the String method of the
// Selector interface.
func (sel selector) String() string {
	return sel.Selector()
}

// GetSelector looks up a Selector by name.
func GetSelector(name string) Selector {
	return selector(name)
}

// selectorWithName looks up a selector by name.
func selectorWithName(name string) unsafe.Pointer {
	return C.GoObjc_RegisterSelector(C.CString(name))
}

// typeInfoForMethod returns the type encoding string for
// selector on obj's Class.
func typeInfoForMethod(obj Object, selector string) string {
	sel := selectorWithName(selector)
	cls := getObjectClass(obj)
	return C.GoString(C.GoObjc_TypeInfoForMethod(unsafe.Pointer(cls.Pointer()), sel))
}

// simpleTypeInfoForMethod returns a simplified typeInfo representation
// with C specifiers and stack information stripped out.
func simpleTypeInfoForMethod(obj Object, selector string) string {
	ti := typeInfoForMethod(obj, selector)
	sti := []byte{}
	for i := 0; i < len(ti); i++ {
		if ti[i] >= '0' && ti[i] <= '9' {
			continue
		}
		if string(ti[i]) == encConst {
			continue
		}
		// fixme(mkrautz): What is V? The NSObject release method uses V.
		if ti[i] == 'V' {
			continue
		}
		sti = append(sti, ti[i])
	}
	return string(sti)
}
<|endoftext|>


<commit_before>package builtin

import (
	"log"
	"sync"

	"github.com/oleiade/lane"

	. "github.com/wanliu/flow/builtin/ai"
	. "github.com/wanliu/flow/context"

	config "github.com/wanliu/flow/builtin/config"
	flow "github.com/wanliu/goflow"
)

func NewApiAi() interface{} {
	return new(ApiAi)
}

type ApiAi struct {
	flow.Component

	MultiField

	sync.RWMutex

	token     string
	sessionId string
	proxyUrl  string

	CtxQueue *lane.Queue
	TxtQueue *lane.Queue

	Echo      <-chan bool
	In        <-chan string
	Token     <-chan string
	SessionId <-chan string
	ProxyUrl  <-chan string
	Out       chan<- Context
	Ctx       <-chan Context
}

func (c *ApiAi) Init() {
	c.CtxQueue = lane.NewQueue()
	c.TxtQueue = lane.NewQueue()
}

func (c *ApiAi) OnIn(input string) {
	// c.SetValue(config.ValueKeyText, input)
	c.Lock()
	c.TxtQueue.Enqueue(input)
	c.Unlock()

	c.SendQuery()
}

func (c *ApiAi) OnToken(token string) {
	c.token = token
}

func (c *ApiAi) OnSessionId(sessionId string) {
	c.sessionId = sessionId
}

func (c *ApiAi) OnProxyUrl(proxy string) {
	c.proxyUrl = proxy
}

func (c *ApiAi) OnCtx(ctx Context) {
	// c.SetValue(config.ValueKeyCtx, ctx)
	c.Lock()
	c.CtxQueue.Enqueue(ctx)
	c.Unlock()

	c.SendQuery()
}

func (c *ApiAi) SendQuery() {
	c.RLock()

	for c.CtxQueue.Head() != nil && c.TxtQueue.Head() != nil {
		txt := c.TxtQueue.Dequeue().(string)
		ctx := c.CtxQueue.Dequeue().(Context)

		res, _ := ApiAiQuery(txt, c.token, c.sessionId, c.proxyUrl)

		ctx.SetValue(config.CtxkeyResult, res)

		intent := res.Metadata.IntentName
		score := res.Score
		query := res.ResolvedQuery

		log.Printf("意图解析\"%s\" -> %s 准确度: %2.2f%%", query, intent, score*100)
		c.Out <- ctx
	}

	c.RUnlock()
}
<commit_msg>error msg for api.ai access<commit_after>package builtin

import (
	"log"
	"sync"

	"github.com/oleiade/lane"

	. "github.com/wanliu/flow/builtin/ai"
	. "github.com/wanliu/flow/context"

	config "github.com/wanliu/flow/builtin/config"
	flow "github.com/wanliu/goflow"
)

func NewApiAi() interface{} {
	return new(ApiAi)
}

type ApiAi struct {
	flow.Component

	MultiField

	sync.RWMutex

	token     string
	sessionId string
	proxyUrl  string

	CtxQueue *lane.Queue
	TxtQueue *lane.Queue

	Echo      <-chan bool
	In        <-chan string
	Token     <-chan string
	SessionId <-chan string
	ProxyUrl  <-chan string
	Out       chan<- Context
	Ctx       <-chan Context
}

func (c *ApiAi) Init() {
	c.CtxQueue = lane.NewQueue()
	c.TxtQueue = lane.NewQueue()
}

func (c *ApiAi) OnIn(input string) {
	// c.SetValue(config.ValueKeyText, input)
	c.Lock()
	c.TxtQueue.Enqueue(input)
	c.Unlock()

	c.SendQuery()
}

func (c *ApiAi) OnToken(token string) {
	c.token = token
}

func (c *ApiAi) OnSessionId(sessionId string) {
	c.sessionId = sessionId
}

func (c *ApiAi) OnProxyUrl(proxy string) {
	c.proxyUrl = proxy
}

func (c *ApiAi) OnCtx(ctx Context) {
	// c.SetValue(config.ValueKeyCtx, ctx)
	c.Lock()
	c.CtxQueue.Enqueue(ctx)
	c.Unlock()

	c.SendQuery()
}

func (c *ApiAi) SendQuery() {
	c.RLock()

	for c.CtxQueue.Head() != nil && c.TxtQueue.Head() != nil {
		txt := c.TxtQueue.Dequeue().(string)
		ctx := c.CtxQueue.Dequeue().(Context)

		res, err := ApiAiQuery(txt, c.token, c.sessionId, c.proxyUrl)

		ctx.SetValue(config.CtxkeyResult, res)

		intent := res.Metadata.IntentName
		score := res.Score
		query := res.ResolvedQuery

		if err != nil {
			log.Printf("意图解析失败:" + err.Error())
		} else {
			log.Printf("意图解析\"%s\" -> %s 准确度: %2.2f%%", query, intent, score*100)
		}

		c.Out <- ctx
	}

	c.RUnlock()
}
<|endoftext|>