<commit_before>// MIT License
//
// Copyright (c) 2019 Stefan Wichmann
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
package main

import (
	"flag"
	"os"
	"os/signal"
	"syscall"
	"time"

	log "github.com/Sirupsen/logrus"
)

var applicationVersion = "development"
var debug = flag.Bool("debug", false, "Enable debug logging")
var logFile = flag.String("log", "", "Redirect log output to specified file")
var configurationFile = flag.String("configuration", absolutePath("config.json"), "Specify the filename of the configuration to load")
var forceUpdate = flag.Bool("forceUpdate", false, "Update to new major version")
var enableWebInterface = flag.Bool("enableWebInterface", false, "Enable the web interface at startup")

var configuration *Configuration
var bridge = &HueBridge{}
var lights []*Light

const lightUpdateIntervalInSeconds = 1
const stateUpdateIntervalInSeconds = 60

func main() {
	flag.Parse()
	configureLogging()
	log.Printf("🤖 Kelvin %v starting up... 🚀", applicationVersion)
	log.Debugf("🤖 Current working directory: %v", workingDirectory())
	go CheckForUpdate(applicationVersion, *forceUpdate)
	go validateSystemTime()
	go handleSIGHUP()

	// Load configuration or create a new one
	conf, err := InitializeConfiguration(*configurationFile, *enableWebInterface)
	if err != nil {
		log.Fatal(err)
	}
	configuration = &conf

	// Start web interface
	go startInterface()

	// Find Hue bridge
	err = bridge.InitializeBridge(configuration)
	if err != nil {
		log.Warning(err)
	}

	// Find geo location
	_, err = InitializeLocation(configuration)
	if err != nil {
		log.Warning(err)
	}

	// Save configuration
	err = configuration.Write()
	if err != nil {
		log.Fatal(err)
	}

	// Initialize lights
	l, err := bridge.Lights()
	if err != nil {
		log.Warning(err)
	}
	printDevices(l)
	for _, light := range l {
		light := light

		// Filter devices we can't control
		if !light.HueLight.supportsColorTemperature() && !light.HueLight.supportsBrightness() {
			log.Printf("💡 Light %s - This device doesn't support any functionality Kelvin uses. Ignoring...", light.Name)
		} else {
			lights = append(lights, light)
			updateScheduleForLight(light)
		}
	}

	// Initialize scenes
	updateScenes()

	// Start cyclic update for all lights and scenes
	log.Debugf("🤖 Starting cyclic update...")
	lightUpdateTimer := time.NewTimer(lightTransitionIntervalInSeconds * time.Second)
	stateUpdateTick := time.Tick(stateUpdateIntervalInSeconds * time.Second)
	newDayTimer := time.After(durationUntilNextDay())
	for {
		select {
		case <-newDayTimer:
			// A new day has begun, calculate new schedule
			log.Printf("🤖 Calculating schedule for %v", time.Now().Format("Jan 2 2006"))
			for _, light := range lights {
				light := light
				updateScheduleForLight(light)
			}
			newDayTimer = time.After(durationUntilNextDay())
		case <-stateUpdateTick:
			// update interval and color every minute
			for _, light := range lights {
				light := light
				light.updateInterval()
				light.updateTargetLightState()
			}
			// update scenes
			updateScenes()
		case <-lightUpdateTimer.C:
			states, err := bridge.LightStates()
			if err != nil {
				log.Warning(err)
			}
			atLeastOneLightUpdated := false
			for _, light := range lights {
				light := light
				currentLightState, found := states[light.ID]
				if found {
					light.updateCurrentLightState(currentLightState)
					updated, err := light.update()
					if err != nil {
						log.Warning(err)
					}
					if updated {
						atLeastOneLightUpdated = true
					}
				} else {
					log.Warningf("💡 No current light state found for light %d", light.ID)
				}
			}
			nextTimerInSeconds := lightUpdateIntervalInSeconds
			if atLeastOneLightUpdated {
				log.Warningf("🤖 Updated at least one light. Awaiting transition...")
				nextTimerInSeconds += lightTransitionIntervalInSeconds
				atLeastOneLightUpdated = false
			}
			lightUpdateTimer.Reset(time.Duration(nextTimerInSeconds) * time.Second)
		}
	}
}

func updateScheduleForLight(light *Light) {
	schedule, err := configuration.lightScheduleForDay(light.ID, time.Now())
	if err != nil {
		log.Printf("💡 Light %s - Light is not associated to any schedule. Ignoring...", light.Name)
		light.Schedule = schedule // Assign empty schedule
		light.Scheduled = false
	} else {
		light.updateSchedule(schedule)
		light.updateInterval()
		light.updateTargetLightState()
	}
}

func printDevices(l []*Light) {
	log.Printf("🤖 Devices found on current bridge:")
	log.Printf("| %-32s | %3v | %-5v | %-8v | %-11v | %-5v |", "Name", "ID", "On", "Dimmable", "Temperature", "Color")
	for _, light := range l {
		log.Printf("| %-32s | %3v | %-5v | %-8v | %-11v | %-5v |", light.Name, light.ID, light.On, light.HueLight.Dimmable, light.HueLight.SupportsColorTemperature, light.HueLight.SupportsXYColor)
	}
}

func handleSIGHUP() {
	sighup := make(chan os.Signal, 1)
	signal.Notify(sighup, syscall.SIGHUP)
	<-sighup // wait for signal
	log.Printf("🤖 Received signal SIGHUP. Restarting...")
	Restart()
}

func configureLogging() {
	formatter := new(log.TextFormatter)
	formatter.FullTimestamp = true
	formatter.TimestampFormat = "2006/02/01 15:04:05"
	log.SetFormatter(formatter)
	if *debug {
		log.SetLevel(log.DebugLevel)
	}
	if logFile != nil && *logFile != "" {
		file, err := os.OpenFile(*logFile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
		if err == nil {
			log.SetOutput(file)
		} else {
			log.Info("🤖 Failed to log to file, using default stderr")
		}
	}
}

func validateSystemTime() {
	// validate local clock as it forms the basis for all time calculations.
	valid, err := IsLocalTimeValid()
	if err != nil {
		log.Fatal(err)
	}
	if !valid {
		log.Warningf("🤖 WARNING: Your local system time seems to be more than one minute off. Timings may be inaccurate.")
	} else {
		log.Debugf("🤖 Local system time validated.")
	}
}
<commit_msg>Add throtteling between api calls<commit_after>// MIT License
//
// Copyright (c) 2019 Stefan Wichmann
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
package main

import (
	"flag"
	"os"
	"os/signal"
	"syscall"
	"time"

	log "github.com/Sirupsen/logrus"
)

var applicationVersion = "development"
var debug = flag.Bool("debug", false, "Enable debug logging")
var logFile = flag.String("log", "", "Redirect log output to specified file")
var configurationFile = flag.String("configuration", absolutePath("config.json"), "Specify the filename of the configuration to load")
var forceUpdate = flag.Bool("forceUpdate", false, "Update to new major version")
var enableWebInterface = flag.Bool("enableWebInterface", false, "Enable the web interface at startup")

var configuration *Configuration
var bridge = &HueBridge{}
var lights []*Light

const lightUpdateInterval = 1 * time.Second
const stateUpdateInterval = 1 * time.Minute
const timeBetweenCalls = 200 * time.Millisecond // see https://developers.meethue.com/develop/application-design-guidance/hue-system-performance/

func main() {
	flag.Parse()
	configureLogging()
	log.Printf("🤖 Kelvin %v starting up... 🚀", applicationVersion)
	log.Debugf("🤖 Current working directory: %v", workingDirectory())
	go CheckForUpdate(applicationVersion, *forceUpdate)
	go validateSystemTime()
	go handleSIGHUP()

	// Load configuration or create a new one
	conf, err := InitializeConfiguration(*configurationFile, *enableWebInterface)
	if err != nil {
		log.Fatal(err)
	}
	configuration = &conf

	// Start web interface
	go startInterface()

	// Find Hue bridge
	err = bridge.InitializeBridge(configuration)
	if err != nil {
		log.Warning(err)
	}

	// Find geo location
	_, err = InitializeLocation(configuration)
	if err != nil {
		log.Warning(err)
	}

	// Save configuration
	err = configuration.Write()
	if err != nil {
		log.Fatal(err)
	}

	// Initialize lights
	l, err := bridge.Lights()
	if err != nil {
		log.Warning(err)
	}
	time.Sleep(timeBetweenCalls)
	printDevices(l)
	for _, light := range l {
		light := light

		// Filter devices we can't control
		if !light.HueLight.supportsColorTemperature() && !light.HueLight.supportsBrightness() {
			log.Printf("🤖 Light %s - This device doesn't support any functionality Kelvin uses. Ignoring...", light.Name)
		} else {
			lights = append(lights, light)
			updateScheduleForLight(light)
		}
	}

	// Initialize scenes
	updateScenes()
	time.Sleep(timeBetweenCalls)

	// Start cyclic update for all lights and scenes
	log.Debugf("🤖 Starting cyclic update...")
	lightUpdateTimer := time.NewTimer(lightUpdateInterval)
	stateUpdateTick := time.Tick(stateUpdateInterval)
	newDayTimer := time.After(durationUntilNextDay())
	for {
		select {
		case <-newDayTimer:
			// A new day has begun, calculate new schedule
			log.Printf("🤖 Calculating schedule for %v", time.Now().Format("Jan 2 2006"))
			for _, light := range lights {
				light := light
				updateScheduleForLight(light)
			}
			newDayTimer = time.After(durationUntilNextDay())
		case <-stateUpdateTick:
			// update interval and color every minute
			for _, light := range lights {
				light := light
				light.updateInterval()
				light.updateTargetLightState()
			}
			// update scenes
			updateScenes()
			time.Sleep(timeBetweenCalls)
		case <-lightUpdateTimer.C:
			states, err := bridge.LightStates()
			if err != nil {
				log.Warningf("🤖 Failed to update light states: %v", err)
			}
			time.Sleep(timeBetweenCalls)

			for _, light := range lights {
				light := light
				currentLightState, found := states[light.ID]
				if found {
					light.updateCurrentLightState(currentLightState)
					updated, err := light.update()
					if err != nil {
						log.Warningf("🤖 Light %s - Failed to update light: %v", light.Name, err)
					}
					if updated {
						log.Debugf("🤖 Light %s - Updated light state. Awaiting transition...", light.Name)
						time.Sleep(timeBetweenCalls)
						updated = false
					}
				} else {
					log.Warningf("🤖 Light %s - No current light state found", light.Name)
				}
			}

			lightUpdateTimer.Reset(lightUpdateInterval)
		}
	}
}

func updateScheduleForLight(light *Light) {
	schedule, err := configuration.lightScheduleForDay(light.ID, time.Now())
	if err != nil {
		log.Printf("🤖 Light %s - Light is not associated to any schedule. Ignoring...", light.Name)
		light.Schedule = schedule // Assign empty schedule
		light.Scheduled = false
	} else {
		light.updateSchedule(schedule)
		light.updateInterval()
		light.updateTargetLightState()
	}
}

func printDevices(l []*Light) {
	log.Printf("🤖 Devices found on current bridge:")
	log.Printf("| %-32s | %3v | %-5v | %-8v | %-11v | %-5v |", "Name", "ID", "On", "Dimmable", "Temperature", "Color")
	for _, light := range l {
		log.Printf("| %-32s | %3v | %-5v | %-8v | %-11v | %-5v |", light.Name, light.ID, light.On, light.HueLight.Dimmable, light.HueLight.SupportsColorTemperature, light.HueLight.SupportsXYColor)
	}
}

func handleSIGHUP() {
	sighup := make(chan os.Signal, 1)
	signal.Notify(sighup, syscall.SIGHUP)
	<-sighup // wait for signal
	log.Printf("🤖 Received signal SIGHUP. Restarting...")
	Restart()
}

func configureLogging() {
	formatter := new(log.TextFormatter)
	formatter.FullTimestamp = true
	formatter.TimestampFormat = "2006/02/01 15:04:05"
	log.SetFormatter(formatter)
	if *debug {
		log.SetLevel(log.DebugLevel)
	}
	if logFile != nil && *logFile != "" {
		file, err := os.OpenFile(*logFile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
		if err == nil {
			log.SetOutput(file)
		} else {
			log.Info("🤖 Failed to log to file, using default stderr")
		}
	}
}

func validateSystemTime() {
	// validate local clock as it forms the basis for all time calculations.
	valid, err := IsLocalTimeValid()
	if err != nil {
		log.Fatal(err)
	}
	if !valid {
		log.Warningf("🤖 WARNING: Your local system time seems to be more than one minute off. Timings may be inaccurate.")
	} else {
		log.Debugf("🤖 Local system time validated.")
	}
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor:
// - Aaron Meihm ameihm@mozilla.com

package main

import (
	"encoding/json"
	"fmt"
	elastigo "github.com/mattbaird/elastigo/lib"
	"sync"
	"time"
)

type integQueue struct {
	queue []eventResult
	sync.Mutex
}

var queue integQueue

func (i *integQueue) getResult() (eventResult, bool) {
	var ret eventResult
	i.Lock()
	defer i.Unlock()
	if len(i.queue) == 0 {
		return ret, false
	}
	ret = i.queue[0]
	i.queue = i.queue[1:]
	return ret, true
}

func (i *integQueue) addResult(n eventResult) {
	i.Lock()
	i.queue = append(i.queue, n)
	i.Unlock()
}

func mergeResults(principal string, res []eventResult) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("mergeResults() -> %v", e)
		}
	}()

	logf("merging and updating for %v", principal)
	o, err := getPrincipalState(principal)
	if err != nil {
		panic(err)
	}

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

	return nil
}

func savePrincipalState(o object) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("savePrincipalState() -> %v", e)
		}
	}()

	conn := elastigo.NewConn()
	conn.Domain = cfg.ES.StateESHost

	_, err = conn.Index(cfg.ES.StateIndex, "geomodel_state", o.ObjectID, nil, o)
	if err != nil {
		panic(err)
	}

	return nil
}

func getPrincipalState(principal string) (ret object, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("getPrincipalState() -> %v", e)
		}
	}()

	objid := getObjectID(principal)

	conn := elastigo.NewConn()
	conn.Domain = cfg.ES.StateESHost

	template := `{
		"query": {
			"term": {
				"object_id": "%v"
			}
		}
	}`
	tempbuf := fmt.Sprintf(template, objid)
	res, err := conn.Search(cfg.ES.StateIndex, "geomodel_state", nil, tempbuf)
	if err != nil {
		panic(err)
	}
	if res.Hits.Len() == 0 {
		logf("no state found for %v, creating", principal)
		ret.newFromPrincipal(principal)
		return ret, nil
	}
	err = json.Unmarshal(*res.Hits.Hits[0].Source, &ret)
	if err != nil {
		panic(err)
	}

	return ret, nil
}

func integrationMerge(exitCh chan bool) {
	defer func() {
		if e := recover(); e != nil {
			logf("integrationMerge() -> %v", e)
		}
		logf("integration merge exiting")
	}()
	logf("integration merge started")

	for {
		select {
		case <-exitCh:
			return
		case <-time.After(time.Duration(cfg.Timer.Merge) * time.Second):
		}
		logf("integration merge process running")

		princemap := make(map[string][]eventResult)
		// Fetch whatever we have queued; for efficiency group results
		// for the same principal together, reducing the number of
		// requests needed later.
		for e, ok := queue.getResult(); ok; e, ok = queue.getResult() {
			ptr, ok := princemap[e.Principal]
			if !ok {
				princemap[e.Principal] = make([]eventResult, 0)
				ptr = princemap[e.Principal]
			}
			ptr = append(ptr, e)
			princemap[e.Principal] = ptr
		}
		failed := false
		var err error
		for k, v := range princemap {
			err = mergeResults(k, v)
			if err != nil {
				failed = true
				break
			}
		}
		if failed {
			logf("integration merge failed, %v", err)
		}
	}
}

func integrate(pr pluginResult) {
	for _, x := range pr.Results {
		queue.addResult(x)
	}
}

func integrator(exitCh chan bool, notifyCh chan bool) {
	defer func() {
		if e := recover(); e != nil {
			logf("integrator() -> %v", e)
		}
		logf("integrator exiting")
	}()
	logf("integrator started")

	var iwg sync.WaitGroup
	mergeExit := make(chan bool, 1)
	iwg.Add(1)
	go func() {
		integrationMerge(mergeExit)
		iwg.Done()
	}()

	for {
		select {
		case p := <-pluginResultCh:
			integrate(p)
		case <-exitCh:
			mergeExit <- true
			iwg.Wait()
			notifyCh <- true
			return
		}
	}
}
<commit_msg>dont send invalid results for processing<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor:
// - Aaron Meihm ameihm@mozilla.com

package main

import (
	"encoding/json"
	"fmt"
	elastigo "github.com/mattbaird/elastigo/lib"
	"sync"
	"time"
)

type integQueue struct {
	queue []eventResult
	sync.Mutex
}

var queue integQueue

func (i *integQueue) getResult() (eventResult, bool) {
	var ret eventResult
	i.Lock()
	defer i.Unlock()
	if len(i.queue) == 0 {
		return ret, false
	}
	ret = i.queue[0]
	i.queue = i.queue[1:]
	return ret, true
}

func (i *integQueue) addResult(n eventResult) {
	i.Lock()
	i.queue = append(i.queue, n)
	i.Unlock()
}

func mergeResults(principal string, res []eventResult) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("mergeResults() -> %v", e)
		}
	}()

	logf("merging and updating for %v", principal)
	o, err := getPrincipalState(principal)
	if err != nil {
		panic(err)
	}

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

	return nil
}

func savePrincipalState(o object) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("savePrincipalState() -> %v", e)
		}
	}()

	conn := elastigo.NewConn()
	conn.Domain = cfg.ES.StateESHost

	_, err = conn.Index(cfg.ES.StateIndex, "geomodel_state", o.ObjectID, nil, o)
	if err != nil {
		panic(err)
	}

	return nil
}

func getPrincipalState(principal string) (ret object, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("getPrincipalState() -> %v", e)
		}
	}()

	objid := getObjectID(principal)

	conn := elastigo.NewConn()
	conn.Domain = cfg.ES.StateESHost

	template := `{
		"query": {
			"term": {
				"object_id": "%v"
			}
		}
	}`
	tempbuf := fmt.Sprintf(template, objid)
	res, err := conn.Search(cfg.ES.StateIndex, "geomodel_state", nil, tempbuf)
	if err != nil {
		panic(err)
	}
	if res.Hits.Len() == 0 {
		logf("no state found for %v, creating", principal)
		ret.newFromPrincipal(principal)
		return ret, nil
	}
	err = json.Unmarshal(*res.Hits.Hits[0].Source, &ret)
	if err != nil {
		panic(err)
	}

	return ret, nil
}

func integrationMerge(exitCh chan bool) {
	defer func() {
		if e := recover(); e != nil {
			logf("integrationMerge() -> %v", e)
		}
		logf("integration merge exiting")
	}()
	logf("integration merge started")

	for {
		select {
		case <-exitCh:
			return
		case <-time.After(time.Duration(cfg.Timer.Merge) * time.Second):
		}
		logf("integration merge process running")

		princemap := make(map[string][]eventResult)
		// Fetch whatever we have queued; for efficiency group results
		// for the same principal together, reducing the number of
		// requests needed later.
		for e, ok := queue.getResult(); ok; e, ok = queue.getResult() {
			ptr, ok := princemap[e.Principal]
			if !ok {
				princemap[e.Principal] = make([]eventResult, 0)
				ptr = princemap[e.Principal]
			}
			ptr = append(ptr, e)
			princemap[e.Principal] = ptr
		}
		failed := false
		var err error
		for k, v := range princemap {
			err = mergeResults(k, v)
			if err != nil {
				failed = true
				break
			}
		}
		if failed {
			logf("integration merge failed, %v", err)
		}
	}
}

func integrate(pr pluginResult) {
	for _, x := range pr.Results {
		if !x.Valid {
			logf("ignoring invalid result from plugin")
			continue
		}
		queue.addResult(x)
	}
}

func integrator(exitCh chan bool, notifyCh chan bool) {
	defer func() {
		if e := recover(); e != nil {
			logf("integrator() -> %v", e)
		}
		logf("integrator exiting")
	}()
	logf("integrator started")

	var iwg sync.WaitGroup
	mergeExit := make(chan bool, 1)
	iwg.Add(1)
	go func() {
		integrationMerge(mergeExit)
		iwg.Done()
	}()

	for {
		select {
		case p := <-pluginResultCh:
			integrate(p)
		case <-exitCh:
			mergeExit <- true
			iwg.Wait()
			notifyCh <- true
			return
		}
	}
}
<|endoftext|>


<commit_before>package iobit

import (
	"encoding/binary"
	"errors"
	"io"
)

type Writer struct {
	cache uint64
	dst   io.Writer
	fill  uint
	err   error
}

type bigEndian struct{}
type littleEndian struct{}

var (
	BigEndian    bigEndian
	LittleEndian littleEndian
)

func NewWriter(dst io.Writer) *Writer {
	return &Writer{dst: dst}
}

func (w *Writer) flushCache(bits uint) {
	if w.fill+bits <= 64 {
		return
	}
	var data [4]uint8
	binary.BigEndian.PutUint32(data[:], uint32(w.cache>>32))
	w.cache <<= 32
	w.fill -= 32
	w.write(data[:])
}

func (w *Writer) writeCache(bits uint, val uint32) {
	u := uint64(val)
	u &= ^(^uint64(0) << bits)
	u <<= 64 - w.fill - bits
	w.cache |= u
	w.fill += bits
}

func (bigEndian) PutUint32(w *Writer, bits uint, val uint32) {
	w.flushCache(bits)
	w.writeCache(bits, val)
}

func (littleEndian) PutUint32(w *Writer, bits uint, val uint32) {
	w.flushCache(bits)
	for bits > 8 {
		w.writeCache(8, val)
		val >>= 8
		bits -= 8
	}
	w.writeCache(bits, val)
}

func (bigEndian) PutUint64(w *Writer, bits uint, val uint64) {
	if bits > 32 {
		BigEndian.PutUint32(w, bits-32, uint32(val>>32))
		bits = 32
		val &= 0xFFFFFFFF
	}
	BigEndian.PutUint32(w, bits, uint32(val))
}

func (littleEndian) PutUint64(w *Writer, bits uint, val uint64) {
	if bits > 32 {
		LittleEndian.PutUint32(w, bits-32, uint32(val&0xFFFFFFFF))
		bits = 32
		val >>= 32
	}
	LittleEndian.PutUint32(w, bits, uint32(val))
}

func (w *Writer) write(data []uint8) {
	if w.err == nil {
		_, w.err = w.dst.Write(data)
	}
}

func (w *Writer) Flush() error {
	var data [8]uint8
	idx := 0
	for w.fill >= 8 {
		data[idx] = uint8(w.cache >> 56)
		w.cache <<= 8
		w.fill -= 8
		idx++
	}
	if w.fill != 0 {
		w.err = errors.New("iobit: unable to flush unaligned output")
	}
	w.write(data[:])
	return w.err
}

func (w *Writer) Write(p []uint8) (int, error) {
	err := w.Flush()
	if err != nil {
		return 0, err
	}
	return w.dst.Write(p)
}
<commit_msg>writer: export ErrUnderflow<commit_after>package iobit

import (
	"encoding/binary"
	"errors"
	"io"
)

type Writer struct {
	cache uint64
	dst   io.Writer
	fill  uint
	err   error
}

type bigEndian struct{}
type littleEndian struct{}

var (
	ErrUnderflow = errors.New("bit underflow")
	BigEndian    bigEndian
	LittleEndian littleEndian
)

func NewWriter(dst io.Writer) *Writer {
	return &Writer{dst: dst}
}

func (w *Writer) flushCache(bits uint) {
	if w.fill+bits <= 64 {
		return
	}
	var data [4]uint8
	binary.BigEndian.PutUint32(data[:], uint32(w.cache>>32))
	w.cache <<= 32
	w.fill -= 32
	w.write(data[:])
}

func (w *Writer) writeCache(bits uint, val uint32) {
	u := uint64(val)
	u &= ^(^uint64(0) << bits)
	u <<= 64 - w.fill - bits
	w.cache |= u
	w.fill += bits
}

func (bigEndian) PutUint32(w *Writer, bits uint, val uint32) {
	w.flushCache(bits)
	w.writeCache(bits, val)
}

func (littleEndian) PutUint32(w *Writer, bits uint, val uint32) {
	w.flushCache(bits)
	for bits > 8 {
		w.writeCache(8, val)
		val >>= 8
		bits -= 8
	}
	w.writeCache(bits, val)
}

func (bigEndian) PutUint64(w *Writer, bits uint, val uint64) {
	if bits > 32 {
		BigEndian.PutUint32(w, bits-32, uint32(val>>32))
		bits = 32
		val &= 0xFFFFFFFF
	}
	BigEndian.PutUint32(w, bits, uint32(val))
}

func (littleEndian) PutUint64(w *Writer, bits uint, val uint64) {
	if bits > 32 {
		LittleEndian.PutUint32(w, bits-32, uint32(val&0xFFFFFFFF))
		bits = 32
		val >>= 32
	}
	LittleEndian.PutUint32(w, bits, uint32(val))
}

func (w *Writer) write(data []uint8) {
	if w.err == nil {
		_, w.err = w.dst.Write(data)
	}
}

func (w *Writer) Flush() error {
	var data [8]uint8
	idx := 0
	for w.fill >= 8 {
		data[idx] = uint8(w.cache >> 56)
		w.cache <<= 8
		w.fill -= 8
		idx++
	}
	if w.fill != 0 {
		w.err = ErrUnderflow
	}
	w.write(data[:])
	return w.err
}

func (w *Writer) Write(p []uint8) (int, error) {
	err := w.Flush()
	if err != nil {
		return 0, err
	}
	return w.dst.Write(p)
}
<|endoftext|>


<commit_before>package indexer

import (
	"fmt"
	"os"
	"path/filepath"
	"sync"

	"github.com/deepfabric/bkdtree"
	"github.com/deepfabric/indexer/cql"
	"github.com/deepfabric/pilosa"
	"github.com/pkg/errors"
)

//Conf originate from config file. It's comman to all indices.
type Conf struct {
	//BKD specific items
	T0mCap   int
	LeafCap  int
	IntraCap int
}

//Indexer shall be singleton
type Indexer struct {
	MainDir string //the main directory where stores all indices
	Conf    Conf   //indexer conf

	rwlock   sync.RWMutex                    //concurrent access of docProts, indices
	docProts map[string]*cql.DocumentWithIdx //index meta, need to persist
	indices  map[string]*Index               //index data, need to persist
}

//NewIndexer creates an Indexer.
func NewIndexer(mainDir string, conf *Conf, overwirte bool) (ir *Indexer, err error) {
	ir = &Indexer{
		MainDir: mainDir,
		Conf:    *conf,
	}
	if err = os.MkdirAll(mainDir, 0700); err != nil {
		err = errors.Wrap(err, "")
		return
	}
	if overwirte {
		ir.docProts = make(map[string]*cql.DocumentWithIdx)
		ir.indices = make(map[string]*Index)
		err = ir.removeIndices()
	} else {
		err = ir.Open()
	}
	return
}

//Destroy close and remove index files
func (ir *Indexer) Destroy() (err error) {
	if err = ir.Close(); err != nil {
		return
	}
	if err = ir.removeIndices(); err != nil {
		return
	}
	return
}

//Open opens all indices. Assumes ir.MainDir is already populated.
func (ir *Indexer) Open() (err error) {
	ir.rwlock.Lock()
	defer ir.rwlock.Unlock()
	if ir.indices != nil || ir.docProts != nil {
		panic("indexer already open")
	}
	ir.docProts = make(map[string]*cql.DocumentWithIdx)
	ir.indices = make(map[string]*Index)
	if err = ir.readMeta(); err != nil {
		return
	}
	var ind *Index
	for name, docProt := range ir.docProts {
		if ind, err = NewIndexExt(ir.MainDir, docProt.Index); err != nil {
			return
		}
		ir.indices[name] = ind
	}
	return
}

// Close close indexer
func (ir *Indexer) Close() (err error) {
	ir.rwlock.Lock()
	defer ir.rwlock.Unlock()
	for _, ind := range ir.indices {
		if err = ind.Close(); err != nil {
			return
		}
	}
	ir.indices = nil
	ir.docProts = nil
	return
}

// GetDocProt returns docProt of given index
func (ir *Indexer) GetDocProt(name string) (docProt *cql.DocumentWithIdx) {
	ir.rwlock.RLock()
	docProt, _ = ir.docProts[name]
	ir.rwlock.RUnlock()
	return
}

// CreateIndex creates index
func (ir *Indexer) CreateIndex(docProt *cql.DocumentWithIdx) (err error) {
	ir.rwlock.Lock()
	defer ir.rwlock.Unlock()
	if _, found := ir.docProts[docProt.Index]; found {
		panic("CreateIndex conflict with existing index")
	}
	if err = indexWriteConf(ir.MainDir, docProt); err != nil {
		return
	}
	var ind *Index
	if ind, err = NewIndex(docProt, ir.MainDir, ir.Conf.T0mCap, ir.Conf.LeafCap, ir.Conf.IntraCap); err != nil {
		return
	}
	ir.indices[docProt.Index] = ind
	ir.docProts[docProt.Index] = docProt
	return
}

//DestroyIndex destroy given index
func (ir *Indexer) DestroyIndex(name string) (err error) {
	ir.rwlock.Lock()
	defer ir.rwlock.Unlock()
	delete(ir.indices, name)
	delete(ir.docProts, name)
	err = ir.removeIndex(name)
	return
}

//Insert executes CqlInsert
func (ir *Indexer) Insert(doc *cql.DocumentWithIdx) (err error) {
	var ind *Index
	var found bool
	ir.rwlock.RLock()
	if ind, found = ir.indices[doc.Index]; !found {
		err = errors.Errorf("failed to insert %v to non-existing index %v", doc, doc.Index)
		ir.rwlock.RUnlock()
		return
	}
	ir.rwlock.RUnlock()
	err = ind.Insert(doc)
	return
}

//Del executes CqlDel.
func (ir *Indexer) Del(idxName string, docID uint64) (found bool, err error) {
	var ind *Index
	var fnd bool
	ir.rwlock.RLock()
	if ind, fnd = ir.indices[idxName]; !fnd {
		err = errors.Errorf("failed to delete %v from non-existing index %v", docID, idxName)
		ir.rwlock.RUnlock()
		return
	}
	ir.rwlock.RUnlock()
	found, err = ind.Del(docID)
	return
}

//Select executes CqlSelect.
func (ir *Indexer) Select(q *cql.CqlSelect) (rb *pilosa.Bitmap, err error) {
	var ind *Index
	var found bool
	ir.rwlock.RLock()
	if ind, found = ir.indices[q.Index]; !found {
		err = errors.Errorf("failed to select %v from non-existing index %v", q, q.Index)
		ir.rwlock.RUnlock()
		return
	}
	ir.rwlock.RUnlock()
	rb, err = ind.Select(q)
	return
}

//writeMeta persists Conf and DocProts to files.
func (ir *Indexer) writeMeta() (err error) {
	for _, docProt := range ir.docProts {
		if err = indexWriteConf(ir.MainDir, docProt); err != nil {
			return
		}
	}
	return
}

//readMeta parses Conf and DocProts from files.
func (ir *Indexer) readMeta() (err error) {
	var matches [][]string
	patt := `^index_(?P<name>[^.]+)\.json$`
	if matches, err = bkdtree.FilepathGlob(ir.MainDir, patt); err != nil {
		return
	}
	for _, match := range matches {
		var doc cql.DocumentWithIdx
		if err = indexReadConf(ir.MainDir, match[1], &doc); err != nil {
			return
		}
		ir.docProts[match[1]] = &doc
	}
	return
}

func (ir *Indexer) removeIndices() (err error) {
	var matches [][]string
	patt := `^index_(?P<name>[^.]+)\.json$`
	if matches, err = bkdtree.FilepathGlob(ir.MainDir, patt); err != nil {
		return
	}
	for _, match := range matches {
		if err = ir.removeIndex(match[1]); err != nil {
			return
		}
	}
	return
}

func (ir *Indexer) removeIndex(name string) (err error) {
	var fp string
	fp = filepath.Join(ir.MainDir, fmt.Sprintf("index_%s.json", name))
	if err = os.Remove(fp); err != nil {
		err = errors.Wrap(err, "")
		return
	}
	fp = filepath.Join(ir.MainDir, name)
	if err = os.RemoveAll(fp); err != nil {
		err = errors.Wrap(err, "")
	}
	return
}
<commit_msg>Made (*Indexer).Insert creates index  implictily<commit_after>package indexer

import (
	"fmt"
	"os"
	"path/filepath"
	"sync"

	"github.com/deepfabric/bkdtree"
	"github.com/deepfabric/indexer/cql"
	"github.com/deepfabric/pilosa"
	"github.com/pkg/errors"
)

//Conf originate from config file. It's comman to all indices.
type Conf struct {
	//BKD specific items
	T0mCap   int
	LeafCap  int
	IntraCap int
}

//Indexer shall be singleton
type Indexer struct {
	MainDir string //the main directory where stores all indices
	Conf    Conf   //indexer conf

	rwlock   sync.RWMutex                    //concurrent access of docProts, indices
	docProts map[string]*cql.DocumentWithIdx //index meta, need to persist
	indices  map[string]*Index               //index data, need to persist
}

//NewIndexer creates an Indexer.
func NewIndexer(mainDir string, conf *Conf, overwirte bool) (ir *Indexer, err error) {
	ir = &Indexer{
		MainDir: mainDir,
		Conf:    *conf,
	}
	if err = os.MkdirAll(mainDir, 0700); err != nil {
		err = errors.Wrap(err, "")
		return
	}
	if overwirte {
		ir.docProts = make(map[string]*cql.DocumentWithIdx)
		ir.indices = make(map[string]*Index)
		err = ir.removeIndices()
	} else {
		err = ir.Open()
	}
	return
}

//Destroy close and remove index files
func (ir *Indexer) Destroy() (err error) {
	if err = ir.Close(); err != nil {
		return
	}
	if err = ir.removeIndices(); err != nil {
		return
	}
	return
}

//Open opens all indices. Assumes ir.MainDir is already populated.
func (ir *Indexer) Open() (err error) {
	ir.rwlock.Lock()
	defer ir.rwlock.Unlock()
	if ir.indices != nil || ir.docProts != nil {
		panic("indexer already open")
	}
	ir.docProts = make(map[string]*cql.DocumentWithIdx)
	ir.indices = make(map[string]*Index)
	if err = ir.readMeta(); err != nil {
		return
	}
	var ind *Index
	for name, docProt := range ir.docProts {
		if ind, err = NewIndexExt(ir.MainDir, docProt.Index); err != nil {
			return
		}
		ir.indices[name] = ind
	}
	return
}

// Close close indexer
func (ir *Indexer) Close() (err error) {
	ir.rwlock.Lock()
	defer ir.rwlock.Unlock()
	for _, ind := range ir.indices {
		if err = ind.Close(); err != nil {
			return
		}
	}
	ir.indices = nil
	ir.docProts = nil
	return
}

// GetDocProt returns docProt of given index
func (ir *Indexer) GetDocProt(name string) (docProt *cql.DocumentWithIdx) {
	ir.rwlock.RLock()
	docProt, _ = ir.docProts[name]
	ir.rwlock.RUnlock()
	return
}

// CreateIndex creates index
func (ir *Indexer) CreateIndex(docProt *cql.DocumentWithIdx) (err error) {
	ir.rwlock.Lock()
	err = ir.createIndex(docProt)
	ir.rwlock.Unlock()
	return
}

//DestroyIndex destroy given index
func (ir *Indexer) DestroyIndex(name string) (err error) {
	ir.rwlock.Lock()
	defer ir.rwlock.Unlock()
	delete(ir.indices, name)
	delete(ir.docProts, name)
	err = ir.removeIndex(name)
	return
}

//Insert executes CqlInsert. If the given index doesn't exist, create it before insertion.
func (ir *Indexer) Insert(doc *cql.DocumentWithIdx) (err error) {
	var ind *Index
	var found bool
	ir.rwlock.RLock()
	if ind, found = ir.indices[doc.Index]; !found {
		if err = ir.createIndex(doc); err != nil {
			ir.rwlock.RUnlock()
			return
		}
		ind, found = ir.indices[doc.Index]
	}
	ir.rwlock.RUnlock()
	err = ind.Insert(doc)
	return
}

//Del executes CqlDel. It's allowed that the given index doesn't exist.
func (ir *Indexer) Del(idxName string, docID uint64) (found bool, err error) {
	var ind *Index
	var fnd bool
	ir.rwlock.RLock()
	if ind, fnd = ir.indices[idxName]; !fnd {
		ir.rwlock.RUnlock()
		return
	}
	ir.rwlock.RUnlock()
	found, err = ind.Del(docID)
	return
}

//Select executes CqlSelect.
func (ir *Indexer) Select(q *cql.CqlSelect) (rb *pilosa.Bitmap, err error) {
	var ind *Index
	var found bool
	ir.rwlock.RLock()
	if ind, found = ir.indices[q.Index]; !found {
		err = errors.Errorf("failed to select %v from non-existing index %v", q, q.Index)
		ir.rwlock.RUnlock()
		return
	}
	ir.rwlock.RUnlock()
	rb, err = ind.Select(q)
	return
}

// createIndex creates index without holding the lock
func (ir *Indexer) createIndex(docProt *cql.DocumentWithIdx) (err error) {
	if _, found := ir.docProts[docProt.Index]; found {
		err = errors.New("CreateIndex conflict with existing index")
		return
	}
	if err = indexWriteConf(ir.MainDir, docProt); err != nil {
		return
	}
	var ind *Index
	if ind, err = NewIndex(docProt, ir.MainDir, ir.Conf.T0mCap, ir.Conf.LeafCap, ir.Conf.IntraCap); err != nil {
		return
	}
	ir.indices[docProt.Index] = ind
	ir.docProts[docProt.Index] = docProt
	return
}

//writeMeta persists Conf and DocProts to files.
func (ir *Indexer) writeMeta() (err error) {
	for _, docProt := range ir.docProts {
		if err = indexWriteConf(ir.MainDir, docProt); err != nil {
			return
		}
	}
	return
}

//readMeta parses Conf and DocProts from files.
func (ir *Indexer) readMeta() (err error) {
	var matches [][]string
	patt := `^index_(?P<name>[^.]+)\.json$`
	if matches, err = bkdtree.FilepathGlob(ir.MainDir, patt); err != nil {
		return
	}
	for _, match := range matches {
		var doc cql.DocumentWithIdx
		if err = indexReadConf(ir.MainDir, match[1], &doc); err != nil {
			return
		}
		ir.docProts[match[1]] = &doc
	}
	return
}

func (ir *Indexer) removeIndices() (err error) {
	var matches [][]string
	patt := `^index_(?P<name>[^.]+)\.json$`
	if matches, err = bkdtree.FilepathGlob(ir.MainDir, patt); err != nil {
		return
	}
	for _, match := range matches {
		if err = ir.removeIndex(match[1]); err != nil {
			return
		}
	}
	return
}

func (ir *Indexer) removeIndex(name string) (err error) {
	var fp string
	fp = filepath.Join(ir.MainDir, fmt.Sprintf("index_%s.json", name))
	if err = os.Remove(fp); err != nil {
		err = errors.Wrap(err, "")
		return
	}
	fp = filepath.Join(ir.MainDir, name)
	if err = os.RemoveAll(fp); err != nil {
		err = errors.Wrap(err, "")
	}
	return
}
<|endoftext|>


<commit_before>package main

/*
#include <stdlib.h>
#include <string.h>
typedef enum {
  PADDLE_ELEMENT_TYPE_INT32   = 0,
  PADDLE_ELEMENT_TYPE_UINT32  = 1,
  PADDLE_ELEMENT_TYPE_INT64   = 2,
  PADDLE_ELEMENT_TYPE_UINT64  = 3,
  PADDLE_ELEMENT_TYPE_FLOAT32 = 4,
  PADDLE_ELEMENT_TYPE_FLOAT64 = 5,
} paddle_element_type;

typedef struct {
  char*               name;
  paddle_element_type element_type;
  char*               content;
  int                 content_len;
} paddle_parameter, paddle_gradient;

static inline void paddle_release_param(paddle_parameter* param) {
  if (param != NULL) {
    if (param->name != NULL) {
      free(param->name);
    }

    if (param->content != NULL) {
      free(param->content);
    }

    free(param);
  }
}

typedef int client;
*/
import "C"

import (
	"log"
	"sync"
	"unsafe"

	"github.com/PaddlePaddle/Paddle/paddle/go/pserver"
)

var nullPtr = unsafe.Pointer(uintptr(0))
var mu sync.Mutex
var handleMap = make(map[C.client]*pserver.Client)
var curHandle C.client

func add(c *pserver.Client) C.client {
	mu.Lock()
	defer mu.Unlock()
	client := curHandle
	curHandle++
	handleMap[client] = c
	return client
}

func get(client C.client) *pserver.Client {
	mu.Lock()
	defer mu.Unlock()
	return handleMap[client]
}

func remove(client C.client) *pserver.Client {
	mu.Lock()
	defer mu.Unlock()
	h := handleMap[client]
	delete(handleMap, client)
	return h
}

func cArrayToSlice(p unsafe.Pointer, len int) []byte {
	if p == nullPtr {
		return nil
	}

	// create a Go clice backed by a C array,
	// reference: https://github.com/golang/go/wiki/cgo#turning-c-arrays-into-go-slices
	return (*[1 << 30]byte)(p)[:len:len]
}

//export paddle_new_pserver_client
func paddle_new_pserver_client(addr *C.char) C.client {
	c := pserver.NewClient(C.GoString(addr))
	return add(c)
}

//export paddle_pserver_client_release
func paddle_pserver_client_release(client C.client) {
	c := remove(client)
	c.Cleanup()
}

//export paddle_begin_init_params
func paddle_begin_init_params(client C.client, pserver_config unsafe.Pointer, config_len C.int) C.int {
	c := get(client)
	b := cArrayToSlice(pserver_config, int(config_len))
	selected, err := c.BeginInitParams(b)
	if err != nil {
		log.Println(err)
		return -1
	}

	if selected {
		return 1
	}
	return 0
}

//export paddle_init_param
func paddle_init_param(client C.client, param C.paddle_parameter, param_config unsafe.Pointer, config_len C.int) C.int {
	et := pserver.ElementType(param.element_type)
	name := C.GoString(param.name)
	content := cArrayToSlice(unsafe.Pointer(param.content), int(param.content_len))
	pc := pserver.ParameterWithConfig{
		Param:  pserver.Parameter{Name: name, ElementType: et, Content: content},
		Config: cArrayToSlice(param_config, int(config_len)),
	}
	c := get(client)
	err := c.InitParam(pc)
	if err != nil {
		log.Println(err)
		return -1
	}

	return 0
}

//export paddle_finish_init_params
func paddle_finish_init_params(client C.client) C.int {
	c := get(client)
	err := c.FinishInitParams()
	if err != nil {
		log.Println(err)
		return -1
	}

	return 0
}

//export paddle_send_grads
func paddle_send_grads(client C.client, grads *C.paddle_gradient, total C.int) C.int {
	var gs []pserver.Gradient
	for i := 0; i < int(total); i++ {
		grad := (*C.paddle_gradient)(unsafe.Pointer((uintptr(unsafe.Pointer(grads)) + uintptr(i)*unsafe.Sizeof(*grads))))
		et := pserver.ElementType(grad.element_type)
		name := C.GoString(grad.name)
		content := cArrayToSlice(unsafe.Pointer(grad.content), int(grad.content_len))
		gs = append(gs, pserver.Gradient{Name: name, ElementType: et, Content: content})
	}

	c := get(client)
	err := c.SendGrads(gs)
	if err != nil {
		log.Println(err)
		return -1
	}

	return 0
}

//export paddle_get_params
func paddle_get_params(client C.client, names **C.char, dst **C.paddle_parameter, total C.int) C.int {
	var ns []string
	for i := 0; i < int(total); i++ {
		name := *(**C.char)(unsafe.Pointer((uintptr(unsafe.Pointer(names)) + uintptr(i)*unsafe.Sizeof(*names))))
		ns = append(ns, C.GoString(name))
	}
	c := get(client)
	ps, err := c.GetParams(ns)
	if err != nil {
		log.Println(err)
		return -1
	}

	for i := 0; i < int(total); i++ {
		if i >= len(ps) {
			break
		}

		p := ps[i]
		param := *(**C.paddle_parameter)(unsafe.Pointer((uintptr(unsafe.Pointer(dst)) + uintptr(i)*unsafe.Sizeof(*dst))))
		nameReady := false
		contentAllocated := false

		if unsafe.Pointer(param) == nullPtr {
			param = (*C.paddle_parameter)(C.calloc(1, C.size_t(unsafe.Sizeof(*param))))
		} else {
			if unsafe.Pointer(param.name) != nullPtr {
				if n := C.GoString(param.name); n != p.Name {
					log.Println("Warning: the pre-allocated parameter name does not match the parameter name, it will be freed.", n, p.Name)
					C.free(unsafe.Pointer(param.name))
				} else {
					nameReady = true
				}
			}

			if unsafe.Pointer(param.content) != nullPtr {
				if int(param.content_len) == len(p.Content) {
					contentAllocated = true
				} else {
					log.Println("Warning: the pre-allocated content len does not match parameter content len, the pre-allocated content will be freed.", param.content_len, len(p.Content))
					C.free(unsafe.Pointer(param.content))
				}
			}
		}

		if !nameReady {
			param.name = C.CString(p.Name)
		}
		if !contentAllocated {
			param.content = (*C.char)(C.malloc(C.size_t(len(p.Content))))
		}
		C.memcpy(unsafe.Pointer(param.content), unsafe.Pointer(&p.Content[0]), C.size_t(len(p.Content)))
		param.content_len = C.int(len(p.Content))
		param.element_type = C.paddle_element_type(p.ElementType)
	}

	return 0
}

//export paddle_save_model
func paddle_save_model(client C.client, path *C.char) C.int {
	p := C.GoString(path)
	c := get(client)
	err := c.SaveModel(p)
	if err != nil {
		log.Println(err)
		return -1
	}

	return 0
}

func main() {} // Required but ignored
<commit_msg>use unsigned char* for parameter.content<commit_after>package main

/*
#include <stdlib.h>
#include <string.h>
typedef enum {
  PADDLE_ELEMENT_TYPE_INT32   = 0,
  PADDLE_ELEMENT_TYPE_UINT32  = 1,
  PADDLE_ELEMENT_TYPE_INT64   = 2,
  PADDLE_ELEMENT_TYPE_UINT64  = 3,
  PADDLE_ELEMENT_TYPE_FLOAT32 = 4,
  PADDLE_ELEMENT_TYPE_FLOAT64 = 5,
} paddle_element_type;

typedef struct {
  char*               name;
  paddle_element_type element_type;
  unsigned char*      content;
  int                 content_len;
} paddle_parameter, paddle_gradient;

static inline void paddle_release_param(paddle_parameter* param) {
  if (param != NULL) {
    if (param->name != NULL) {
      free(param->name);
    }

    if (param->content != NULL) {
      free(param->content);
    }

    free(param);
  }
}

typedef int client;
*/
import "C"

import (
	"log"
	"sync"
	"unsafe"

	"github.com/PaddlePaddle/Paddle/paddle/go/pserver"
)

var nullPtr = unsafe.Pointer(uintptr(0))
var mu sync.Mutex
var handleMap = make(map[C.client]*pserver.Client)
var curHandle C.client

func add(c *pserver.Client) C.client {
	mu.Lock()
	defer mu.Unlock()
	client := curHandle
	curHandle++
	handleMap[client] = c
	return client
}

func get(client C.client) *pserver.Client {
	mu.Lock()
	defer mu.Unlock()
	return handleMap[client]
}

func remove(client C.client) *pserver.Client {
	mu.Lock()
	defer mu.Unlock()
	h := handleMap[client]
	delete(handleMap, client)
	return h
}

func cArrayToSlice(p unsafe.Pointer, len int) []byte {
	if p == nullPtr {
		return nil
	}

	// create a Go clice backed by a C array,
	// reference: https://github.com/golang/go/wiki/cgo#turning-c-arrays-into-go-slices
	return (*[1 << 30]byte)(p)[:len:len]
}

//export paddle_new_pserver_client
func paddle_new_pserver_client(addr *C.char) C.client {
	c := pserver.NewClient(C.GoString(addr))
	return add(c)
}

//export paddle_pserver_client_release
func paddle_pserver_client_release(client C.client) {
	c := remove(client)
	c.Cleanup()
}

//export paddle_begin_init_params
func paddle_begin_init_params(client C.client, pserver_config unsafe.Pointer, config_len C.int) C.int {
	c := get(client)
	b := cArrayToSlice(pserver_config, int(config_len))
	selected, err := c.BeginInitParams(b)
	if err != nil {
		log.Println(err)
		return -1
	}

	if selected {
		return 1
	}
	return 0
}

//export paddle_init_param
func paddle_init_param(client C.client, param C.paddle_parameter, param_config unsafe.Pointer, config_len C.int) C.int {
	et := pserver.ElementType(param.element_type)
	name := C.GoString(param.name)
	content := cArrayToSlice(unsafe.Pointer(param.content), int(param.content_len))
	pc := pserver.ParameterWithConfig{
		Param:  pserver.Parameter{Name: name, ElementType: et, Content: content},
		Config: cArrayToSlice(param_config, int(config_len)),
	}
	c := get(client)
	err := c.InitParam(pc)
	if err != nil {
		log.Println(err)
		return -1
	}

	return 0
}

//export paddle_finish_init_params
func paddle_finish_init_params(client C.client) C.int {
	c := get(client)
	err := c.FinishInitParams()
	if err != nil {
		log.Println(err)
		return -1
	}

	return 0
}

//export paddle_send_grads
func paddle_send_grads(client C.client, grads *C.paddle_gradient, total C.int) C.int {
	var gs []pserver.Gradient
	for i := 0; i < int(total); i++ {
		grad := (*C.paddle_gradient)(unsafe.Pointer((uintptr(unsafe.Pointer(grads)) + uintptr(i)*unsafe.Sizeof(*grads))))
		et := pserver.ElementType(grad.element_type)
		name := C.GoString(grad.name)
		content := cArrayToSlice(unsafe.Pointer(grad.content), int(grad.content_len))
		gs = append(gs, pserver.Gradient{Name: name, ElementType: et, Content: content})
	}

	c := get(client)
	err := c.SendGrads(gs)
	if err != nil {
		log.Println(err)
		return -1
	}

	return 0
}

//export paddle_get_params
func paddle_get_params(client C.client, names **C.char, dst **C.paddle_parameter, total C.int) C.int {
	var ns []string
	for i := 0; i < int(total); i++ {
		name := *(**C.char)(unsafe.Pointer((uintptr(unsafe.Pointer(names)) + uintptr(i)*unsafe.Sizeof(*names))))
		ns = append(ns, C.GoString(name))
	}
	c := get(client)
	ps, err := c.GetParams(ns)
	if err != nil {
		log.Println(err)
		return -1
	}

	for i := 0; i < int(total); i++ {
		if i >= len(ps) {
			break
		}

		p := ps[i]
		param := *(**C.paddle_parameter)(unsafe.Pointer((uintptr(unsafe.Pointer(dst)) + uintptr(i)*unsafe.Sizeof(*dst))))
		nameReady := false
		contentAllocated := false

		if unsafe.Pointer(param) == nullPtr {
			param = (*C.paddle_parameter)(C.calloc(1, C.size_t(unsafe.Sizeof(*param))))
		} else {
			if unsafe.Pointer(param.name) != nullPtr {
				if n := C.GoString(param.name); n != p.Name {
					log.Println("Warning: the pre-allocated parameter name does not match the parameter name, it will be freed.", n, p.Name)
					C.free(unsafe.Pointer(param.name))
				} else {
					nameReady = true
				}
			}

			if unsafe.Pointer(param.content) != nullPtr {
				if int(param.content_len) == len(p.Content) {
					contentAllocated = true
				} else {
					log.Println("Warning: the pre-allocated content len does not match parameter content len, the pre-allocated content will be freed.", param.content_len, len(p.Content))
					C.free(unsafe.Pointer(param.content))
				}
			}
		}

		if !nameReady {
			param.name = C.CString(p.Name)
		}
		if !contentAllocated {
			param.content = (*C.uchar)(C.malloc(C.size_t(len(p.Content))))
		}
		C.memcpy(unsafe.Pointer(param.content), unsafe.Pointer(&p.Content[0]), C.size_t(len(p.Content)))
		param.content_len = C.int(len(p.Content))
		param.element_type = C.paddle_element_type(p.ElementType)
	}

	return 0
}

//export paddle_save_model
func paddle_save_model(client C.client, path *C.char) C.int {
	p := C.GoString(path)
	c := get(client)
	err := c.SaveModel(p)
	if err != nil {
		log.Println(err)
		return -1
	}

	return 0
}

func main() {} // Required but ignored
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/user"
	"path"
	"strings"

	"github.com/davidrjonas/ssh-iam-bridge/string_array"
)

func check(err error) {
	if err == nil {
		return
	}

	if exerr, ok := err.(*exec.ExitError); ok {
		os.Stderr.Write(exerr.Stderr)
	}

	panic(err)
}

func backupFile(filename string) {
	fmt.Println("Backing up", filename, "to", filename+".orig")
	err := exec.Command("cp", "-f", filename, filename+".orig").Run()
	check(err)
}

func install(selfPath, username string) {
	cmdName := installAuthorizedKeysCommandScript(selfPath)
	installUser(username)
	installToSshd(cmdName, username)
	installToPam(selfPath)
	installToCron(selfPath)
}

// ssh is picky about AuthorizedKeysCommand, see man sshd_config
func installAuthorizedKeysCommandScript(selfPath string) string {
	cmdName := "/usr/sbin/ssh-iam-bridge-public-keys"
	fmt.Println("Writing AuthorizedKeysCommand script", cmdName)

	script := fmt.Sprintf("#!/bin/sh\nexec %s authorized_keys \"$@\"\n", selfPath)

	check(ioutil.WriteFile(cmdName, []byte(script), 0755))

	return cmdName
}

func installUser(username string) {
	_, err := user.Lookup(username)

	if err == nil {
		// User already exists
		return
	}

	if _, ok := err.(user.UnknownUserError); !ok {
		panic(err)
	}

	fmt.Println("Creating SSH authorized keys lookup user", username)

	args := []string{
		"--system",
		"--shell", "/usr/sbin/nologin",
		"--comment", "SSH authorized keys lookup",
		username,
	}

	_, err = exec.Command("useradd", args...).Output()
	check(err)
}

func installToSshd(cmd, username string) {

	filename := "/etc/ssh/sshd_config"

	// TODO: Ensure 'PasswordAuthentication no' and 'UsePAM yes'

	lines_to_add := []string{
		"AuthorizedKeysCommand " + cmd + "\n",
		"AuthorizedKeysCommandUser " + username + "\n",
		"ChallengeResponseAuthentication yes\n",
		"AuthenticationMethods publickey keyboard-interactive:pam,publickey\n",
	}

	lines := string_array.ReadFile(filename)

	if string_array.ContainsAll(lines, lines_to_add) {
		return
	}

	fmt.Println("Updating", filename)

	// Comment out specific lines
	lines_to_comment := []string{
		"AuthorizedKeysCommand ",
		"AuthorizedKeysCommandUser ",
		"ChallengeResponseAuthentication ",
		"AuthenticationMethods ",
	}

	for idx, line := range lines {
		for _, check := range lines_to_comment {
			if strings.HasPrefix(line, check) {
				lines[idx] = "# " + line
			}
		}
	}

	backupFile(filename)

	check(string_array.WriteFile(filename, lines, lines_to_add))

	err := exec.Command("sshd", "-t").Run()

	if err != nil {
		if exerr, ok := err.(*exec.ExitError); ok {
			os.Stderr.Write(exerr.Stderr)
			os.Exit(1)
		}

		panic(err)
	}
}

func installToPam(selfPath string) {

	filename := "/etc/pam.d/sshd"
	fmt.Println("Updating", filename)

	pam_exec := "auth requisite pam_exec.so stdout quiet " + selfPath + " pam_create_user\n"

	lines := string_array.ReadFile(filename)

	for _, line := range lines {
		if line == pam_exec {
			return
		}
	}

	backupFile(filename)
	check(string_array.WriteFile(filename, []string{"# Next line added by " + selfPath + "\n", pam_exec}, lines))
}

func installToCron(selfPath string) {

	filename := "/etc/cron.d/" + path.Base(selfPath)

	fmt.Println("Installing crontab", filename)

	contents := "*/10 * * * * root " + selfPath + " sync_groups"

	check(ioutil.WriteFile(filename, []byte(contents), 0644))
}
<commit_msg>Rename lines_to_add and lines_to_comment to camelCase<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/user"
	"path"
	"strings"

	"github.com/davidrjonas/ssh-iam-bridge/string_array"
)

func check(err error) {
	if err == nil {
		return
	}

	if exerr, ok := err.(*exec.ExitError); ok {
		os.Stderr.Write(exerr.Stderr)
	}

	panic(err)
}

func backupFile(filename string) {
	fmt.Println("Backing up", filename, "to", filename+".orig")
	err := exec.Command("cp", "-f", filename, filename+".orig").Run()
	check(err)
}

func install(selfPath, username string) {
	cmdName := installAuthorizedKeysCommandScript(selfPath)
	installUser(username)
	installToSshd(cmdName, username)
	installToPam(selfPath)
	installToCron(selfPath)
}

// ssh is picky about AuthorizedKeysCommand, see man sshd_config
func installAuthorizedKeysCommandScript(selfPath string) string {
	cmdName := "/usr/sbin/ssh-iam-bridge-public-keys"
	fmt.Println("Writing AuthorizedKeysCommand script", cmdName)

	script := fmt.Sprintf("#!/bin/sh\nexec %s authorized_keys \"$@\"\n", selfPath)

	check(ioutil.WriteFile(cmdName, []byte(script), 0755))

	return cmdName
}

func installUser(username string) {
	_, err := user.Lookup(username)

	if err == nil {
		// User already exists
		return
	}

	if _, ok := err.(user.UnknownUserError); !ok {
		panic(err)
	}

	fmt.Println("Creating SSH authorized keys lookup user", username)

	args := []string{
		"--system",
		"--shell", "/usr/sbin/nologin",
		"--comment", "SSH authorized keys lookup",
		username,
	}

	_, err = exec.Command("useradd", args...).Output()
	check(err)
}

func installToSshd(cmd, username string) {

	filename := "/etc/ssh/sshd_config"

	// TODO: Ensure 'PasswordAuthentication no' and 'UsePAM yes'

	linesToAdd := []string{
		"AuthorizedKeysCommand " + cmd + "\n",
		"AuthorizedKeysCommandUser " + username + "\n",
		"ChallengeResponseAuthentication yes\n",
		"AuthenticationMethods publickey keyboard-interactive:pam,publickey\n",
	}

	lines := string_array.ReadFile(filename)

	if string_array.ContainsAll(lines, linesToAdd) {
		return
	}

	fmt.Println("Updating", filename)

	// Comment out specific lines
	linesToComment := []string{
		"AuthorizedKeysCommand ",
		"AuthorizedKeysCommandUser ",
		"ChallengeResponseAuthentication ",
		"AuthenticationMethods ",
	}

	for idx, line := range lines {
		for _, check := range linesToComment {
			if strings.HasPrefix(line, check) {
				lines[idx] = "# " + line
			}
		}
	}

	backupFile(filename)

	check(string_array.WriteFile(filename, lines, linesToAdd))

	err := exec.Command("sshd", "-t").Run()

	if err != nil {
		if exerr, ok := err.(*exec.ExitError); ok {
			os.Stderr.Write(exerr.Stderr)
			os.Exit(1)
		}

		panic(err)
	}
}

func installToPam(selfPath string) {

	filename := "/etc/pam.d/sshd"
	fmt.Println("Updating", filename)

	pam_exec := "auth requisite pam_exec.so stdout quiet " + selfPath + " pam_create_user\n"

	lines := string_array.ReadFile(filename)

	for _, line := range lines {
		if line == pam_exec {
			return
		}
	}

	backupFile(filename)
	check(string_array.WriteFile(filename, []string{"# Next line added by " + selfPath + "\n", pam_exec}, lines))
}

func installToCron(selfPath string) {

	filename := "/etc/cron.d/" + path.Base(selfPath)

	fmt.Println("Installing crontab", filename)

	contents := "*/10 * * * * root " + selfPath + " sync_groups"

	check(ioutil.WriteFile(filename, []byte(contents), 0644))
}
<|endoftext|>


<commit_before>package kd

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"koding/newkite/kd/util"
	"koding/newkite/kodingkey"
	"log"
	"net/http"
	"time"
)

const KeyLength = 64

var (
	AuthServer      = "https://koding.com"
	AuthServerLocal = "http://localhost:3020"
)

type Register struct {
	authServer string
}

func NewRegister() *Register {
	return &Register{
		authServer: AuthServer,
	}
}

func (r *Register) Definition() string {
	return "Register this host to Koding"
}

func (r *Register) Exec(args []string) error {
	// change authServer address if debug mode is enabled
	if len(args) == 1 && (args[0] == "--debug" || args[0] == "-d") {
		r.authServer = AuthServerLocal
	}

	hostID, err := util.HostID()
	if err != nil {
		return err
	}

	var key string
	keyExist := false

	key, err = util.GetKey()
	if err != nil {
		k, err := kodingkey.NewKodingKey()
		if err != nil {
			return err
		}

		key = k.String()
	} else {
		fmt.Printf("Found a key under '%s'. Going to use it to register\n", util.GetKdPath())
		keyExist = true
	}

	registerUrl := fmt.Sprintf("%s/-/auth/register/%s/%s", r.authServer, hostID, key)

	fmt.Printf("Please open the following url for authentication:\n\n")
	fmt.Println(registerUrl)
	fmt.Printf("\nwaiting . ")

	err = r.checker(key)
	if err != nil {
		return err
	}
	fmt.Println("successfully authenticated.")

	if keyExist {
		return nil
	}

	err = util.WriteKey(key)
	if err != nil {
		return err
	}

	return nil
}

// checker checks if the user has browsed the register URL by polling the check URL.
func (r *Register) checker(key string) error {
	checkUrl := fmt.Sprintf("%s/-/auth/check/%s", r.authServer, key)

	// check the result every two seconds
	ticker := time.NewTicker(2 * time.Second).C

	// wait for three minutes, if not successfull abort it
	timeout := time.After(3 * time.Minute)

	for {
		select {
		case <-ticker:
			err := checkResponse(checkUrl)
			if err != nil {
				// we didn't get OK message, continue until timout
				continue
			}

			return nil
		case <-timeout:
			return errors.New("timeout")
		}
	}
}

func checkResponse(checkUrl string) error {
	resp, err := http.Get(checkUrl)
	if err != nil {
		return err
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	resp.Body.Close()

	fmt.Printf(". ") // animation

	if resp.StatusCode != 200 {
		return errors.New("non 200 response")
	}

	type Result struct {
		Result string `json:"result"`
	}

	res := Result{}
	err = json.Unmarshal(bytes.TrimSpace(body), &res)
	if err != nil {
		log.Fatalln(err) // this should not happen, exit here
	}

	return nil
}
<commit_msg>kd/register: don't start checker if the user is already registered.<commit_after>package kd

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"koding/newkite/kd/util"
	"koding/newkite/kodingkey"
	"log"
	"net/http"
	"time"
)

const KeyLength = 64

var (
	AuthServer      = "https://koding.com"
	AuthServerLocal = "http://localhost:3020"
)

type Register struct {
	authServer string
}

func NewRegister() *Register {
	return &Register{
		authServer: AuthServer,
	}
}

func (r *Register) Definition() string {
	return "Register this host to Koding"
}

func (r *Register) Exec(args []string) error {
	// change authServer address if debug mode is enabled
	if len(args) == 1 && (args[0] == "--debug" || args[0] == "-d") {
		r.authServer = AuthServerLocal
	}

	hostID, err := util.HostID()
	if err != nil {
		return err
	}

	var key string
	keyExist := false

	key, err = util.GetKey()
	if err != nil {
		k, err := kodingkey.NewKodingKey()
		if err != nil {
			return err
		}

		key = k.String()
	} else {
		fmt.Printf("Found a key under '%s'. Going to use it to register\n", util.GetKdPath())
		keyExist = true
	}

	registerUrl := fmt.Sprintf("%s/-/auth/register/%s/%s", r.authServer, hostID, key)
	checkUrl := fmt.Sprintf("%s/-/auth/check/%s", r.authServer, key)

	// first check if the user is alrady registered
	err = checkResponse(checkUrl)
	if err == nil {
		fmt.Printf("... you are already registered.\n")
		return nil
	}

	fmt.Printf("Please open the following url for authentication:\n\n")
	fmt.Println(registerUrl)
	fmt.Printf("\nwaiting . ")

	// .. if not let the user register himself
	err = checker(checkUrl)
	if err != nil {
		return err
	}
	fmt.Println("successfully authenticated.")

	if keyExist {
		return nil
	}

	err = util.WriteKey(key)
	if err != nil {
		return err
	}

	return nil
}

// checker checks if the user has browsed the register URL by polling the check URL.
func checker(checkUrl string) error {
	// check the result every two seconds
	ticker := time.NewTicker(2 * time.Second).C

	// wait for three minutes, if not successfull abort it
	timeout := time.After(3 * time.Minute)

	for {
		select {
		case <-ticker:
			err := checkResponse(checkUrl)
			if err != nil {
				// we didn't get OK message, continue until timout
				fmt.Printf(". ") // animation
				continue
			}

			return nil
		case <-timeout:
			return errors.New("timeout")
		}
	}
}

func checkResponse(checkUrl string) error {
	resp, err := http.Get(checkUrl)
	if err != nil {
		return err
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	resp.Body.Close()

	if resp.StatusCode != 200 {
		return errors.New("non 200 response")
	}

	type Result struct {
		Result string `json:"result"`
	}

	res := Result{}
	err = json.Unmarshal(bytes.TrimSpace(body), &res)
	if err != nil {
		log.Fatalln(err) // this should not happen, exit here
	}

	return nil
}
<|endoftext|>


<commit_before>package api

import (
	"net/http"

	"github.com/golang/protobuf/jsonpb"
	"github.com/julienschmidt/httprouter"

	"github.com/octavore/press/proto/press/api"
	"github.com/octavore/press/proto/press/models"
	"github.com/octavore/press/server/router"
	"github.com/octavore/press/util/errors"
)

func (m *Module) ListRoutes(rw http.ResponseWriter, req *http.Request, par httprouter.Params) error {
	routes, err := m.DB.ListRoutes()
	if err != nil {
		return err
	}
	return router.Proto(rw, &api.ListRouteResponse{
		Routes: routes,
	})
}

func (m *Module) ListRoutesByPage(rw http.ResponseWriter, req *http.Request, par httprouter.Params) error {
	routes, err := m.DB.ListRoutes()
	if err != nil {
		return err
	}
	pageUUID := par.ByName("uuid")
	filteredRoutes := []*models.Route{}
	for _, route := range routes {
		if route.GetPageUuid() == pageUUID {
			filteredRoutes = append(filteredRoutes, route)
		}
	}
	return router.Proto(rw, &api.ListRouteResponse{
		Routes: filteredRoutes,
	})
}

func (m *Module) UpdateRoute(rw http.ResponseWriter, req *http.Request, par httprouter.Params) error {
	route := &models.Route{}
	err := jsonpb.Unmarshal(req.Body, route)
	if err != nil {
		return errors.Wrap(err)
	}
	err = m.DB.UpdateRoute(route)
	if err != nil {
		return err
	}
	err = m.Content.ReloadRouter()
	if err != nil {
		return nil
	}
	return router.Proto(rw, route)
}

func (m *Module) DeleteRoute(rw http.ResponseWriter, req *http.Request, par httprouter.Params) error {
	routeUUID := par.ByName("uuid")
	r, err := m.DB.GetRoute(routeUUID)
	if err != nil {
		return err
	}
	return m.DB.DeleteRoute(r)
}

func (m *Module) UpdateRoutesByPage(rw http.ResponseWriter, req *http.Request, par httprouter.Params) error {
	pageUUID := par.ByName("uuid")
	pb := &api.UpdatePageRoutesRequest{}
	err := jsonpb.Unmarshal(req.Body, pb)
	if err != nil {
		return errors.Wrap(err)
	}

	newList := map[string]*models.Route{}
	for _, route := range pb.GetRoutes() {
		route.Target = &models.Route_PageUuid{PageUuid: pageUUID}
		newList[route.GetUuid()] = route
	}

	routes, err := m.DB.ListRoutes()
	if err != nil {
		return err
	}

	filteredRoutes := []*models.Route{}
	for _, route := range routes {
		if route.GetPageUuid() == pageUUID {
			if newList[route.GetUuid()] == nil {
				err = m.DB.DeleteRoute(route)
				if err != nil {
					return err
				}
			}
		}
	}

	for _, route := range newList {
		err = m.DB.UpdateRoute(route)
		if err != nil {
			return err
		}
	}

	err = m.Content.ReloadRouter()
	if err != nil {
		return err
	}

	return router.Proto(rw, &api.ListRouteResponse{
		Routes: filteredRoutes,
	})
}
<commit_msg>server: Server side route cleaning.<commit_after>package api

import (
	"encoding/json"
	"net/http"
	"regexp"
	"strings"

	"github.com/julienschmidt/httprouter"

	"io/ioutil"

	"github.com/octavore/press/proto/press/api"
	"github.com/octavore/press/proto/press/models"
	"github.com/octavore/press/server/router"
	"github.com/octavore/press/util/errors"
)

var (
	re1 = regexp.MustCompile(`[^a-zA-Z0-9\/]`)
	re2 = regexp.MustCompile(`^-+`)
	re3 = regexp.MustCompile(`-+$`)
	re4 = regexp.MustCompile(`\/\/+`)
)

func formatRoute(r *models.Route) *models.Route {
	if r.Path == nil {
		return r
	}
	p := "/" + strings.Trim(r.GetPath(), "/")
	p = strings.ToLower(p)
	p = re1.ReplaceAllString(p, "-")
	p = re2.ReplaceAllString(p, "")
	p = re3.ReplaceAllString(p, "")
	p = re4.ReplaceAllString(p, "/")
	r.Path = &p
	return r
}

func (m *Module) ListRoutes(rw http.ResponseWriter, req *http.Request, par httprouter.Params) error {
	routes, err := m.DB.ListRoutes()
	if err != nil {
		return err
	}
	return router.Proto(rw, &api.ListRouteResponse{
		Routes: routes,
	})
}

func (m *Module) ListRoutesByPage(rw http.ResponseWriter, req *http.Request, par httprouter.Params) error {
	routes, err := m.DB.ListRoutes()
	if err != nil {
		return err
	}
	pageUUID := par.ByName("uuid")
	filteredRoutes := []*models.Route{}
	for _, route := range routes {
		if route.GetPageUuid() == pageUUID {
			filteredRoutes = append(filteredRoutes, route)
		}
	}
	return router.Proto(rw, &api.ListRouteResponse{
		Routes: filteredRoutes,
	})
}

func (m *Module) UpdateRoute(rw http.ResponseWriter, req *http.Request, par httprouter.Params) error {
	route := &models.Route{}
	b, err := ioutil.ReadAll(req.Body)
	if err != nil {
		return errors.Wrap(err)
	}
	err = json.Unmarshal(b, route)
	if err != nil {
		return errors.Wrap(err)
	}
	err = m.DB.UpdateRoute(formatRoute(route))
	if err != nil {
		return err
	}
	err = m.Content.ReloadRouter()
	if err != nil {
		return nil
	}
	return router.Proto(rw, route)
}

func (m *Module) DeleteRoute(rw http.ResponseWriter, req *http.Request, par httprouter.Params) error {
	routeUUID := par.ByName("uuid")
	r, err := m.DB.GetRoute(routeUUID)
	if err != nil {
		return err
	}
	return m.DB.DeleteRoute(r)
}

func (m *Module) UpdateRoutesByPage(rw http.ResponseWriter, req *http.Request, par httprouter.Params) error {
	pageUUID := par.ByName("uuid")
	pb := &api.UpdatePageRoutesRequest{}
	b, err := ioutil.ReadAll(req.Body)
	if err != nil {
		return errors.Wrap(err)
	}
	err = json.Unmarshal(b, pb)
	if err != nil {
		return errors.Wrap(err)
	}

	newList := map[string]*models.Route{}
	for _, route := range pb.GetRoutes() {
		route.Target = &models.Route_PageUuid{PageUuid: pageUUID}
		newList[route.GetUuid()] = route
	}

	routes, err := m.DB.ListRoutes()
	if err != nil {
		return err
	}

	filteredRoutes := []*models.Route{}
	for _, route := range routes {
		if route.GetPageUuid() == pageUUID {
			if newList[route.GetUuid()] == nil {
				err = m.DB.DeleteRoute(route)
				if err != nil {
					return err
				}
			}
		}
	}

	for _, route := range newList {
		err = m.DB.UpdateRoute(formatRoute(route))
		if err != nil {
			return err
		}
	}

	err = m.Content.ReloadRouter()
	if err != nil {
		return err
	}

	return router.Proto(rw, &api.ListRouteResponse{
		Routes: filteredRoutes,
	})
}
<|endoftext|>


<commit_before>package flowcontrol

import (
	"fmt"

	"github.com/lucas-clemente/quic-go/internal/congestion"
	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/utils"
	"github.com/lucas-clemente/quic-go/qerr"
)

type streamFlowController struct {
	baseFlowController

	streamID protocol.StreamID

	connection              connectionFlowControllerI
	contributesToConnection bool // does the stream contribute to connection level flow control

	receivedFinalOffset bool
}

var _ StreamFlowController = &streamFlowController{}

// NewStreamFlowController gets a new flow controller for a stream
func NewStreamFlowController(
	streamID protocol.StreamID,
	contributesToConnection bool,
	cfc ConnectionFlowController,
	receiveWindow protocol.ByteCount,
	maxReceiveWindow protocol.ByteCount,
	initialSendWindow protocol.ByteCount,
	rttStats *congestion.RTTStats,
	logger utils.Logger,
) StreamFlowController {
	return &streamFlowController{
		streamID:                streamID,
		contributesToConnection: contributesToConnection,
		connection:              cfc.(connectionFlowControllerI),
		baseFlowController: baseFlowController{
			rttStats:             rttStats,
			receiveWindow:        receiveWindow,
			receiveWindowSize:    receiveWindow,
			maxReceiveWindowSize: maxReceiveWindow,
			sendWindow:           initialSendWindow,
			logger:               logger,
		},
	}
}

// UpdateHighestReceived updates the highestReceived value, if the byteOffset is higher
// it returns an ErrReceivedSmallerByteOffset if the received byteOffset is smaller than any byteOffset received before
func (c *streamFlowController) UpdateHighestReceived(byteOffset protocol.ByteCount, final bool) error {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	// when receiving a final offset, check that this final offset is consistent with a final offset we might have received earlier
	if final && c.receivedFinalOffset && byteOffset != c.highestReceived {
		return qerr.Error(qerr.StreamDataAfterTermination, fmt.Sprintf("Received inconsistent final offset for stream %d (old: %d, new: %d bytes)", c.streamID, c.highestReceived, byteOffset))
	}
	// if we already received a final offset, check that the offset in the STREAM frames is below the final offset
	if c.receivedFinalOffset && byteOffset > c.highestReceived {
		return qerr.StreamDataAfterTermination
	}
	if final {
		c.receivedFinalOffset = true
	}
	if byteOffset == c.highestReceived {
		return nil
	}
	if byteOffset <= c.highestReceived {
		// a STREAM_FRAME with a higher offset was received before.
		if final {
			// If the current byteOffset is smaller than the offset in that STREAM_FRAME, this STREAM_FRAME contained data after the end of the stream
			return qerr.StreamDataAfterTermination
		}
		// this is a reordered STREAM_FRAME
		return nil
	}

	increment := byteOffset - c.highestReceived
	c.highestReceived = byteOffset
	if c.checkFlowControlViolation() {
		return qerr.Error(qerr.FlowControlReceivedTooMuchData, fmt.Sprintf("Received %d bytes on stream %d, allowed %d bytes", byteOffset, c.streamID, c.receiveWindow))
	}
	if c.contributesToConnection {
		return c.connection.IncrementHighestReceived(increment)
	}
	return nil
}

func (c *streamFlowController) AddBytesRead(n protocol.ByteCount) {
	c.baseFlowController.AddBytesRead(n)
	if c.contributesToConnection {
		c.connection.AddBytesRead(n)
	}
}

func (c *streamFlowController) AddBytesSent(n protocol.ByteCount) {
	c.baseFlowController.AddBytesSent(n)
	if c.contributesToConnection {
		c.connection.AddBytesSent(n)
	}
}

func (c *streamFlowController) SendWindowSize() protocol.ByteCount {
	window := c.baseFlowController.sendWindowSize()
	if c.contributesToConnection {
		window = utils.MinByteCount(window, c.connection.SendWindowSize())
	}
	return window
}

// IsBlocked says if it is blocked by stream-level flow control.
// If it is blocked, the offset is returned.
func (c *streamFlowController) IsBlocked() (bool, protocol.ByteCount) {
	if c.sendWindowSize() != 0 {
		return false, 0
	}
	return true, c.sendWindow
}

func (c *streamFlowController) HasWindowUpdate() bool {
	c.mutex.Lock()
	hasWindowUpdate := !c.receivedFinalOffset && c.hasWindowUpdate()
	c.mutex.Unlock()
	return hasWindowUpdate
}

func (c *streamFlowController) GetWindowUpdate() protocol.ByteCount {
	// don't use defer for unlocking the mutex here, GetWindowUpdate() is called frequently and defer shows up in the profiler
	c.mutex.Lock()
	// if we already received the final offset for this stream, the peer won't need any additional flow control credit
	if c.receivedFinalOffset {
		c.mutex.Unlock()
		return 0
	}

	oldWindowSize := c.receiveWindowSize
	offset := c.baseFlowController.getWindowUpdate()
	if c.receiveWindowSize > oldWindowSize { // auto-tuning enlarged the window size
		c.logger.Debugf("Increasing receive flow control window for the connection to %d kB", c.receiveWindowSize/(1<<10))
		if c.contributesToConnection {
			c.connection.EnsureMinimumWindowSize(protocol.ByteCount(float64(c.receiveWindowSize) * protocol.ConnectionFlowControlMultiplier))
		}
	}
	c.mutex.Unlock()
	return offset
}
<commit_msg>fix logging of stream flow control window increases<commit_after>package flowcontrol

import (
	"fmt"

	"github.com/lucas-clemente/quic-go/internal/congestion"
	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/utils"
	"github.com/lucas-clemente/quic-go/qerr"
)

type streamFlowController struct {
	baseFlowController

	streamID protocol.StreamID

	connection              connectionFlowControllerI
	contributesToConnection bool // does the stream contribute to connection level flow control

	receivedFinalOffset bool
}

var _ StreamFlowController = &streamFlowController{}

// NewStreamFlowController gets a new flow controller for a stream
func NewStreamFlowController(
	streamID protocol.StreamID,
	contributesToConnection bool,
	cfc ConnectionFlowController,
	receiveWindow protocol.ByteCount,
	maxReceiveWindow protocol.ByteCount,
	initialSendWindow protocol.ByteCount,
	rttStats *congestion.RTTStats,
	logger utils.Logger,
) StreamFlowController {
	return &streamFlowController{
		streamID:                streamID,
		contributesToConnection: contributesToConnection,
		connection:              cfc.(connectionFlowControllerI),
		baseFlowController: baseFlowController{
			rttStats:             rttStats,
			receiveWindow:        receiveWindow,
			receiveWindowSize:    receiveWindow,
			maxReceiveWindowSize: maxReceiveWindow,
			sendWindow:           initialSendWindow,
			logger:               logger,
		},
	}
}

// UpdateHighestReceived updates the highestReceived value, if the byteOffset is higher
// it returns an ErrReceivedSmallerByteOffset if the received byteOffset is smaller than any byteOffset received before
func (c *streamFlowController) UpdateHighestReceived(byteOffset protocol.ByteCount, final bool) error {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	// when receiving a final offset, check that this final offset is consistent with a final offset we might have received earlier
	if final && c.receivedFinalOffset && byteOffset != c.highestReceived {
		return qerr.Error(qerr.StreamDataAfterTermination, fmt.Sprintf("Received inconsistent final offset for stream %d (old: %d, new: %d bytes)", c.streamID, c.highestReceived, byteOffset))
	}
	// if we already received a final offset, check that the offset in the STREAM frames is below the final offset
	if c.receivedFinalOffset && byteOffset > c.highestReceived {
		return qerr.StreamDataAfterTermination
	}
	if final {
		c.receivedFinalOffset = true
	}
	if byteOffset == c.highestReceived {
		return nil
	}
	if byteOffset <= c.highestReceived {
		// a STREAM_FRAME with a higher offset was received before.
		if final {
			// If the current byteOffset is smaller than the offset in that STREAM_FRAME, this STREAM_FRAME contained data after the end of the stream
			return qerr.StreamDataAfterTermination
		}
		// this is a reordered STREAM_FRAME
		return nil
	}

	increment := byteOffset - c.highestReceived
	c.highestReceived = byteOffset
	if c.checkFlowControlViolation() {
		return qerr.Error(qerr.FlowControlReceivedTooMuchData, fmt.Sprintf("Received %d bytes on stream %d, allowed %d bytes", byteOffset, c.streamID, c.receiveWindow))
	}
	if c.contributesToConnection {
		return c.connection.IncrementHighestReceived(increment)
	}
	return nil
}

func (c *streamFlowController) AddBytesRead(n protocol.ByteCount) {
	c.baseFlowController.AddBytesRead(n)
	if c.contributesToConnection {
		c.connection.AddBytesRead(n)
	}
}

func (c *streamFlowController) AddBytesSent(n protocol.ByteCount) {
	c.baseFlowController.AddBytesSent(n)
	if c.contributesToConnection {
		c.connection.AddBytesSent(n)
	}
}

func (c *streamFlowController) SendWindowSize() protocol.ByteCount {
	window := c.baseFlowController.sendWindowSize()
	if c.contributesToConnection {
		window = utils.MinByteCount(window, c.connection.SendWindowSize())
	}
	return window
}

// IsBlocked says if it is blocked by stream-level flow control.
// If it is blocked, the offset is returned.
func (c *streamFlowController) IsBlocked() (bool, protocol.ByteCount) {
	if c.sendWindowSize() != 0 {
		return false, 0
	}
	return true, c.sendWindow
}

func (c *streamFlowController) HasWindowUpdate() bool {
	c.mutex.Lock()
	hasWindowUpdate := !c.receivedFinalOffset && c.hasWindowUpdate()
	c.mutex.Unlock()
	return hasWindowUpdate
}

func (c *streamFlowController) GetWindowUpdate() protocol.ByteCount {
	// don't use defer for unlocking the mutex here, GetWindowUpdate() is called frequently and defer shows up in the profiler
	c.mutex.Lock()
	// if we already received the final offset for this stream, the peer won't need any additional flow control credit
	if c.receivedFinalOffset {
		c.mutex.Unlock()
		return 0
	}

	oldWindowSize := c.receiveWindowSize
	offset := c.baseFlowController.getWindowUpdate()
	if c.receiveWindowSize > oldWindowSize { // auto-tuning enlarged the window size
		c.logger.Debugf("Increasing receive flow control window for stream %d to %d kB", c.streamID, c.receiveWindowSize/(1<<10))
		if c.contributesToConnection {
			c.connection.EnsureMinimumWindowSize(protocol.ByteCount(float64(c.receiveWindowSize) * protocol.ConnectionFlowControlMultiplier))
		}
	}
	c.mutex.Unlock()
	return offset
}
<|endoftext|>


<commit_before>package action

import (
	"fmt"
	"strings"

	"github.com/fatih/color"
	"github.com/justwatchcom/gopass/termwiz"
	"github.com/urfave/cli"
)

// Find a string in the secret file's name
func (s *Action) Find(c *cli.Context) error {
	if !c.Args().Present() {
		return fmt.Errorf("Usage: gopass find arg")
	}

	l, err := s.Store.List(0)
	if err != nil {
		return err
	}
	needle := strings.ToLower(c.Args().First())
	choices := make([]string, 0, 10)
	for _, value := range l {
		if strings.Contains(strings.ToLower(value), needle) {
			choices = append(choices, value)
		}
	}

	if len(choices) == 1 {
		fmt.Println(color.GreenString("Found exact match in '%s'", choices[0]))
		return s.show(c, choices[0], "", false, false, false)
	}

	if !s.isTerm {
		for _, value := range choices {
			fmt.Println(value)
		}
		return nil
	}

	act, sel := termwiz.GetSelection(choices)
	switch act {
	case "copy":
		return s.show(c, choices[sel], "", true, false, false)
	case "show":
		return s.show(c, choices[sel], "", false, false, false)
	default:
		return fmt.Errorf("User aborted")
	}
}
<commit_msg>Only show selection if results were found<commit_after>package action

import (
	"fmt"
	"strings"

	"github.com/fatih/color"
	"github.com/justwatchcom/gopass/termwiz"
	"github.com/urfave/cli"
)

// Find a string in the secret file's name
func (s *Action) Find(c *cli.Context) error {
	if !c.Args().Present() {
		return fmt.Errorf("Usage: gopass find arg")
	}

	l, err := s.Store.List(0)
	if err != nil {
		return err
	}
	needle := strings.ToLower(c.Args().First())
	choices := make([]string, 0, 10)
	for _, value := range l {
		if strings.Contains(strings.ToLower(value), needle) {
			choices = append(choices, value)
		}
	}

	if len(choices) < 1 {
		return fmt.Errorf("not results found")
	}

	if len(choices) == 1 {
		fmt.Println(color.GreenString("Found exact match in '%s'", choices[0]))
		return s.show(c, choices[0], "", false, false, false)
	}

	if !s.isTerm {
		for _, value := range choices {
			fmt.Println(value)
		}
		return nil
	}

	act, sel := termwiz.GetSelection(choices)
	switch act {
	case "copy":
		return s.show(c, choices[sel], "", true, false, false)
	case "show":
		return s.show(c, choices[sel], "", false, false, false)
	default:
		return fmt.Errorf("User aborted")
	}
}
<|endoftext|>


<commit_before>/*
OstreeDeploy Action

Deploy the OSTree branch to the image.
If any preparation has been done for rootfs, it can be overwritten
during this step.

Action 'image-partition' must be called prior to OSTree deploy.

Yaml syntax:
 - action: ostree-deploy
   repository: repository name
   remote_repository: URL
   branch: branch name
   os: os name
   tls-client-cert-path: path to client certificate
   tls-client-key-path: path to client certificate key
   setup-fstab: bool
   setup-kernel-cmdline: bool
   appendkernelcmdline: arguments
   collection-id: org.apertis.example

Mandatory properties:

- remote_repository -- URL to remote OSTree repository for pulling stateroot branch.
Currently not implemented, please prepare local repository instead.

- repository -- path to repository with OSTree structure.
This path is relative to 'artifact' directory.

- os -- os deployment name, as explained in:
https://ostree.readthedocs.io/en/latest/manual/deployment/

- branch -- branch of the repository to use for populating the image.

Optional properties:

- setup-fstab -- create '/etc/fstab' file for image

- setup-kernel-cmdline -- add the information from the 'image-partition'
action to the configured commandline.

- append-kernel-cmdline -- additional kernel command line arguments passed to kernel.

- tls-client-cert-path -- path to client certificate to use for the remote repository

- tls-client-key-path -- path to client certificate key to use for the remote repository

- collection-id -- Collection ID ref binding (require libostree 2018.6).
*/
package actions

import (
	"fmt"
	"io"
	"os"
	"path"
	"runtime"
	"strings"

	"github.com/go-debos/debos"
	ostree "github.com/sjoerdsimons/ostree-go/pkg/otbuiltin"
)

type OstreeDeployAction struct {
	debos.BaseAction    `yaml:",inline"`
	Repository          string
	RemoteRepository    string "remote_repository"
	Branch              string
	Os                  string
	SetupFSTab          bool   `yaml:"setup-fstab"`
	SetupKernelCmdline  bool   `yaml:"setup-kernel-cmdline"`
	AppendKernelCmdline string `yaml:"append-kernel-cmdline"`
	TlsClientCertPath   string `yaml:"tls-client-cert-path"`
	TlsClientKeyPath    string `yaml:"tls-client-key-path"`
	CollectionID        string `yaml:"collection-id"`
}

func NewOstreeDeployAction() *OstreeDeployAction {
	ot := &OstreeDeployAction{SetupFSTab: true, SetupKernelCmdline: true}
	ot.Description = "Deploying from ostree"
	return ot
}

func (ot *OstreeDeployAction) setupFSTab(deployment *ostree.Deployment, context *debos.DebosContext) error {
	deploymentDir := fmt.Sprintf("ostree/deploy/%s/deploy/%s.%d",
		deployment.Osname(), deployment.Csum(), deployment.Deployserial())

	etcDir := path.Join(context.Rootdir, deploymentDir, "etc")

	err := os.Mkdir(etcDir, 0755)
	if err != nil && !os.IsExist(err) {
		return err
	}

	dst, err := os.OpenFile(path.Join(etcDir, "fstab"), os.O_WRONLY|os.O_CREATE, 0755)
	defer dst.Close()
	if err != nil {
		return err
	}

	_, err = io.Copy(dst, &context.ImageFSTab)

	return err
}

func (ot *OstreeDeployAction) Run(context *debos.DebosContext) error {
	ot.LogStart()

	// This is to handle cases there we didn't partition an image
	if len(context.ImageMntDir) != 0 {
		/* First deploy the current rootdir to the image so it can seed e.g.
		 * bootloader configuration */
		err := debos.Command{}.Run("Deploy to image", "cp", "-a", context.Rootdir+"/.", context.ImageMntDir)
		if err != nil {
			return fmt.Errorf("rootfs deploy failed: %v", err)
		}
		context.Rootdir = context.ImageMntDir
		context.Origins["filesystem"] = context.ImageMntDir
	}

	repoPath := "file://" + path.Join(context.Artifactdir, ot.Repository)

	sysroot := ostree.NewSysroot(context.Rootdir)
	err := sysroot.InitializeFS()
	if err != nil {
		return err
	}

	err = sysroot.InitOsname(ot.Os, nil)
	if err != nil {
		return err
	}

	/* HACK: Getting the repository form the sysroot gets ostree confused on
	 * whether it should configure /etc/ostree or the repo configuration,
	   so reopen by hand */
	/* dstRepo, err := sysroot.Repo(nil) */
	dstRepo, err := ostree.OpenRepo(path.Join(context.Rootdir, "ostree/repo"))
	if err != nil {
		return err
	}

	/* FIXME: add support for gpg signing commits so this is no longer needed */
	opts := ostree.RemoteOptions{NoGpgVerify: true,
		TlsClientCertPath: ot.TlsClientCertPath,
		TlsClientKeyPath:  ot.TlsClientKeyPath,
		CollectionId:      ot.CollectionID,
	}

	err = dstRepo.RemoteAdd("origin", ot.RemoteRepository, opts, nil)
	if err != nil {
		return err
	}

	var options ostree.PullOptions
	options.OverrideRemoteName = "origin"
	options.Refs = []string{ot.Branch}

	err = dstRepo.PullWithOptions(repoPath, options, nil, nil)
	if err != nil {
		return err
	}

	/* Required by ostree to make sure a bunch of information was pulled in  */
	sysroot.Load(nil)

	revision, err := dstRepo.ResolveRev(ot.Branch, false)
	if err != nil {
		return err
	}

	var kargs []string
	if ot.SetupKernelCmdline {
		kargs = append(kargs, context.ImageKernelRoot)
	}

	if ot.AppendKernelCmdline != "" {
		s := strings.Split(ot.AppendKernelCmdline, " ")
		kargs = append(kargs, s...)
	}

	origin := sysroot.OriginNewFromRefspec("origin:" + ot.Branch)
	deployment, err := sysroot.DeployTree(ot.Os, revision, origin, nil, kargs, nil)
	if err != nil {
		return err
	}

	if ot.SetupFSTab {
		err = ot.setupFSTab(deployment, context)
		if err != nil {
			return err
		}
	}

	err = sysroot.SimpleWriteDeployment(ot.Os, deployment, nil, 0, nil)
	if err != nil {
		return err
	}

	runtime.GC()
	return nil
}
<commit_msg>ostree-deploy: fix umount error<commit_after>/*
OstreeDeploy Action

Deploy the OSTree branch to the image.
If any preparation has been done for rootfs, it can be overwritten
during this step.

Action 'image-partition' must be called prior to OSTree deploy.

Yaml syntax:
 - action: ostree-deploy
   repository: repository name
   remote_repository: URL
   branch: branch name
   os: os name
   tls-client-cert-path: path to client certificate
   tls-client-key-path: path to client certificate key
   setup-fstab: bool
   setup-kernel-cmdline: bool
   appendkernelcmdline: arguments
   collection-id: org.apertis.example

Mandatory properties:

- remote_repository -- URL to remote OSTree repository for pulling stateroot branch.
Currently not implemented, please prepare local repository instead.

- repository -- path to repository with OSTree structure.
This path is relative to 'artifact' directory.

- os -- os deployment name, as explained in:
https://ostree.readthedocs.io/en/latest/manual/deployment/

- branch -- branch of the repository to use for populating the image.

Optional properties:

- setup-fstab -- create '/etc/fstab' file for image

- setup-kernel-cmdline -- add the information from the 'image-partition'
action to the configured commandline.

- append-kernel-cmdline -- additional kernel command line arguments passed to kernel.

- tls-client-cert-path -- path to client certificate to use for the remote repository

- tls-client-key-path -- path to client certificate key to use for the remote repository

- collection-id -- Collection ID ref binding (require libostree 2018.6).
*/
package actions

import (
	"fmt"
	"io"
	"os"
	"path"
	"strings"

	"github.com/go-debos/debos"
	ostree "github.com/sjoerdsimons/ostree-go/pkg/otbuiltin"
)

type OstreeDeployAction struct {
	debos.BaseAction    `yaml:",inline"`
	Repository          string
	RemoteRepository    string "remote_repository"
	Branch              string
	Os                  string
	SetupFSTab          bool   `yaml:"setup-fstab"`
	SetupKernelCmdline  bool   `yaml:"setup-kernel-cmdline"`
	AppendKernelCmdline string `yaml:"append-kernel-cmdline"`
	TlsClientCertPath   string `yaml:"tls-client-cert-path"`
	TlsClientKeyPath    string `yaml:"tls-client-key-path"`
	CollectionID        string `yaml:"collection-id"`
}

func NewOstreeDeployAction() *OstreeDeployAction {
	ot := &OstreeDeployAction{SetupFSTab: true, SetupKernelCmdline: true}
	ot.Description = "Deploying from ostree"
	return ot
}

func (ot *OstreeDeployAction) setupFSTab(deployment *ostree.Deployment, context *debos.DebosContext) error {
	deploymentDir := fmt.Sprintf("ostree/deploy/%s/deploy/%s.%d",
		deployment.Osname(), deployment.Csum(), deployment.Deployserial())

	etcDir := path.Join(context.Rootdir, deploymentDir, "etc")

	err := os.Mkdir(etcDir, 0755)
	if err != nil && !os.IsExist(err) {
		return err
	}

	dst, err := os.OpenFile(path.Join(etcDir, "fstab"), os.O_WRONLY|os.O_CREATE, 0755)
	defer dst.Close()
	if err != nil {
		return err
	}

	_, err = io.Copy(dst, &context.ImageFSTab)

	return err
}

func (ot *OstreeDeployAction) Run(context *debos.DebosContext) error {
	ot.LogStart()

	// This is to handle cases there we didn't partition an image
	if len(context.ImageMntDir) != 0 {
		/* First deploy the current rootdir to the image so it can seed e.g.
		 * bootloader configuration */
		err := debos.Command{}.Run("Deploy to image", "cp", "-a", context.Rootdir+"/.", context.ImageMntDir)
		if err != nil {
			return fmt.Errorf("rootfs deploy failed: %v", err)
		}
		context.Rootdir = context.ImageMntDir
		context.Origins["filesystem"] = context.ImageMntDir
	}

	repoPath := "file://" + path.Join(context.Artifactdir, ot.Repository)

	sysroot := ostree.NewSysroot(context.Rootdir)
	err := sysroot.InitializeFS()
	if err != nil {
		return err
	}

	err = sysroot.InitOsname(ot.Os, nil)
	if err != nil {
		return err
	}

	/* HACK: Getting the repository form the sysroot gets ostree confused on
	 * whether it should configure /etc/ostree or the repo configuration,
	   so reopen by hand */
	/* dstRepo, err := sysroot.Repo(nil) */
	dstRepo, err := ostree.OpenRepo(path.Join(context.Rootdir, "ostree/repo"))
	if err != nil {
		return err
	}

	/* FIXME: add support for gpg signing commits so this is no longer needed */
	opts := ostree.RemoteOptions{NoGpgVerify: true,
		TlsClientCertPath: ot.TlsClientCertPath,
		TlsClientKeyPath:  ot.TlsClientKeyPath,
		CollectionId:      ot.CollectionID,
	}

	err = dstRepo.RemoteAdd("origin", ot.RemoteRepository, opts, nil)
	if err != nil {
		return err
	}

	var options ostree.PullOptions
	options.OverrideRemoteName = "origin"
	options.Refs = []string{ot.Branch}

	err = dstRepo.PullWithOptions(repoPath, options, nil, nil)
	if err != nil {
		return err
	}

	/* Required by ostree to make sure a bunch of information was pulled in  */
	sysroot.Load(nil)

	revision, err := dstRepo.ResolveRev(ot.Branch, false)
	if err != nil {
		return err
	}

	var kargs []string
	if ot.SetupKernelCmdline {
		kargs = append(kargs, context.ImageKernelRoot)
	}

	if ot.AppendKernelCmdline != "" {
		s := strings.Split(ot.AppendKernelCmdline, " ")
		kargs = append(kargs, s...)
	}

	origin := sysroot.OriginNewFromRefspec("origin:" + ot.Branch)
	deployment, err := sysroot.DeployTree(ot.Os, revision, origin, nil, kargs, nil)
	if err != nil {
		return err
	}

	if ot.SetupFSTab {
		err = ot.setupFSTab(deployment, context)
		if err != nil {
			return err
		}
	}

	err = sysroot.SimpleWriteDeployment(ot.Os, deployment, nil, 0, nil)
	if err != nil {
		return err
	}

	/* libostree keeps some information, like repo lock file descriptor, in
	 * thread specific variables. As GC can be run from another thread, it
	 * may not been able to access this, preventing to free them correctly.
	 * To prevent this, explicitly dereference libostree objects. */
	dstRepo.Unref()
	sysroot.Unref()
	return nil
}
<|endoftext|>